skypilot-nightly 1.0.0.dev20250502__py3-none-any.whl → 1.0.0.dev20251203__py3-none-any.whl

sky/{clouds/service_catalog → catalog}/data_fetchers/fetch_aws.py

@@ -46,7 +46,7 @@ ALL_REGIONS = [
     'eu-west-1',
     'eu-west-2',
     'eu-south-1',
-    # 'eu-south-2', # no supported AMI
+    'eu-south-2',
     'eu-west-3',
     'eu-north-1',
     'me-south-1',
@@ -60,6 +60,7 @@ ALL_REGIONS = [
     'ap-northeast-2',
     'ap-southeast-1',
     'ap-southeast-2',
+    'ap-southeast-4',
     'ap-northeast-1',
 ]
 US_REGIONS = ['us-east-1', 'us-east-2', 'us-west-1', 'us-west-2']
@@ -67,17 +68,17 @@ US_REGIONS = ['us-east-1', 'us-east-2', 'us-west-1', 'us-west-2']
 # The following columns will be included in the final catalog.
 USEFUL_COLUMNS = [
     'InstanceType', 'AcceleratorName', 'AcceleratorCount', 'vCPUs', 'MemoryGiB',
-    'GpuInfo', 'Price', 'SpotPrice', 'Region', 'AvailabilityZone'
+    'GpuInfo', 'Price', 'SpotPrice', 'Region', 'AvailabilityZone', 'Arch'
 ]
 
 # NOTE: the hard-coded us-east-1 URL is not a typo. AWS pricing endpoint is
 # only available in this region, but it serves pricing information for all
 # regions.
 PRICING_TABLE_URL_FMT = 'https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonEC2/current/{region}/index.csv'  # pylint: disable=line-too-long
-# Hardcode the regions that offer p4de.24xlarge as our credential does not have
-# the permission to query the offerings of the instance.
-# Ref: https://aws.amazon.com/ec2/instance-types/p4/
-P4DE_REGIONS = ['us-east-1', 'us-west-2']
+# g6f instances have fractional GPUs, but the API returns Count: 1 under
+# GpuInfo. However, the GPU memory is properly scaled. Taking the instance GPU
+# divided by the total memory of an L4 will give us the fraction of the GPU.
+L4_GPU_MEMORY = 22888
 
 regions_enabled: Optional[Set[str]] = None
 
@@ -210,35 +211,6 @@ def _get_spot_pricing_table(region: str) -> 'pd.DataFrame':
     return df
 
 
-def _patch_p4de(region: str, df: 'pd.DataFrame',
-                pricing_df: 'pd.DataFrame') -> 'pd.DataFrame':
-    # Hardcoded patch for p4de.24xlarge, as our credentials doesn't have access
-    # to the instance type.
-    # Columns:
-    # InstanceType,AcceleratorName,AcceleratorCount,vCPUs,MemoryGiB,GpuInfo,
-    # Price,SpotPrice,Region,AvailabilityZone
-    records = []
-    for zone in df[df['Region'] == region]['AvailabilityZone'].unique():
-        records.append({
-            'InstanceType': 'p4de.24xlarge',
-            'AcceleratorName': 'A100-80GB',
-            'AcceleratorCount': 8,
-            'vCPUs': 96,
-            'MemoryGiB': 1152,
-            'GpuInfo':
-                ('{\'Gpus\': [{\'Name\': \'A100-80GB\', \'Manufacturer\': '
-                 '\'NVIDIA\', \'Count\': 8, \'MemoryInfo\': {\'SizeInMiB\': '
-                 '81920}}], \'TotalGpuMemoryInMiB\': 655360}'),
-            'AvailabilityZone': zone,
-            'Region': region,
-            'Price': pricing_df[pricing_df['InstanceType'] == 'p4de.24xlarge']
-                     ['Price'].values[0],
-            'SpotPrice': np.nan,
-        })
-    df = pd.concat([df, pd.DataFrame.from_records(records)])
-    return df
-
-
 def _get_instance_types_df(region: str) -> Union[str, 'pd.DataFrame']:
     try:
         # Fetch the zone info first to make sure the account has access to the
@@ -262,7 +234,7 @@ def _get_instance_types_df(region: str) -> Union[str, 'pd.DataFrame']:
         def get_acc_info(row) -> Tuple[Optional[str], float]:
             accelerator = None
             for col, info_key in [('GpuInfo', 'Gpus'),
-                                  ('InferenceAcceleratorInfo', 'Accelerators'),
+                                  ('NeuronInfo', 'NeuronDevices'),
                                   ('FpgaInfo', 'Fpgas')]:
                 info = row.get(col)
                 if isinstance(info, dict):
@@ -271,13 +243,24 @@ def _get_instance_types_df(region: str) -> Union[str, 'pd.DataFrame']:
                 return None, np.nan
             return accelerator['Name'], accelerator['Count']
 
+        def get_arch(row) -> Optional[str]:
+            if 'ProcessorInfo' in row:
+                processor = row['ProcessorInfo']
+                if 'SupportedArchitectures' in processor:
+                    archs = processor['SupportedArchitectures']
+                    if isinstance(archs, list):
+                        return archs[0]
+                    elif isinstance(archs, str):
+                        return archs
+            return None
+
         def get_vcpus(row) -> float:
             if not np.isnan(row['vCPU']):
                 return float(row['vCPU'])
             try:
                 return float(row['VCpuInfo']['DefaultVCpus'])
             except Exception as e:  # pylint: disable=broad-except
-                print('Error occured for row:', row)
+                print('Error occurred for row:', row)
                 print('Error:', e)
                 raise
 
@@ -295,29 +278,33 @@ def _get_instance_types_df(region: str) -> Union[str, 'pd.DataFrame']:
             if row['InstanceType'] == 'p4de.24xlarge':
                 acc_name = 'A100-80GB'
                 acc_count = 8
-            if row['InstanceType'].startswith('trn1'):
-                # Trainium instances does not have a field for information of
-                # the accelerators. We need to infer the accelerator info from
-                # the instance type name.
-                # aws ec2 describe-instance-types --region us-east-1
-                # https://aws.amazon.com/ec2/instance-types/trn1/
-                acc_name = 'Trainium'
-                find_num_in_name = re.search(r'(\d+)xlarge',
-                                             row['InstanceType'])
-                assert find_num_in_name is not None, row['InstanceType']
-                num_in_name = find_num_in_name.group(1)
-                acc_count = int(num_in_name) // 2
             if row['InstanceType'] == 'p5en.48xlarge':
                 # TODO(andyl): Check if this workaround still needed after
                 # v0.10.0 released. Currently, the acc_name returned by the
                 # AWS API is 'NVIDIA', which is incorrect. See #4652.
                 acc_name = 'H200'
                 acc_count = 8
+            if (row['InstanceType'].startswith('g6f') or
+                    row['InstanceType'].startswith('gr6f')):
+                # These instance actually have only fractional GPUs, but the API
+                # returns Count: 1 or Count: 0 under GpuInfo. We need to
+                # directly check the GPU memory to get the actual fraction of
+                # the GPU. Note that TotalGpuMemoryInMiB seems unreliable here -
+                # sometimes it is unexpectedly 0.
+                # See also Standard_NV{vcpu}ads_A10_v5 support on Azure.
+                assert len(row['GpuInfo']['Gpus']) == 1
+                assert row['GpuInfo']['Gpus'][0]['Name'] == 'L4'
+                fraction = row['GpuInfo']['Gpus'][0]['MemoryInfo'][
+                    'SizeInMiB'] / L4_GPU_MEMORY
+                acc_count = round(fraction, 3)
+            if row['InstanceType'] == 'p5.4xlarge':
+                acc_count = 1
             return pd.Series({
                 'AcceleratorName': acc_name,
                 'AcceleratorCount': acc_count,
                 'vCPUs': get_vcpus(row),
                 'MemoryGiB': get_memory_gib(row),
+                'Arch': get_arch(row),
             })
 
         # The AWS API may not have all the instance types in the pricing table,
@@ -341,11 +328,21 @@ def _get_instance_types_df(region: str) -> Union[str, 'pd.DataFrame']:
         df = pd.concat(
             [df, df.apply(get_additional_columns, axis='columns')],
             axis='columns')
-        # patch the df for p4de.24xlarge
-        if region in P4DE_REGIONS:
-            df = _patch_p4de(region, df, pricing_df)
         if 'GpuInfo' not in df.columns:
             df['GpuInfo'] = np.nan
+        if 'NeuronInfo' in df.columns:
+            # The AWS Neuron API uses 'NeuronDevices' instead of 'Gpus'
+            # in its dict; for consistency with GPU handling, rename key.
+            def map_neuroninfo(neuroninfo):
+                if isinstance(neuroninfo,
+                              dict) and 'NeuronDevices' in neuroninfo:
+                    # Rename 'NeuronDevices' to 'Gpus'
+                    neuroninfo = neuroninfo.copy()
+                    neuroninfo['Gpus'] = neuroninfo.pop('NeuronDevices')
+                return neuroninfo
+
+            df['NeuronInfo'] = df['NeuronInfo'].apply(map_neuroninfo)
+            df['GpuInfo'] = df['GpuInfo'].fillna(df['NeuronInfo'])
         df = df[USEFUL_COLUMNS]
     except Exception as e:  # pylint: disable=broad-except
         print(traceback.format_exc())
@@ -393,44 +390,70 @@ def get_all_regions_instance_types_df(regions: Set[str]) -> 'pd.DataFrame':
 # TODO(tian): find out the driver version.
 #   Neuron driver:
 _GPU_DESC_UBUNTU_DATE = [
-    ('gpu', 'AMI GPU PyTorch 2.1.0', '20.04', '20231103'),
-    ('gpu', 'AMI GPU PyTorch 1.10.0', '18.04', '20221114'),
-    ('k80', 'AMI GPU PyTorch 1.10.0', '20.04', '20211208'),
-    ('k80', 'AMI GPU PyTorch 1.10.0', '18.04', '20211208'),
-    ('neuron', 'Base Neuron AMI', '22.04', '20240923'),
+    ('neuron', '/aws/service/neuron/dlami/multi-framework', '22.04'),
 ]
 
 
-def _fetch_image_id(region: str, description: str, ubuntu_version: str,
-                    creation_date: str) -> Optional[str]:
+def _fetch_image_creation_date(region: str,
+                               image_id: Optional[str]) -> Optional[str]:
+    if image_id is None:
+        return None
     try:
         image = subprocess.check_output(f"""\
-            aws ec2 describe-images --region {region} --owners amazon \\
-                --filters 'Name=name,Values="Deep Learning {description} (Ubuntu {ubuntu_version}) {creation_date}"' \\
-                    'Name=state,Values=available' --query 'Images[:1].ImageId' --output text
+            aws ec2 describe-images --region {region} --image-ids {image_id} \\
+                --query 'Images[0].Name' --output text
             """,
                                         shell=True)
     except subprocess.CalledProcessError as e:
-        print(f'Failed {region}, {description}, {ubuntu_version}, '
-              f'{creation_date}. Trying next date.')
+        print(f'Failed to fetch image creation date for {region}, {image_id}')
         print(f'{type(e)}: {e}')
         image_id = None
+    else:
+        assert image is not None
+        image_name = image.decode('utf-8').strip()
+        match = re.search(r'(\d+)$', image_name)
+        if match:
+            return match.group(1)
+    return None
+
+
+def _fetch_image_id_from_ssm_param(
+        region: str,
+        ssm_prefix: str,
+        ubuntu_version: str = '22.04') -> Optional[str]:
+    try:
+        image = subprocess.check_output(f"""\
+            aws ssm get-parameter --region {region} --name "{ssm_prefix}/ubuntu-{ubuntu_version}/latest/image_id" \\
+                --query 'Parameter.Value' --output text
+            """,
+                                        shell=True)
+    except subprocess.CalledProcessError as e:
+        print(
+            f'Failed to fetch image ID from SSM parameter for {region}, {ssm_prefix}, {ubuntu_version}'
+        )
+        print(f'{type(e)}: {e}')
+        return None
     else:
         assert image is not None
         image_id = image.decode('utf-8').strip()
     return image_id
 
 
-def _get_image_row(region: str, gpu: str, description: str, ubuntu_version: str,
-                   date: str) -> Tuple[str, str, str, str, Optional[str], str]:
-    print(f'Getting image for {region}, {description}, {ubuntu_version}, {gpu}')
-    image_id = _fetch_image_id(region, description, ubuntu_version, date)
-    if image_id is None:
-        # not found
-        print(f'Failed to find image for {region}, {description}, '
-              f'{ubuntu_version}, {gpu}')
+def _get_image_row(
+    region: str,
+    gpu: str,
+    ssm_prefix: str,
+    ubuntu_version: str = '22.04'
+) -> Tuple[str, str, str, str, Optional[str], Optional[str]]:
+    print(f'Getting image for {region}, {ssm_prefix}, {ubuntu_version}, {gpu}')
+    image_id = _fetch_image_id_from_ssm_param(region, ssm_prefix,
+                                              ubuntu_version)
+    if image_id is not None:
+        creation_date = _fetch_image_creation_date(region, image_id)
+    else:
+        creation_date = None
     tag = f'skypilot:{gpu}-ubuntu-{ubuntu_version.replace(".", "")}'
-    return tag, region, 'ubuntu', ubuntu_version, image_id, date
+    return tag, region, 'ubuntu', ubuntu_version, image_id, creation_date
 
 
 def get_all_regions_images_df(regions: Set[str]) -> 'pd.DataFrame':
@@ -545,13 +568,26 @@ if __name__ == '__main__':
     instance_df.to_csv('aws/vms.csv', index=False)
     print('AWS Service Catalog saved to aws/vms.csv')
 
-    # Disable refreshing images.csv as we are using skypilot custom AMIs
-    # See sky/clouds/service_catalog/images/README.md for more details.
-    # image_df = get_all_regions_images_df(user_regions)
-    # _check_regions_integrity(image_df, 'images')
-
-    # image_df.to_csv('aws/images.csv', index=False)
-    # print('AWS Images saved to aws/images.csv')
+    # Disable refreshing images.csv for skypilot custom AMIs
+    # refresh only the neuron based images
+    # See sky/clouds/catalog/images/README.md for more details.
+    image_df = get_all_regions_images_df(user_regions)
+    _check_regions_integrity(image_df, 'images')
+    # filter out rows where ImageId is None
+    image_df = image_df[image_df['ImageId'].notna()]
+
+    # check if aws/images.csv exists
+    if os.path.exists('aws/images.csv'):
+        # load the data from aws/images.csv
+        existing_image_df = pd.read_csv('aws/images.csv')
+        # filter out the neuron based images
+        existing_image_df = existing_image_df[~existing_image_df['Tag'].
+                                              eq('skypilot:neuron-ubuntu-2204')]
+        # concat the new neuron based images with the existing images
+        image_df = pd.concat([existing_image_df, image_df])
+
+    image_df.to_csv('aws/images.csv', index=False)
+    print('AWS Images saved to aws/images.csv')
 
     if args.az_mappings:
         az_mappings_df = fetch_availability_zone_mappings()

sky/{clouds/service_catalog → catalog}/data_fetchers/fetch_cudo.py

@@ -9,18 +9,11 @@ import os
 
 import cudo_compute
 
-import sky.provision.cudo.cudo_utils as utils
+from sky.provision.cudo import cudo_utils as utils
 
 VMS_CSV = 'cudo/vms.csv'
 
 
-def cudo_api():
-    configuration = cudo_compute.Configuration()
-    configuration.host = 'https://rest.compute.cudo.org'
-    client = cudo_compute.ApiClient(configuration)
-    return cudo_compute.VirtualMachinesApi(client)
-
-
 def get_gpu_info(count, model):
     mem = utils.cudo_gpu_mem[model]
     # pylint: disable=line-too-long
@@ -45,39 +38,46 @@ def get_instance_type(machine_type, vcpu, mem, gpu):
         mem) + 'gb'
 
 
