lws integration

runner for lws added

added updatig condition

added spec for pathways in lws

updated to variables

refactored to match JobSet Pathways implementation

added new runner for jetstream-pathways

bug fixes, made num_replicas a parameter for an LWS object

ran precommits

added pathways utils tests

added more unit tests

removed mentions of jetstream, fixed bug in runner

added tests for gke runner

minor fixes

cleaned up tests

removed jetstream from runner name

addressed comments

test

service and PR changes

Update axlearn/cloud/gcp/job.py

Co-authored-by: Meng (Ethan) Li <[email protected]>

Update axlearn/cloud/gcp/job.py

Co-authored-by: Meng (Ethan) Li <[email protected]>

Update axlearn/cloud/gcp/job.py

Co-authored-by: Meng (Ethan) Li <[email protected]>

Update axlearn/cloud/gcp/job.py

Co-authored-by: Meng (Ethan) Li <[email protected]>

added flag

added patch

flatten logic, removed inner

cleanup

Leader on a seperate CPU node

changed to subgroup exclusive policy

added image

fixed node selector on head pod

added backend platforms env variable

added jax backend target variable

add rm port

fixed issue with proxy error

addressed comments

Removed service, will add on different PR

added env variable to fix handshake issue

Update node_pool_provisioner.py

fixed pre-commit error

refactored to add TPUJobBuilder, addressed comments

addressed final comments

added resource requirements

added worker image

get accelerator from inner

added lws test, refactored other tests

rebase

added builder type

changed builder name

changed builder name to add _

addressed comments

fix for list command

remove propagation policy from list command

fixed LWS runner test

Author: Edwinhr716

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(),
+        )