Support ct6e-standard-8t

GitOrigin-RevId: 217890c603e3b59b6f77ffd8141291d2884065b0

Author: muyangyuapple

axlearn/cloud/gcp/jobs/launch.py

@@ -402,7 +402,6 @@ def submit(self) -> JobSpec:
         """Submits the command to bastion."""
         cfg: BaseBastionManagedJob.Config = self.config
         self._bundler.bundle(cfg.name)
-
         logging.info("Starting run for job name %s", cfg.name)
         logging.info("Command: %s", cfg.command)
         with tempfile.NamedTemporaryFile("w") as f:

axlearn/cloud/gcp/jobset_utils.py

@@ -479,6 +479,7 @@ def _build_container(self) -> Nested[Any]:
         if cfg.enable_tpu_ici_resiliency is not None:
             env_vars["ENABLE_ICI_RESILIENCY"] = str(cfg.enable_tpu_ici_resiliency).lower()
 
+        # This label will be used by TPU provisioner to select machine type.
         resources = {"limits": {"google.com/tpu": system.chips_per_vm}}
         # Set request memory by host machine type.
         machine_memory_gi = GCE_MACHINE_TYPE_TO_MEMORY_CHARACTERISTICS.get(

axlearn/cloud/gcp/node_pool_provisioner.py

@@ -153,7 +153,15 @@ def create_for(self, job: GKEJob):
             additional_labels_list.append(additional_labels)
 
         start_time = time.perf_counter()
-        topology = None if isinstance(job, _INFERENCE_JOBS) else job_sys_property.topology
+        topology = job_sys_property.topology
+        if job_sys_property.gce_machine_type == "ct6e-standard-8t":
+            # If we customize chips_per_vm to use ct6e-standard-8t for v6e,
+            # it is required not to set topology.
+            topology = None
+        if isinstance(job, _INFERENCE_JOBS):
+            # Inference jobs like Flink/Beam jobs use node pool as single
+            # host nodes, we don't set topology for them
+            topology = None
         create_node_pools(
             node_pool_names,
             project=cfg.project,

axlearn/cloud/gcp/system_characteristics.py

@@ -395,6 +395,12 @@ class _SystemCharacteristics:
     "v6e-8": _SystemCharacteristics(
         "2x4", 2, "tpu-v6e-slice", "ct6e-standard-4t", 4, AcceleratorType["TPU"], "v6e-8"
     ),
+    # Naming convention for TPU: {version}-{cores}[-{variant}]
+    # "-{variant}" is optional. It is used to define a spec that is different from standard.
+    # The value can be anything as long as it is unique.
+    "v6e-8-1": _SystemCharacteristics(
+        "2x4", 1, "tpu-v6e-slice", "ct6e-standard-8t", 8, AcceleratorType["TPU"], "v6e-8"
+    ),
     "v6e-16": _SystemCharacteristics(
         "4x4", 4, "tpu-v6e-slice", "ct6e-standard-4t", 4, AcceleratorType["TPU"], "v6e-16"
     ),

axlearn/cloud/gcp/tpu.py

@@ -35,23 +35,29 @@ def infer_tpu_cores(tpu_type: str) -> int:
     """Infer the number of TPU cores from the TPU type.
 
     Args:
-        tpu_type: A string of the format {version}-{cores}.
+        tpu_type: A string of the format {version}-{cores}[-{variant}].
+        -variant is optional.
 
     Returns:
         Inferred number of TPU cores.
     """
+    if tpu_type.count("-") == 2:
+        tpu_type = tpu_type[: tpu_type.rfind("-")]
     return int(tpu_type.rsplit("-", 1)[1])
 
 
 def infer_tpu_workers(tpu_type: str) -> int:
     """Infer the number of worker processes for the given TPU type.
 
     Args:
-        tpu_type: A string of the format {version}-{cores}.
+        tpu_type: A string of the format {version}-{cores}[-{variant}].
+        -variant is optional.
 
     Returns:
         Inferred number of TPU workers.
     """
+    if tpu_type.count("-") == 2:
+        tpu_type = tpu_type[: tpu_type.rfind("-")]
     tpu_pattern = r"(.+)*-(\d+)"
     match = re.search(tpu_pattern, tpu_type)
     try:

axlearn/cloud/gcp/tpu_test.py

@@ -23,6 +23,7 @@ class TpuUtilsTest(parameterized.TestCase):
         dict(tpu_type="v4-128", version="v4", cores=128, workers=16),
         dict(tpu_type="v5litepod-16", version="v5litepod", cores=16, workers=4),
         dict(tpu_type="v3-64", version="v3", cores=64, workers=8),
+        dict(tpu_type="v6e-8-1", version="v6e", cores=8, workers=2),
     )
     def test_infer_utils(self, tpu_type: str, version: str, cores: int, workers: int):
         self.assertEqual(version, infer_tpu_version(tpu_type))

axlearn/common/compiler_options.py

@@ -263,7 +263,8 @@ class NotTpuError(ValueError):
 
 # TODO(markblee): Generalize to other accelerators.
 def infer_tpu_type(instance_type: str) -> str:
-    """Infers tpu type (e.g. v4-8) from instance type (e.g. tpu-v4-8 or v4-8)."""
+    """Infers tpu type (e.g. v4-8 or v6e-8-1) from instance type
+    (e.g. tpu-v4-8, v4-8, tpu-v6e-8-1 or v6e-8-1)."""
     if not (instance_type and re.fullmatch(r"(tpu-)?v.+-\d+", instance_type)):
         raise NotTpuError(f"Invalid TPU instance: {instance_type}")
     return instance_type.replace("tpu-", "")
@@ -282,6 +283,8 @@ def infer_tpu_version(tpu_type: str) -> str:
     Raises:
         ValueError: if the TPU version string is unknown.
     """
+    if tpu_type.count("-") == 2:
+        tpu_type = tpu_type[: tpu_type.rfind("-")]
     tpu_type = infer_tpu_type(tpu_type)
     tpu_version = tpu_type.rsplit("-", 1)[0]  # split from the last occurrence of '-'
     # Resolve aliases like v5e to v5litepod, since in some cases (e.g. aot compilation) v5e is

axlearn/common/compiler_options_test.py

@@ -88,6 +88,7 @@ def test_xsc_compiler_options(self):
 
     @parameterized.parameters(
         dict(tpu_type="v5e-16", expected="v5litepod"),
+        dict(tpu_type="v6e-8-1", expected="v6e"),
     )
     def test_tpu_version_alias(self, tpu_type: str, expected: str):
         self.assertEqual(expected, compiler_options.infer_tpu_version(tpu_type))