[Catalog] Add TPU V6e.

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)