Merge pull request #1174 from Edwinhr716:lws-integration

GitOrigin-RevId: 6fe455e23b395d623898e702e0449bcd38fccf00

Author: changlan

axlearn/cloud/gcp/job.py

@@ -19,7 +19,13 @@
 from axlearn.cloud.common.utils import generate_job_name, subprocess_run
 from axlearn.cloud.gcp.config import default_env_id, default_project, default_zone
 from axlearn.cloud.gcp.jobset_utils import BaseReplicatedJob
-from axlearn.cloud.gcp.utils import custom_jobset_kwargs, delete_k8s_jobset
+from axlearn.cloud.gcp.lws_utils import BaseLeaderWorkerTemplate
+from axlearn.cloud.gcp.utils import (
+    custom_jobset_kwargs,
+    custom_leaderworkerset_kwargs,
+    delete_k8s_jobset,
+    delete_k8s_leaderworkerset,
+)
 from axlearn.common.config import REQUIRED, ConfigOr, Required, config_class, maybe_instantiate
 from axlearn.common.utils import Nested
 
@@ -267,3 +273,103 @@ def docker_command(
     )
     logging.debug("Docker run command: %s", cmd)
     return cmd
+
+
+class GKELeaderWorkerSet(GCPJob):
+    """Base GKE LeaderWorkerSet interface"""
+
+    @config_class
+    class Config(GCPJob.Config):
+        """Configures GKELeaderWorkerSet.
+        Attributes:
+            builder: A builder that returns one or more statefulset specs.
+            namespace: The namespace to use within the k8s cluster.
+            annotations: LeaderWorkerSet annotations.
+            num_replicas: number of LWS replicas.
+        """
+
+        builder: Required[BaseLeaderWorkerTemplate.Config] = REQUIRED
+        namespace: str = "default"
+        annotations: Optional[ConfigOr[dict]] = None
+        num_replicas: int = 1
+
+    @classmethod
+    def set_defaults(cls, fv):
+        super().set_defaults(fv)
+        fv.set_default("max_tries", fv.max_tries or 10)
+        fv.set_default("retry_interval", fv.retry_interval or 60)
+
+    @classmethod
+    def define_flags(cls, fv: flags.FlagValues):
+        super().define_flags(fv)
+        common_kwargs = dict(flag_values=fv, allow_override=True)
+        flags.DEFINE_string("name", None, "Name of the LeaderWorkerSet.", **common_kwargs)
+
+    @classmethod
+    def from_flags(cls, fv: flags.FlagValues, **kwargs):
+        cfg: GKELeaderWorkerSet.Config = super().from_flags(fv, **kwargs)
+        cfg.num_replicas = fv.num_replicas
+        return cfg
+
+    def __init__(self, cfg: Config, *, bundler: BaseDockerBundler):
+        super().__init__(cfg)
+        cfg: GKELeaderWorkerSet.Config = self.config
+        self._bundler = bundler
+        # This instantiatees a builder for constructing replicated job specs, which will be managed
+        # together under the leaderworkerset represented by this class.
+        # Note the distinction from bundlers, which are responsible for bundling any code assets
+        # required to run the job.
+        self._builder: BaseLeaderWorkerTemplate = cfg.builder.instantiate(bundler=bundler)
+
+    def _delete(self):
+        cfg: GKELeaderWorkerSet.Config = self.config
+        # Issues a delete request for the LeaderWorkerSet and proactively delete its descendants.
+        # This is not fully blocking; after the call returns there can be a delay before
+        # everything is deleted.
+        delete_k8s_leaderworkerset(cfg.name, namespace=cfg.namespace)
+
+    def _build_leaderworkerset(self) -> Nested[Any]:
+        """
+        Builds a config for a LeaderWorkerSet, which is a set for multi-host inference
+
+        Returns:
+            A nested dict corresponding to a k8s LWS config
+        """
+        cfg: GKELeaderWorkerSet.Config = self.config
+        annotations = maybe_instantiate(cfg.annotations or {})
+
+        return dict(
+            metadata=dict(name=cfg.name, annotations=annotations),
+            spec=dict(
+                replicas=cfg.num_replicas,
+                leaderWorkerTemplate=self._builder(),
+            ),
+        )
+
+    def _execute(self):
+        cfg: GKELeaderWorkerSet.Config = self.config
+
+        api_kwargs = custom_leaderworkerset_kwargs()
+        custom_object = dict(
+            apiVersion=f"{api_kwargs['group']}/{api_kwargs['version']}",
+            kind="LeaderWorkerSet",
+            **self._build_leaderworkerset(),
+        )
+        logging.info("submitting LeaderWorkerSet: %s", custom_object)
+        return k8s.client.CustomObjectsApi().create_namespaced_custom_object(
+            namespace=cfg.namespace,
+            body=custom_object,
+            **api_kwargs,
+        )
+
+
+def exclusive_topology_annotations_leaderworkerset() -> dict:
+    """Used for TPU GKELeaderWorkerSet.
+
+    The exclusive topology annotation will ensure that all Pods will have affinity
+    rules added that will ensure that they are fully scheduled on the same pod-slice
+    node-pools.
+    """
+    return {
+        "leaderworkerset.sigs.k8s.io/subgroup-exclusive-topology": "cloud.google.com/gke-nodepool"
+    }

