Add command rewrite support for LWS based jobs

GitOrigin-RevId: b2d4da4

Author: nataren

axlearn/cloud/gcp/jobset_utils.py

@@ -190,6 +190,9 @@ def define_flags(cls, fv):
 
     def __init__(self, cfg: Config, *, bundler: Bundler):
         super().__init__(cfg)
+        self._user_command_patcher = (
+            cfg.user_command_patcher.instantiate() if cfg.user_command_patcher else None
+        )
         self._bundler = bundler
 
     def __call__(self) -> Sequence[Nested[Any]]:
@@ -446,9 +449,6 @@ def from_flags(cls, fv: flags.FlagValues, **kwargs) -> Config:
 
     def __init__(self, cfg: Config, *, bundler: Bundler):
         super().__init__(cfg, bundler=bundler)
-        self._user_command_patcher = (
-            cfg.user_command_patcher.instantiate() if cfg.user_command_patcher else None
-        )
         cfg: TPUJobBuilder.Config = self.config
         if cfg.output_dir is None:
             raise ValueError("cfg.output_dir is required.")

axlearn/cloud/gcp/lws_utils.py

@@ -7,6 +7,7 @@
 from absl import flags
 
 from axlearn.cloud.common.bundler import Bundler
+from axlearn.cloud.common.user_command_patcher import UserCommandPatcher
 from axlearn.cloud.common.utils import FlagConfigurable
 from axlearn.cloud.gcp.config import gcp_settings
 from axlearn.cloud.gcp.jobset_utils import TPUJobBuilder
@@ -40,6 +41,7 @@ class Config(FlagConfigurable.Config):
         service_account: Optional[str] = None
         output_dir: Optional[str] = None
         image_id: Optional[str] = None
+        user_command_patcher: Optional[UserCommandPatcher.Config] = None
 
     @classmethod
     def define_flags(cls, fv):
@@ -79,6 +81,9 @@ def from_flags(cls, fv: flags.FlagValues, **kwargs):
     def __init__(self, cfg: Config, *, bundler: Bundler):
         super().__init__(cfg)
         self._bundler = bundler
