From b3ff0f280240d7f6a569f558bcd3f38a4e1053bf Mon Sep 17 00:00:00 2001 From: cblmemo Date: Wed, 30 Oct 2024 15:13:38 -0700 Subject: [PATCH 1/2] [Catalog] Add TPU V6e. --- .../data_fetchers/fetch_gcp.py | 52 ++++++++++++++++++- 1 file changed, 50 insertions(+), 2 deletions(-) diff --git a/sky/clouds/service_catalog/data_fetchers/fetch_gcp.py b/sky/clouds/service_catalog/data_fetchers/fetch_gcp.py index 6550c6bbe64..fb7af48a420 100644 --- a/sky/clouds/service_catalog/data_fetchers/fetch_gcp.py +++ b/sky/clouds/service_catalog/data_fetchers/fetch_gcp.py @@ -47,6 +47,10 @@ TPU_V4_ZONES = ['us-central2-b'] # TPU v3 pods are available in us-east1-d, but hidden in the skus. # We assume the TPU prices are the same as us-central1. +# TPU v6e's pricing info is not available on the SKUs. However, in +# https://cloud.google.com/tpu/pricing, it listed the price for 4 regions: +# us-east1, us-east5, europe-west4, and asia-northeast1. We hardcode them here +# and filtered out the other regions (us-central{1,2}, us-south1). HIDDEN_TPU_DF = pd.read_csv( io.StringIO( textwrap.dedent("""\ @@ -58,8 +62,50 @@ ,tpu-v3-512,1,,,tpu-v3-512,512.0,153.6,us-east1,us-east1-d ,tpu-v3-1024,1,,,tpu-v3-1024,1024.0,307.2,us-east1,us-east1-d ,tpu-v3-2048,1,,,tpu-v3-2048,2048.0,614.4,us-east1,us-east1-d + ,tpu-v6e-1,1,,,tpu-v6e-1,2.7,,us-east5,us-east5-b + ,tpu-v6e-1,1,,,tpu-v6e-1,2.7,,us-east5,us-east5-c + ,tpu-v6e-1,1,,,tpu-v6e-1,2.97,,europe-west4,europe-west4-a + ,tpu-v6e-1,1,,,tpu-v6e-1,3.24,,asia-northeast1,asia-northeast1-b + ,tpu-v6e-1,1,,,tpu-v6e-1,2.7,,us-east1,us-east1-d + ,tpu-v6e-4,1,,,tpu-v6e-4,10.8,,us-east5,us-east5-b + ,tpu-v6e-4,1,,,tpu-v6e-4,10.8,,us-east5,us-east5-c + ,tpu-v6e-4,1,,,tpu-v6e-4,11.88,,europe-west4,europe-west4-a + ,tpu-v6e-4,1,,,tpu-v6e-4,12.96,,asia-northeast1,asia-northeast1-b + ,tpu-v6e-4,1,,,tpu-v6e-4,10.8,,us-east1,us-east1-d + ,tpu-v6e-8,1,,,tpu-v6e-8,21.6,,us-east5,us-east5-b + ,tpu-v6e-8,1,,,tpu-v6e-8,21.6,,us-east5,us-east5-c + ,tpu-v6e-8,1,,,tpu-v6e-8,23.76,,europe-west4,europe-west4-a + ,tpu-v6e-8,1,,,tpu-v6e-8,25.92,,asia-northeast1,asia-northeast1-b + ,tpu-v6e-8,1,,,tpu-v6e-8,21.6,,us-east1,us-east1-d + ,tpu-v6e-16,1,,,tpu-v6e-16,43.2,,us-east5,us-east5-b + ,tpu-v6e-16,1,,,tpu-v6e-16,43.2,,us-east5,us-east5-c + ,tpu-v6e-16,1,,,tpu-v6e-16,47.52,,europe-west4,europe-west4-a + ,tpu-v6e-16,1,,,tpu-v6e-16,51.84,,asia-northeast1,asia-northeast1-b + ,tpu-v6e-16,1,,,tpu-v6e-16,43.2,,us-east1,us-east1-d + ,tpu-v6e-32,1,,,tpu-v6e-32,86.4,,us-east5,us-east5-b + ,tpu-v6e-32,1,,,tpu-v6e-32,86.4,,us-east5,us-east5-c + ,tpu-v6e-32,1,,,tpu-v6e-32,95.04,,europe-west4,europe-west4-a + ,tpu-v6e-32,1,,,tpu-v6e-32,103.68,,asia-northeast1,asia-northeast1-b + ,tpu-v6e-32,1,,,tpu-v6e-32,86.4,,us-east1,us-east1-d + ,tpu-v6e-64,1,,,tpu-v6e-64,172.8,,us-east5,us-east5-b + ,tpu-v6e-64,1,,,tpu-v6e-64,172.8,,us-east5,us-east5-c + ,tpu-v6e-64,1,,,tpu-v6e-64,190.08,,europe-west4,europe-west4-a + ,tpu-v6e-64,1,,,tpu-v6e-64,207.36,,asia-northeast1,asia-northeast1-b + ,tpu-v6e-64,1,,,tpu-v6e-64,172.8,,us-east1,us-east1-d + ,tpu-v6e-128,1,,,tpu-v6e-128,345.6,,us-east5,us-east5-b + ,tpu-v6e-128,1,,,tpu-v6e-128,345.6,,us-east5,us-east5-c + ,tpu-v6e-128,1,,,tpu-v6e-128,380.16,,europe-west4,europe-west4-a + ,tpu-v6e-128,1,,,tpu-v6e-128,414.72,,asia-northeast1,asia-northeast1-b + ,tpu-v6e-128,1,,,tpu-v6e-128,345.6,,us-east1,us-east1-d + ,tpu-v6e-256,1,,,tpu-v6e-256,691.2,,us-east5,us-east5-b + ,tpu-v6e-256,1,,,tpu-v6e-256,691.2,,us-east5,us-east5-c + ,tpu-v6e-256,1,,,tpu-v6e-256,760.32,,europe-west4,europe-west4-a + ,tpu-v6e-256,1,,,tpu-v6e-256,829.44,,asia-northeast1,asia-northeast1-b + ,tpu-v6e-256,1,,,tpu-v6e-256,691.2,,us-east1,us-east1-d """))) +TPU_V6E_MISSING_REGIONS = ['us-central1', 'us-central2', 'us-south1'] + # TPU V5 is not visible in specific zones. We hardcode the missing zones here. # NOTE(dev): Keep the zones and the df in sync. TPU_V5_MISSING_ZONES_DF = { @@ -683,11 +729,13 @@ def get_tpu_price(row: pd.Series, spot: bool) -> Optional[float]: 'not found in SKUs or hidden TPU price DF.') # TODO(tian): Hack. Should investigate how to retrieve the price # for TPU-v6e. - if not tpu_name.startswith('tpu-v6e'): + if not (tpu_name.startswith('tpu-v6e') and + tpu_region in TPU_V6E_MISSING_REGIONS): assert spot or tpu_price is not None, (row, hidden_tpu, HIDDEN_TPU_DF) else: - tpu_price = 0.0 + if not spot: + tpu_price = 0.0 return tpu_price df['Price'] = df.apply(lambda row: get_tpu_price(row, spot=False), axis=1) From 52cf39d7e9ee42257c0c5ab6ee99453d8bd39829 Mon Sep 17 00:00:00 2001 From: cblmemo Date: Wed, 30 Oct 2024 15:21:16 -0700 Subject: [PATCH 2/2] swap if else branch --- sky/clouds/service_catalog/data_fetchers/fetch_gcp.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/sky/clouds/service_catalog/data_fetchers/fetch_gcp.py b/sky/clouds/service_catalog/data_fetchers/fetch_gcp.py index fb7af48a420..8cc9fc6c127 100644 --- a/sky/clouds/service_catalog/data_fetchers/fetch_gcp.py +++ b/sky/clouds/service_catalog/data_fetchers/fetch_gcp.py @@ -729,13 +729,13 @@ def get_tpu_price(row: pd.Series, spot: bool) -> Optional[float]: 'not found in SKUs or hidden TPU price DF.') # TODO(tian): Hack. Should investigate how to retrieve the price # for TPU-v6e. - if not (tpu_name.startswith('tpu-v6e') and + if (tpu_name.startswith('tpu-v6e') and tpu_region in TPU_V6E_MISSING_REGIONS): - assert spot or tpu_price is not None, (row, hidden_tpu, - HIDDEN_TPU_DF) - else: if not spot: tpu_price = 0.0 + else: + assert spot or tpu_price is not None, (row, hidden_tpu, + HIDDEN_TPU_DF) return tpu_price df['Price'] = df.apply(lambda row: get_tpu_price(row, spot=False), axis=1)