Fix TPU head resource name for v4 and v5e (#165)

richardsliu · web-flow · commit f49451449dff · 2024-08-01T14:34:31.000-07:00
* Fix TPU head resource name for v4 and v5e

* fix format
diff --git a/run_ray_serve_interleave.py b/run_ray_serve_interleave.py
@@ -40,7 +40,11 @@
 
 
 def create_head_resource_name(generation, tpu_chips):
-  return f"TPU-{generation}-{tpu_chips}-head"
+  if generation == "v5litepod":
+    return f"TPU-{generation}-{tpu_chips}-head"
+  else:
+    tpu_cores = tpu_chips * 2
+    return f"TPU-{generation}-{tpu_cores}-head"
 
 
 def create_engine(**kwargs):
