[GCP] Add H100 mega (#4099)

* Add H100 mega support on GCP

* fix for some other regions

* format

* fix resource type

* fix catalog fetching

Author: Michaelvll

sky/clouds/gcp.py

@@ -483,7 +483,7 @@ def _failover_disk_tier() -> Optional[resources_utils.DiskTier]:
                 if acc in ('A100-80GB', 'L4'):
                     # A100-80GB and L4 have a different name pattern.
                     resources_vars['gpu'] = f'nvidia-{acc.lower()}'
-                elif acc == 'H100':
+                elif acc in ('H100', 'H100-MEGA'):
                     resources_vars['gpu'] = f'nvidia-{acc.lower()}-80gb'
                 else:
                     resources_vars['gpu'] = 'nvidia-tesla-{}'.format(

sky/clouds/service_catalog/data_fetchers/fetch_gcp.py

@@ -419,6 +419,11 @@ def _get_gpus_for_zone(zone: str) -> 'pd.DataFrame':
                 if count != 8:
                     # H100 only has 8 cards.
                     continue
+            if 'H100-MEGA-80GB' in gpu_name:
+                gpu_name = 'H100-MEGA'
+                if count != 8:
+                    # H100-MEGA only has 8 cards.
+                    continue
             if 'VWS' in gpu_name:
                 continue
             if gpu_name.startswith('TPU-'):
@@ -447,6 +452,7 @@ def _gpu_info_from_name(name: str) -> Optional[Dict[str, List[Dict[str, Any]]]]:
         'A100-80GB': 80 * 1024,
         'A100': 40 * 1024,
         'H100': 80 * 1024,
+        'H100-MEGA': 80 * 1024,
         'P4': 8 * 1024,
         'T4': 16 * 1024,
         'V100': 16 * 1024,
@@ -491,12 +497,17 @@ def get_gpu_price(row: pd.Series, spot: bool) -> Optional[float]:
             if sku['category']['usageType'] != ondemand_or_spot:
                 continue
 
-            gpu_name = row['AcceleratorName']
-            if gpu_name == 'A100-80GB':
-                gpu_name = 'A100 80GB'
-            if gpu_name == 'H100':
-                gpu_name = 'H100 80GB'
-            if f'{gpu_name} GPU' not in sku['description']:
+            gpu_names = [row['AcceleratorName']]
+            if gpu_names[0] == 'A100-80GB':
+                gpu_names = ['A100 80GB']
+            if gpu_names[0] == 'H100':
+                gpu_names = ['H100 80GB']
+            if gpu_names[0] == 'H100-MEGA':
+                # Seems that H100-MEGA has two different descriptions in SKUs in
+                # different regions: 'H100 80GB Mega' and 'H100 80GB Plus'.
+                gpu_names = ['H100 80GB Mega', 'H100 80GB Plus']
+            if not any(f'{gpu_name} GPU' in sku['description']
+                       for gpu_name in gpu_names):
                 continue
 
             unit_price = _get_unit_price(sku)

sky/clouds/service_catalog/gcp_catalog.py

@@ -98,6 +98,9 @@
     },
     'H100': {
         8: ['a3-highgpu-8g'],
+    },
+    'H100-MEGA': {
+        8: ['a3-megagpu-8g'],
     }
 }