axlearn/cloud/gcp/job_test.py

@@ -11,7 +11,7 @@
 
 from axlearn.cloud.common.bundler import Bundler
 from axlearn.cloud.common.utils import define_flags, from_flags
-from axlearn.cloud.gcp import bundler, job, jobset_utils
+from axlearn.cloud.gcp import bundler, job, jobset_utils, pathways_utils
 from axlearn.cloud.gcp.bundler import ArtifactRegistryBundler, CloudBuildBundler
 from axlearn.cloud.gcp.test_utils import default_mock_settings, mock_gcp_settings
 from axlearn.common.config import REQUIRED, Required, config_class
@@ -211,3 +211,90 @@ def test_build_jobset(
             self.assertNotIn("kueue.x-k8s.io/queue-name", jobset_annotations)
         else:
             self.assertEqual(jobset_annotations["kueue.x-k8s.io/queue-name"], queue)
+
+
+class TPUGKELeaderWorkerSetTest(TestCase):
+    """Tests GKELeaderWorkerSet with TPU."""
+
+    def run(self, result=None):
+        # Run tests under mock user and settings.
+        self._settings = default_mock_settings()
+        with mock_gcp_settings(
+            [jobset_utils.__name__, bundler.__name__],
+            settings=self._settings,
+        ):
+            return super().run(result)
+
+    def _job_config(
+        self,
+        *,
+        command: str,
+        bundler_cls: type[Bundler],
+        **kwargs,
+    ) -> tuple[job.GKELeaderWorkerSet.Config, Bundler.Config]:
+        fv = flags.FlagValues()
+        cfg = job.GKELeaderWorkerSet.default_config().set(
+            builder=pathways_utils.PathwaysLeaderWorkerTemplate.default_config()
+        )
+        define_flags(cfg, fv)
+        for key, value in kwargs.items():
+            if value is not None:
+                # Use setattr rather than set_default to set flags.
+                setattr(fv, key, value)
+        fv.name = "fake-name"
+        fv.output_dir = "FAKE"
+        fv.instance_type = "tpu-v4-8"
+        fv.mark_as_parsed()
+        from_flags(cfg, fv, command=command)
+        # Test that retries are configured on fv by default.
+        self.assertIsNotNone(fv["max_tries"].default)
+        self.assertIsNotNone(fv["retry_interval"].default)
+        bundler_cfg = bundler_cls.from_spec([], fv=fv).set(image="test-image")
+        return cfg, bundler_cfg
+
+    @parameterized.product(
+        reservation=[None, "test"],
+        bundler_cls=[ArtifactRegistryBundler, CloudBuildBundler],
+        wrap_bundler=[False, True],
+    )
+    def test_instantiate(
+        self,
+        reservation,
+        bundler_cls: type[Bundler],
+        wrap_bundler,
+    ):
+        class WrappedBundler(Bundler):
+            @config_class
+            class Config(Bundler.Config):
+                inner: Required[Bundler.Config] = REQUIRED
+
+        cfg, bundler_cfg = self._job_config(
+            command="test-command",
+            bundler_cls=bundler_cls,
+            reservation=reservation,
+            num_replicas=1,
+        )
+
+        self.assertIsInstance(cfg.builder, pathways_utils.PathwaysLeaderWorkerTemplate.Config)
+        cfg.builder = cast(pathways_utils.PathwaysLeaderWorkerTemplate.Config, cfg.builder)
+
+        self.assertEqual(cfg.name, cfg.builder.name)
+        self.assertEqual(cfg.project, self._settings["project"])
+        self.assertEqual(cfg.zone, self._settings["zone"])
+        self.assertEqual(
+            cfg.builder.inner.reservation, reservation or self._settings["gke_reservation"]
+        )
+        self.assertEqual(cfg.num_replicas, 1)
+        # Should work with wrapped bundlers.
+        if wrap_bundler:
+            bundler_cfg = WrappedBundler.default_config().set(inner=bundler_cfg)
+        gke_job = cfg.instantiate(bundler=bundler_cfg.instantiate())
+        self.assertEqual("v4-8", gke_job._builder._tpu_type)
+
+    def test_delete(self):
+        patch_delete = mock.patch(f"{job.__name__}.delete_k8s_leaderworkerset")
+        with patch_delete as mock_delete:
+            cfg, _ = self._job_config(command="test-command", bundler_cls=CloudBuildBundler)
+            gke_job = cfg.instantiate(bundler=mock.Mock())
+            gke_job._delete()  # pylint: disable=protected-access
+            mock_delete.assert_called()

axlearn/cloud/gcp/jobset_utils.py

@@ -308,12 +308,12 @@ def from_flags(cls, fv: flags.FlagValues, **kwargs):
         return cfg
 
 
-class TPUReplicatedJob(SingleReplicatedJob):
-    """Builds a replicated jobspec for TPU, to be used with JobSet API."""
+class TPUJobBuilder(SingleReplicatedJob):
+    """Common base class for TPU Specs"""
 
     @config_class
     class Config(SingleReplicatedJob.Config):
-        """Configures TPUReplicatedJob.
+        """Configures TPUJobBuilder.
 
         Attributes:
             reservation: If specified, the TPU reservation name. This is not necessarily specific to
@@ -380,7 +380,7 @@ def define_flags(cls, fv: flags.FlagValues):
 
     @classmethod
     def from_flags(cls, fv: flags.FlagValues, **kwargs) -> Config:
-        cfg: TPUReplicatedJob.Config = super().from_flags(fv, **kwargs)
+        cfg: TPUJobBuilder.Config = super().from_flags(fv, **kwargs)
         default_env = get_default_env(
             tpu_type=infer_tpu_type(fv.instance_type),
             num_tpu_slices=fv.num_replicas,
@@ -404,7 +404,7 @@ def from_flags(cls, fv: flags.FlagValues, **kwargs) -> Config:
 
     def __init__(self, cfg: Config, *, bundler: Bundler):
         super().__init__(cfg, bundler=bundler)
-        cfg: TPUReplicatedJob.Config = self.config
+        cfg: TPUJobBuilder.Config = self.config
         if cfg.output_dir is None:
             raise ValueError("cfg.output_dir is required.")
         self._tpu_type = infer_tpu_type(cfg.accelerator.instance_type)
@@ -433,7 +433,7 @@ def _build_container(self) -> Nested[Any]:
         Returns:
             A nested dict corresponding to a k8s Container config.
         """
-        cfg: TPUReplicatedJob.Config = self.config
+        cfg: TPUJobBuilder.Config = self.config
         system = USER_FACING_NAME_TO_SYSTEM_CHARACTERISTICS[self._tpu_type]
         volume_mounts = [self._output_volume_mount]
 
@@ -503,7 +503,7 @@ def _build_uploader_container(
         Returns:
             A nested dict corresponding to a k8s Container config.
         """
-        cfg: TPUReplicatedJob.Config = self.config
+        cfg: TPUJobBuilder.Config = self.config
         output_volume_mount = output_volume_mount or self._output_volume_mount
         dst = f"{cfg.output_dir}/output/$HOSTNAME/"
         interval_s = 60
@@ -538,7 +538,7 @@ def _build_pod(self) -> Nested[Any]:
         Returns:
             A nested dict corresponding to a k8s Pod template, including the pod metadata and spec.
         """
-        cfg: TPUReplicatedJob.Config = self.config
+        cfg: TPUJobBuilder.Config = self.config
         system = USER_FACING_NAME_TO_SYSTEM_CHARACTERISTICS[self._tpu_type]
         annotations, labels, selector, volumes, tolerations = {}, {}, {}, [], []
 
@@ -727,6 +727,12 @@ def _build_pod(self) -> Nested[Any]:
             spec=spec,
         )
 
+
+class TPUReplicatedJob(TPUJobBuilder):
+    """Builds a replicated job spec for a generic TPU job to be used with the JobSet API"""
+
+    Config = TPUJobBuilder.Config
+
     def __call__(self) -> Sequence[Nested[Any]]:
         """See `BaseReplicatedJob` docstring for details."""
         cfg: TPUReplicatedJob.Config = self.config

axlearn/cloud/gcp/lws_utils.py

@@ -0,0 +1,102 @@
+# Copyright © 2025 Apple Inc.
+
+"""Utilities for building LeaderWorkerSet specs"""
+
+from typing import Any, Optional, Sequence
+
+from absl import flags
+
+from axlearn.cloud.common.bundler import Bundler
+from axlearn.cloud.common.utils import AcceleratorConfig, FlagConfigurable, accelerator_flags
+from axlearn.cloud.gcp.config import gcp_settings
+from axlearn.cloud.gcp.jobset_utils import TPUJobBuilder
+from axlearn.cloud.gcp.system_characteristics import USER_FACING_NAME_TO_SYSTEM_CHARACTERISTICS
+from axlearn.common.config import REQUIRED, Required, config_class
+from axlearn.common.utils import Nested
+
+
+class BaseLeaderWorkerTemplate(FlagConfigurable):
+    """
+    Common base class for LeaderWorker Templates
+    """
+
+    @config_class
+    class Config(FlagConfigurable.Config):
+        """
+        Configures BaseLeaderWorker.
+        Attributes:
+            name: Name of the LeaderWorkerSet
+            command: Command to be executed.
+            accelerator: Accelerator configuration.
+            env_vars: Optional env vars to set.
+            service_account: Optional service account to execute the job as.
+            output_dir: An optional GCS path to upload LWS outputs to.
+        """
+
+        name: Required[str] = REQUIRED
+        # TODO: Change this to be a list of str[], to support different commands
+        # between leader and workers
+        command: Required[str] = REQUIRED
+        accelerator: AcceleratorConfig = AcceleratorConfig()
+        env_vars: dict[str, str] = {}
+        service_account: Optional[str] = None
+        output_dir: Optional[str] = None
+
+    @classmethod
+    def define_flags(cls, fv):
+        super().define_flags(fv)
+        common_kwargs = dict(flag_values=fv, allow_override=True)
+        accelerator_flags(**common_kwargs)
+        # NOTE: the parent typically sets these flags, so we leave them as None.
+        flags.DEFINE_string("name", None, "Name of the LWS.", **common_kwargs)
+        flags.DEFINE_string("command", None, "Command to execute.", **common_kwargs)
+        flags.DEFINE_multi_string("env", [], "Env var in the format key:value.", **common_kwargs)
+        flags.DEFINE_string(
+            "service_account",
+            None,
+            "If specified, will run job as the service account.",
+            **common_kwargs,
+        )
+        flags.DEFINE_string(
+            "output_dir",
+            None,
+            "If specified, the directory to store outputs (such as logs).",
+            **common_kwargs,
+        )
+        flags.DEFINE_boolean(
+            "enable_pre_provisioner", None, "Whether to enable pre-provisioner.", **common_kwargs
+        )
+
+    @classmethod
+    def from_flags(cls, fv: flags.FlagValues, **kwargs):
+        cfg: BaseLeaderWorkerTemplate.Config = super().from_flags(fv, **kwargs)
+        cfg.service_account = cfg.service_account or gcp_settings(
+            "k8s_service_account", default="default", fv=fv
+        )
+        cfg.accelerator.set(instance_type=fv.instance_type, num_replicas=fv.num_replicas)
+        return cfg
+
+    def __init__(self, cfg: Config, *, bundler: Bundler):
+        super().__init__(cfg)
+        self._bundler = bundler
+
+    def __call__(self) -> Sequence[Nested[Any]]:
+        """Builds LeaderWorkerTemplate for the LWS API.
+
+        Returns:
+        A nested dict corresponding to a LeaderWorkerTemplate config.
+        """
+        raise NotImplementedError(type(self))
+
+
+class TPULeaderWorkerTemplate(TPUJobBuilder):
+    """Builds a LeaderWorkerTemplate spec for a generic TPU workload"""
+
+    Config = TPUJobBuilder.Config
+
+    def __call__(self) -> Sequence[Nested[Any]]:
+        system = USER_FACING_NAME_TO_SYSTEM_CHARACTERISTICS[self._tpu_type]
+        return dict(
+            size=system.vms_per_slice,
+            workerTemplate=self._build_pod(),
+        )