-def machine_types(gpu_model, mem_gib, vcpu_count, gpu_count):
-    try:
-        api = cudo_api()
-        types = api.list_vm_machine_types(mem_gib,
-                                          vcpu_count,
-                                          gpu=gpu_count,
-                                          gpu_model=gpu_model)
-        return types.to_dict()
-    except cudo_compute.rest.ApiException as e:
-        raise e
-
-
 def update_prices():
     rows = []
+
+    api = cudo_compute.cudo_api.virtual_machines()
+    all_types = api.list_vm_machine_types2()
+    all_machine_types = all_types.to_dict()['machine_types']
+
     for spec in utils.machine_specs:
-        mts = machine_types('', spec['mem'], spec['vcpu'], spec['gpu'])
-        for hc in mts['host_configs']:
-            if not utils.gpu_exists(hc['gpu_model']):
-                continue
-            accelerator_name = utils.cudo_gpu_to_skypilot_gpu(hc['gpu_model'])
-            row = {
-                'instance_type': get_instance_type(hc['machine_type'],
-                                                   spec['vcpu'], spec['mem'],
-                                                   spec['gpu']),
-                'accelerator_name': accelerator_name,
-                'accelerator_count': str(spec['gpu']) + '.0',
-                'vcpus': str(spec['vcpu']),
-                'memory_gib': str(spec['mem']),
-                'price': hc['total_price_hr']['value'],
-                'region': hc['data_center_id'],
-                'gpu_info': get_gpu_info(spec['gpu'], accelerator_name),
-            }
-            rows.append(row)
+        for machine_type in all_machine_types:
+            if (machine_type['min_vcpu'] <= spec['vcpu'] and
+                    machine_type['min_memory_gib'] <= spec['mem'] and
+                    utils.gpu_exists(machine_type['gpu_model'])):
+
+                accelerator_name = utils.cudo_gpu_to_skypilot_gpu(
+                    machine_type['gpu_model'])
+
+                # Calculate total price per hour based on the given spec
+                vcpu_price = float(
+                    machine_type['vcpu_price_hr']['value']) * spec['vcpu']
+                memory_price = float(
+                    machine_type['memory_gib_price_hr']['value']) * spec['mem']
+                gpu_price = float(
+                    machine_type['gpu_price_hr']['value']) * spec['gpu']
+                # Note: Not including storage and IPv4 prices
+                # for now as they may be optional
+                total_price = vcpu_price + memory_price + gpu_price
+
+                row = {
+                    'instance_type': get_instance_type(
+                        machine_type['machine_type'], spec['vcpu'], spec['mem'],
+                        spec['gpu']),
+                    'accelerator_name': accelerator_name,
+                    'accelerator_count': str(spec['gpu']) + '.0',
+                    'vcpus': str(spec['vcpu']),
+                    'memory_gib': str(spec['mem']),
+                    'price': str(total_price),
+                    'region': machine_type['data_center_id'],
+                    'gpu_info': get_gpu_info(spec['gpu'], accelerator_name),
+                }
+                rows.append(row)
     path = VMS_CSV
     with open(path, 'w', encoding='utf-8') as file:
         file.write(

sky/{clouds/service_catalog → catalog}/data_fetchers/fetch_gcp.py

@@ -179,9 +179,13 @@ TPU_V4_HOST_DF = pd.read_csv(
 # TODO(woosuk): Make this more robust.
 # Refer to: https://github.com/skypilot-org/skypilot/issues/1006
 # Unsupported Series: 'f1', 'm2'
-SERIES_TO_DISCRIPTION = {
+SERIES_TO_DESCRIPTION = {
     'a2': 'A2 Instance',
     'a3': 'A3 Instance',
+    # NOTE: GCP does not provide separate CPU/RAM pricing for A4 instances.
+    # The B200 GPU pricing includes the full VM cost. See special handling in
+    # get_vm_price() which sets A4 VM price to 0.
+    'a4': 'A4 Instance',
     'c2': 'Compute optimized',
     'c2d': 'C2D AMD Instance',
     'c3': 'C3 Instance',
@@ -195,9 +199,11 @@ SERIES_TO_DISCRIPTION = {
     'n1': 'N1 Predefined Instance',
     'n2': 'N2 Instance',
     'n2d': 'N2D AMD Instance',
+    'n4': 'N4 Instance',
     't2a': 'T2A Arm Instance',
     't2d': 'T2D AMD Instance',
 }
+
 creds, project_id = google.auth.default()
 gcp_client = discovery.build('compute', 'v1')
 tpu_client = discovery.build('tpu', 'v1')
@@ -334,7 +340,7 @@ def get_vm_df(skus: List[Dict[str, Any]], region_prefix: str) -> 'pd.DataFrame':
 
     # Drop the unsupported series.
     df = df[df['InstanceType'].str.startswith(
-        tuple(f'{series}-' for series in SERIES_TO_DISCRIPTION))]
+        tuple(f'{series}-' for series in SERIES_TO_DESCRIPTION))]
     df = df[~df['AvailabilityZone'].str.startswith(tuple(TPU_V4_ZONES))]
 
     # TODO(woosuk): Make this more efficient.
@@ -352,7 +358,7 @@ def get_vm_df(skus: List[Dict[str, Any]], region_prefix: str) -> 'pd.DataFrame':
 
             # Check if the SKU is for the correct series.
             description = sku['description']
-            if SERIES_TO_DISCRIPTION[series].lower() not in description.lower():
+            if SERIES_TO_DESCRIPTION[series].lower() not in description.lower():
                 continue
             # Special check for M1 instances.
             if series == 'm1' and 'M3' in description:
@@ -389,6 +395,15 @@ def get_vm_df(skus: List[Dict[str, Any]], region_prefix: str) -> 'pd.DataFrame':
         if series in ['f1', 'g1']:
             memory_price = 0.0
 
+        # Special case for A4 instances.
+        # GCP does not provide separate CPU/RAM pricing for A4 instances in the
+        # SKUs API. The GPU pricing (B200) includes the full VM cost.
+        # We set the VM price to 0 so the entry is not dropped, and the GPU
+        # pricing will provide the total cost.
+        if series == 'a4':
+            cpu_price = 0.0
+            memory_price = 0.0
+
         # TODO(tian): (2024/11/10) Some SKUs are missing in the SKUs API. We
         # skip them in the catalog for now. We should investigate why they are
         # missing and add them back.
@@ -434,10 +449,18 @@ def _get_gpus_for_zone(zone: str) -> 'pd.DataFrame':
             gpu_name = gpu_name.upper()
             if 'H100-80GB' in gpu_name:
                 gpu_name = 'H100'
-            if 'H100-MEGA-80GB' in gpu_name:
+
+            if 'H100-MEGA' in gpu_name:
                 gpu_name = 'H100-MEGA'
                 if count != 8:
-                    # H100-MEGA only has 8 cards.
+                    continue
+            elif 'H200' in gpu_name:
+                gpu_name = 'H200'
+                if count != 8:
+                    continue
+            elif 'B200' in gpu_name:
+                gpu_name = 'B200'
+                if count != 8:
                     continue
             if 'VWS' in gpu_name:
                 continue
@@ -468,6 +491,8 @@ def _gpu_info_from_name(name: str) -> Optional[Dict[str, List[Dict[str, Any]]]]:
         'A100': 40 * 1024,
         'H100': 80 * 1024,
         'H100-MEGA': 80 * 1024,
+        'H200': 141 * 1024,
+        'B200': 180 * 1024,
         'P4': 8 * 1024,
         'T4': 16 * 1024,
         'V100': 16 * 1024,
@@ -507,22 +532,47 @@ def get_gpu_df(skus: List[Dict[str, Any]],
         ondemand_or_spot = 'OnDemand' if not spot else 'Preemptible'
         gpu_price = None
         for sku in gpu_skus:
+            row_gpu_name = row['AcceleratorName']
             if row['Region'] not in sku['serviceRegions']:
                 continue
-            if sku['category']['usageType'] != ondemand_or_spot:
+
+            # Check usageType matches, with special handling for B200 spot.
+            # GCP has a bug where some B200 spot SKUs have usageType='OnDemand'
+            # but the description contains 'Spot Preemptible'.
+            usage_type = sku['category']['usageType']
+            description = sku['description']
+            is_spot_description = 'spot preemptible' in description.lower()
+
+            if usage_type != ondemand_or_spot:
+                # For B200 spot pricing, also accept SKUs where description
+                # says "Spot Preemptible" even if usageType is wrong.
+                if not (spot and row_gpu_name == 'B200' and
+                        is_spot_description):
+                    continue
+
+            # For B200 on-demand, skip SKUs that are actually spot (description
+            # says "Spot Preemptible" but usageType is incorrectly 'OnDemand').
+            if not spot and row_gpu_name == 'B200' and is_spot_description:
                 continue
 
-            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':
+            gpu_names = [f'{row_gpu_name} GPU']
+            if row_gpu_name == 'A100-80GB':
+                gpu_names = ['A100 80GB GPU']
+            elif row_gpu_name == 'H100':
+                gpu_names = ['H100 80GB GPU']
+            elif row_gpu_name == '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):
+                gpu_names = [
+                    'H100 80GB Mega GPU', 'H100 Mega 80GB GPU',
+                    'H100 80GB Plus GPU'
+                ]
+            elif row_gpu_name == 'H200':
+                gpu_names = ['H200 141GB GPU']
+            elif row_gpu_name == 'B200':
+                gpu_names = ['Nvidia B200 (1 gpu slice)']
+            if not any(
+                    gpu_name in sku['description'] for gpu_name in gpu_names):
                 continue
 
             unit_price = _get_unit_price(sku)
@@ -554,7 +604,7 @@ def _get_tpu_response_for_zone(zone: str) -> list:
     # Sometimes the response is empty ({}) even for enabled zones. Here we
     # retry the request for a few times.
     backoff = common_utils.Backoff(initial_backoff=1)
-    for _ in range(TPU_RETRY_CNT):
+    for retry_cnt in range(TPU_RETRY_CNT):
         tpus_request = (
             tpu_client.projects().locations().acceleratorTypes().list(
                 parent=parent))
@@ -570,6 +620,10 @@ def _get_tpu_response_for_zone(zone: str) -> list:
                 print(f'  An error occurred: {error}')
             # If error happens, fail early.
             return []
+        except TimeoutError:
+            print(f'  TimeoutError: Failed to fetch TPUs for zone {zone!r}, '
+                  f'retry {retry_cnt + 1} of {TPU_RETRY_CNT}')
+
         time_to_sleep = backoff.current_backoff()
         print(f'  Retry zone {zone!r} in {time_to_sleep} seconds...')
         time.sleep(time_to_sleep)