Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Catalog] Add TPU V6e. #4218

Merged
merged 2 commits into from
Oct 30, 2024
Merged

[Catalog] Add TPU V6e. #4218

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b3ff0f2
[Catalog] Add TPU V6e.
cblmemo Oct 30, 2024
52cf39d
swap if else branch
cblmemo Oct 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 51 additions & 3 deletions sky/clouds/service_catalog/data_fetchers/fetch_gcp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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("""\
Expand All @@ -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 = {
Expand Down Expand Up @@ -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 (tpu_name.startswith('tpu-v6e') and
tpu_region in TPU_V6E_MISSING_REGIONS):
if not spot:
tpu_price = 0.0
else:
assert spot or tpu_price is not None, (row, hidden_tpu,
HIDDEN_TPU_DF)
else:
tpu_price = 0.0
return tpu_price

df['Price'] = df.apply(lambda row: get_tpu_price(row, spot=False), axis=1)
Expand Down
Loading