From 1a711dbcc880519ec0173c333c0546f8b6b24ef1 Mon Sep 17 00:00:00 2001 From: Vaibhav Agrawal Date: Thu, 11 Jan 2024 10:47:54 +0530 Subject: [PATCH] [UI] Now show-gpus shows the combined price for GCP GPU instance (#2946) --- sky/cli.py | 5 --- sky/clouds/service_catalog/gcp_catalog.py | 41 +++++++++++++---------- 2 files changed, 24 insertions(+), 22 deletions(-) diff --git a/sky/cli.py b/sky/cli.py index a66dffbee64..389476848bc 100644 --- a/sky/cli.py +++ b/sky/cli.py @@ -3543,11 +3543,6 @@ def _output(): yield 'to show available accelerators.' return - if cloud is None or cloud.lower() == 'gcp': - yield '*NOTE*: for most GCP accelerators, ' - yield 'INSTANCE_TYPE == (attachable) means ' - yield 'the host VM\'s cost is not included.\n\n' - import pandas as pd # pylint: disable=import-outside-toplevel for i, (gpu, items) in enumerate(result.items()): accelerator_table_headers = [ diff --git a/sky/clouds/service_catalog/gcp_catalog.py b/sky/clouds/service_catalog/gcp_catalog.py index efdeb6b4c14..e6a9c6a39d8 100644 --- a/sky/clouds/service_catalog/gcp_catalog.py +++ b/sky/clouds/service_catalog/gcp_catalog.py @@ -385,21 +385,29 @@ def list_accelerators( new_results[acc_name] = acc_info results = new_results - # Unlike other GPUs that can be attached to different sizes of N1 VMs, - # A100 GPUs can only be attached to fixed-size A2 VMs, - # and L4 GPUs can only be attached to G2 VMs. - # Thus, we can show their exact cost including the host VM prices. - - acc_infos: List[common.InstanceTypeInfo] = sum( - [results.get(a, []) for a in _ACC_INSTANCE_TYPE_DICTS], []) - if not acc_infos: - return results + # Figure out which instance type to use. + infos_with_instance_type: List[common.InstanceTypeInfo] = sum( + results.values(), []) new_infos = defaultdict(list) - for info in acc_infos: - assert pd.isna(info.instance_type) and pd.isna(info.memory), acc_infos - vm_types = _ACC_INSTANCE_TYPE_DICTS[info.accelerator_name][ - info.accelerator_count] + for info in infos_with_instance_type: + assert pd.isna(info.instance_type) and pd.isna(info.memory), info + # We show TPU-VM prices here, so no instance type is needed. + # Set it as `TPU-VM` to be shown in the table. + if info.accelerator_name.startswith('tpu'): + new_infos[info.accelerator_name].append( + info._replace(instance_type='TPU-VM')) + continue + vm_types, _ = get_instance_type_for_accelerator(info.accelerator_name, + info.accelerator_count, + region=region_filter) + # The acc name & count in `info` are retrieved from the table so we + # could definitely find a column in the original table + # Additionally the way get_instance_type_for_accelerator works + # we will always get either a specialized instance type + # or a default instance type. So we can safely assume that + # vm_types is not None. + assert vm_types is not None for vm_type in vm_types: df = _df[_df['InstanceType'] == vm_type] cpu_count = df['vCPUs'].iloc[0] @@ -407,12 +415,12 @@ def list_accelerators( vm_price = common.get_hourly_cost_impl(_df, vm_type, use_spot=False, - region=None, + region=region_filter, zone=None) vm_spot_price = common.get_hourly_cost_impl(_df, vm_type, use_spot=True, - region=None, + region=region_filter, zone=None) new_infos[info.accelerator_name].append( info._replace( @@ -423,8 +431,7 @@ def list_accelerators( price=info.price + vm_price, spot_price=info.spot_price + vm_spot_price, )) - results.update(new_infos) - return results + return new_infos def get_region_zones_for_accelerators(