From c1464e14358f7eecb1c3b7e1de69b17e107a06fe Mon Sep 17 00:00:00 2001 From: Romil Bhardwaj Date: Fri, 13 Sep 2024 08:43:25 -0700 Subject: [PATCH] [TPUs] Update sky show-gpus for TPUs (#3941) * update TPU * update TPU * update TPU * lint --- sky/clouds/service_catalog/__init__.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/sky/clouds/service_catalog/__init__.py b/sky/clouds/service_catalog/__init__.py index acc6fa0aa8b..f2301bac466 100644 --- a/sky/clouds/service_catalog/__init__.py +++ b/sky/clouds/service_catalog/__init__.py @@ -337,10 +337,13 @@ def get_common_gpus() -> List[str]: def get_tpus() -> List[str]: """Returns a list of TPU names.""" # TODO(wei-lin): refactor below hard-coded list. + # There are many TPU configurations available, we show the three smallest + # and the largest configuration for the latest gen TPUs. return [ - 'tpu-v2-8', 'tpu-v2-32', 'tpu-v2-128', 'tpu-v2-256', 'tpu-v2-512', - 'tpu-v3-8', 'tpu-v3-32', 'tpu-v3-64', 'tpu-v3-128', 'tpu-v3-256', - 'tpu-v3-512', 'tpu-v3-1024', 'tpu-v3-2048' + 'tpu-v2-512', 'tpu-v3-2048', 'tpu-v4-8', 'tpu-v4-16', 'tpu-v4-32', + 'tpu-v4-3968', 'tpu-v5litepod-1', 'tpu-v5litepod-4', 'tpu-v5litepod-8', + 'tpu-v5litepod-256', 'tpu-v5p-8', 'tpu-v5p-32', 'tpu-v5p-128', + 'tpu-v5p-12288' ]