Add support for MTC phase-2 drivers.

Author: abhinavclemson

xlml/apis/task.py

@@ -173,6 +173,7 @@ def run(
       use_pathways: bool = False,
       skip_post_process: bool = False,
       ramdisk_directory: str = "",
+
   ) -> DAGNode:
     """Run a test job within a docker image.
 
@@ -248,6 +249,7 @@ def run_model(
       use_vertex_tensorboard: bool = False,
       use_pathways: bool = False,
       ramdisk_directory: str = "",
+      mtc_enabled: bool = False,
   ) -> DAGNode:
     """Run the TPU/GPU test in `task_test_config` using xpk.
 
@@ -274,6 +276,7 @@ def run_model(
           use_vertex_tensorboard,
           use_pathways,
           ramdisk_directory,
+          mtc_enabled,
       )
       wait_for_workload_completion = xpk.wait_for_workload_completion.override(
           timeout=int(self.task_test_config.timeout.total_seconds()),
@@ -306,6 +309,7 @@ def launch_workload(
       use_vertex_tensorboard: bool,
       use_pathways: bool = False,
       ramdisk_directory: str = "",
+      mtc_enabled: bool = False,
   ) -> DAGNode:
     """Create the workload and wait for it to provision."""
     with TaskGroup(group_id="launch_workload") as group:
@@ -326,6 +330,7 @@ def launch_workload(
           use_vertex_tensorboard=use_vertex_tensorboard,
           use_pathways=use_pathways,
           ramdisk_directory=ramdisk_directory,
+          mtc_enabled=mtc_enabled,
       )
       wait_for_workload_start = xpk.wait_for_workload_start.override(
           timeout=self.workload_provision_timeout.total_seconds()

xlml/utils/xpk.py

@@ -77,6 +77,7 @@ def run_workload(
     use_vertex_tensorboard: bool = False,
     use_pathways: bool = False,
     ramdisk_directory: str = "",  # Directory for enabling emergency checkpointing
+    mtc_enabled: bool = False, # It enables MTC phase-2 drivers
 ):
   """Run workload through xpk tool."""
 
@@ -103,6 +104,8 @@ def run_workload(
     )
     if ramdisk_directory:
       workload_create_cmd += f" --ramdisk-directory={ramdisk_directory}"
+    if mtc_enabled:
+      workload_create_cmd += " --mtc-enabled"
     cmds = get_xpk_setup_cmd(tmpdir)
     if accelerator_type == GpuVersion.XPK_H100_MEGA.value:
       workload_create_cmd += " --scheduler=gke.io/topology-aware-auto"