Add labels to LWS configs

GitOrigin-RevId: aa81577

Author: Weiwei Yang

axlearn/cloud/gcp/job.py

@@ -547,12 +547,14 @@ class Config(GCPJob.Config):
             builder: A builder that returns one or more statefulset specs.
             namespace: The namespace to use within the k8s cluster.
             annotations: LeaderWorkerSet annotations.
+            labels: LeaderWorkerSet labels.
             num_replicas: number of LWS replicas.
         """
 
         builder: Required[BaseLeaderWorkerTemplate.Config] = REQUIRED
         namespace: str = "default"
         annotations: Optional[ConfigOr[dict]] = None
+        labels: Optional[ConfigOr[dict]] = None
         num_replicas: int = 1
         enable_service: bool = False
         ports: list[str] = None
@@ -700,7 +702,7 @@ def _build_leaderworkerset(self) -> Nested[Any]:
         """
         cfg: GKELeaderWorkerSet.Config = self.config
         annotations = maybe_instantiate(cfg.annotations or {})
-        labels = {}
+        labels = maybe_instantiate(cfg.labels or {})
 
         # If the topology is set and slice auto provisioning is configured
         # set the necessary annotations

axlearn/cloud/gcp/job_test.py

@@ -654,3 +654,38 @@ def test_build_leaderworkerset(
             self.assertIn("spec", lws_spec)
             self.assertIn("replicas", lws_spec["spec"])
             self.assertIn("leaderWorkerTemplate", lws_spec["spec"])
+
+    @parameterized.product(
+        bundler_cls=[ArtifactRegistryBundler, CloudBuildBundler],
+        labels=[None, {"env": "tpu-test"}, {"team": "research", "experiment": "training"}],
+    )
+    def test_build_leaderworkerset_labels(
+        self,
+        bundler_cls,
+        labels: Optional[dict] = None,
+    ):
+        """Test that labels are properly set in LeaderWorkerSet metadata."""
+        cfg, bundler_cfg = self._job_config(
+            command="test-command",
+            bundler_cls=bundler_cls,
+        )
+        # Set labels on the config
+        cfg = cfg.set(labels=labels)
+        gke_job: job.GKELeaderWorkerSet = cfg.instantiate(bundler=bundler_cfg.instantiate())
+        # pylint: disable-next=protected-access
+        lws_spec = gke_job._build_leaderworkerset()
+        lws_metadata = lws_spec["metadata"]
+        lws_labels = lws_metadata.get("labels", {})
+
+        # Test basic metadata
+        self.assertEqual(lws_metadata["name"], cfg.name)
+
+        # Test labels
+        if labels is None:
+            # When labels is None, labels dict should be empty
+            self.assertEqual(lws_labels, {})
+        else:
+            # When labels is provided, they should be present in metadata
+            for key, value in labels.items():
+                self.assertIn(key, lws_labels)
+                self.assertEqual(lws_labels[key], value)