+        self._user_command_patcher = (
+            cfg.user_command_patcher.instantiate() if cfg.user_command_patcher else None
+        )
 
     def __call__(self) -> Sequence[Nested[Any]]:
         """Builds LeaderWorkerTemplate for the LWS API.

axlearn/cloud/gcp/pathways_utils.py

@@ -3,13 +3,19 @@
 """Utilities for building Pathways Jobset specs."""
 
 import copy
+import io
 import logging
 import os
 from typing import Any, Optional, Sequence, Union
 
 from absl import flags
 
-from axlearn.cloud.common.bastion import BASTION_JOB_VERSION_ENV_VAR
+from axlearn.cloud.common.bastion import (
+    _BASTION_SERIALIZED_JOBSPEC_ENV_VAR,
+    BASTION_JOB_VERSION_ENV_VAR,
+    ValidationError,
+    deserialize_jobspec,
+)
 from axlearn.cloud.common.bundler import Bundler
 from axlearn.cloud.common.utils import parse_kv_flags
 from axlearn.cloud.gcp.jobset_utils import (
@@ -985,6 +991,8 @@ def __init__(self, cfg: Config, *, bundler):
 
         self._bundler = bundler
         self._inner: TPULeaderWorkerTemplate = cfg.inner.instantiate(bundler=self._bundler)
+        self._inner._user_command_patcher = self._user_command_patcher
+
         self._tpu_type = infer_tpu_type(cfg.inner.accelerator.instance_type)
         if self._tpu_type not in USER_FACING_NAME_TO_SYSTEM_CHARACTERISTICS:
             raise NotImplementedError(f"Missing system characteristics for {self._tpu_type}")
@@ -1102,10 +1110,23 @@ def _build_head_container(self) -> dict:
             "requests": {"cpu": cpu_req, "memory": mem_req},
             "limits": {"cpu": cpu_req, "memory": mem_req},
         }
+        command = inner_cfg.command
+
+        if self._user_command_patcher is not None:
+            user_id = None
+            job_spec = os.environ.get(_BASTION_SERIALIZED_JOBSPEC_ENV_VAR)
+            if job_spec:
+                try:
+                    spec = deserialize_jobspec(io.StringIO(job_spec))
+                    user_id = spec.metadata.user_id
+                except ValidationError:
+                    logging.debug("Failed to deserialize job spec for user command patching.")
+            command = self._user_command_patcher.patch(command, user_id=user_id)
+
         container = dict(
             name=inner_cfg.name,
             image=inner_cfg.image_id or self._bundler.id(inner_cfg.name),
-            command=["bash", "-c", inner_cfg.command],
+            command=["bash", "-c", command],
             env=[
                 {
                     "name": "XCLOUD_ENVIRONMENT",

axlearn/cloud/gcp/pathways_utils_test.py

@@ -57,6 +57,11 @@ def test_v6e_default_options_split_megascale_and_xla(self):
         self.assertEqual(len(megascale_options) + len(xla_options), len(default_options))
 
 
+class MockUserCommandPatcher(jobset_utils.UserCommandPatcher):
+    def patch(self, command: str, **kwargs: dict) -> str:
+        return f"./prefix_command.sh; {command}"
+
+
 class PathwaysReplicatedJobTest(TestCase):
     """Tests PathwaysReplicatedJob."""
 
@@ -408,18 +413,23 @@ def _job_config(self, bundler_cls: type[Bundler], num_replicas: int, **kwargs):
             bundler_cfg = bundler_cls.from_spec([], fv=fv).set(image="test-image")
             yield cfg, bundler_cfg
 
-    @parameterized.parameters([1, 2])
-    def test_replicated_job(self, num_replicas):
+    @parameterized.product(
+        num_replicas=[1, 2], user_command_patcher=[None, MockUserCommandPatcher.default_config()]
+    )
+    def test_replicated_job(self, num_replicas, user_command_patcher):
         with (self._job_config(CloudBuildBundler, num_replicas) as (cfg, bundler_cfg),):
+            command = "test_command"
             cfg.inner.set(
                 project="test-project",
                 name="test",
-                command="test_command",
+                command=command,
                 output_dir="FAKE",
-            ).instantiate(bundler=bundler_cfg.instantiate())
+            )
+            if user_command_patcher:
+                cfg.inner.set(user_command_patcher=user_command_patcher)
+            cfg.instantiate(bundler=bundler_cfg.instantiate())
 
             builder = cfg.instantiate(bundler=bundler_cfg.instantiate())
-
             replicated_jobs = builder()
 
             self.assertEqual(len(replicated_jobs), num_replicas + 1)
@@ -449,6 +459,29 @@ def test_replicated_job(self, num_replicas):
                 elif replicated_job_name.startswith("pwwk"):
                     self.assertEqual(replicated_job["replicas"], 1)
 
+                command_was_patched = False
+
+                if replicated_job_name.startswith("pwhd"):
+                    containers = (
+                        job_spec.get("spec", {})
+                        .get("template", {})
+                        .get("spec", {})
+                        .get("containers", [])
+                    )
+                    if containers and len(containers) > 0:
+                        command_parts = containers[0].get("command", [])
+                        if command_parts and len(command_parts) > 0:
+                            the_command = command_parts[-1]
+                            if user_command_patcher:
+                                self.assertEqual(f"./prefix_command.sh; {command}", the_command)
+                                command_was_patched = True
+                            else:
+                                self.assertEqual(command, the_command)
+                self.assertEqual(
+                    user_command_patcher is not None and replicated_job_name.startswith("pwhd"),
+                    command_was_patched,
+                )
+
     def test_validate_head_name(self):
         with self._job_config(CloudBuildBundler, 2) as (cfg, bundler_cfg):
             cfg.inner.set(