Add Shared Pathways Service support to PathwaysJobSet.

PiperOrigin-RevId: 949604465

Author: lukebaumann

pathwaysutils/experimental/gke/jobset.py

@@ -103,6 +103,7 @@ def __init__(
       elastic_slices: int = 0,
       labels: Mapping[str, str] | None = None,
       annotations: Mapping[str, str] | None = None,
+      shared_pathways_service: bool = False,
   ):
     """Initializes the instance.
 
@@ -124,6 +125,13 @@ def __init__(
       labels: Optional labels for the JobSet.
       annotations: Optional annotations for the JobSet.
     """
+    if shared_pathways_service and user_pod_template:
+      raise ValueError(
+          "Cannot enable shared_pathways_service when user_pod_template is"
+          " provided."
+      )
+    self._shared_pathways_service = shared_pathways_service
+
     self._name = name
     self._namespace = namespace
     self._jobset_api_version = jobset_api_version
@@ -161,6 +169,7 @@ def __init__(
         user_pod_template=user_pod_template,
         main_container_name=main_container_name,
         elastic_slices=elastic_slices,
+        shared_pathways_service=shared_pathways_service,
     )
 
     # Build worker template.
@@ -176,7 +185,7 @@ def __init__(
     )
 
     self._success_policy = None
-    if user_pod_template:
+    if user_pod_template or shared_pathways_service:
       self._success_policy = {
           "operator": "All",
           "targetReplicatedJobs": [PATHWAYS_HEAD_JOB_NAME],
@@ -199,6 +208,7 @@ def _build_head_job_template(
       user_pod_template: Mapping[str, Any] | None,
       main_container_name: str,
       elastic_slices: int,
+      shared_pathways_service: bool,
   ) -> client.V1JobTemplateSpec:
     """Builds the head job template for the JobSet.
 
@@ -357,10 +367,13 @@ def _build_head_job_template(
       labels = user_pod_template.get("metadata", {}).get("labels", {})
     else:
       # Headless mode.
+      containers = [rm_container]
+      if not shared_pathways_service:
+        containers.append(proxy_container)
       head_pod_spec = client.V1PodSpec(
           host_network=True,
           dns_policy="ClusterFirstWithHostNet",
-          containers=[rm_container, proxy_container],
+          containers=containers,
       )
       annotations = {}
       labels = {}

pathwaysutils/test/experimental/gke/jobset_test.py

@@ -921,6 +921,57 @@ def test_failure_policy(self):
     config = pw_jobset.to_dict()
     self.assertEqual(config["spec"]["failurePolicy"]["maxRestarts"], 5)
 
+  def test_shared_pathways_service(self):
+    pw_jobset = self._create_jobset(
+        name="test-sps",
+        shared_pathways_service=True,
+    )
+
+    config = pw_jobset.to_dict()
+    helper = JobSetManifestHelper(config)
+
+    self.assertEqual(
+        config["spec"]["successPolicy"],
+        {
+            "operator": "All",
+            "targetReplicatedJobs": ["pathways-head"],
+        },
+    )
+
+    self.assertTrue(config["spec"]["network"]["enableDNSHostnames"])
+    self.assertTrue(config["spec"]["network"]["publishNotReadyAddresses"])
+
+    self.assertIn("pathways-head", helper.jobs)
+    pod_spec = helper.pod_specs["pathways-head"]
 
+    # Head job should only have pathways-rm container, no pathways-proxy.
+    self.assertIn("pathways-rm", helper.containers["pathways-head"])
+    self.assertNotIn("pathways-proxy", helper.containers["pathways-head"])
+    self.assertLen(pod_spec["containers"], 1)
+
+  def test_shared_pathways_service_with_user_template_fails(self):
+    user_pod_template = {
+        "spec": {
+            "containers": [{
+                "name": "jax-tpu",
+                "image": "gcr.io/my-project/jax-tpu:latest",
+            }]
+        }
+    }
+    with self.assertRaisesRegex(
+        ValueError,
+        "Cannot enable shared_pathways_service when user_pod_template is"
+        " provided.",
+    ):
+      jobset.PathwaysJobSet(
+          name="test-sps",
+          namespace="default",
+          pathways_dir="gs://bucket/scratch",
+          tpu_type="v5e",
+          topology="2x2",
+          num_slices=2,
+          shared_pathways_service=True,
+          user_pod_template=user_pod_template,
+      )
 if __name__ == "__main__":
   absltest.main()