Add Colocated Python support to PathwaysJobSet.

PiperOrigin-RevId: 948118787

Author: lukebaumann

pathwaysutils/experimental/gke/jobset.py

@@ -574,6 +574,77 @@ def _add_volume_to_pod_spec(
       volumes.append(volume)
       pod_spec.volumes = volumes
 
+  def add_colocated_python(
+      self,
+      image: str,
+      shm_mount_path: str = "/tmp/shared-memory",
+      shm_size_limit: str | None = None,
+  ) -> "PathwaysJobSet":
+    """Adds colocated python sidecar to the worker pods."""
+    pod_spec = self._worker_job_template.spec.template.spec
+
+    # Add shared memory volume if not exists.
+    volumes = pod_spec.volumes or []
+    shm_volume_name = "shared-memory"
+    shm_exists = any(v.name == shm_volume_name for v in volumes)
+    if not shm_exists:
+      volumes.append(
+          client.V1Volume(
+              name=shm_volume_name,
+              empty_dir=client.V1EmptyDirVolumeSource(
+                  medium="Memory", size_limit=shm_size_limit
+              ),
+          )
+      )
+      pod_spec.volumes = volumes
+
+    # Add colocated python container.
+    colocated_container = client.V1Container(
+        name="colocated-python-sidecar",
+        image=image,
+        image_pull_policy="Always",
+        env=[
+            client.V1EnvVar(name="GRPC_SERVER_ADDRESS", value="0.0.0.0:50051"),
+            client.V1EnvVar(
+                name="CLOUD_PATHWAYS_SIDECAR_SHM_DIRECTORY",
+                value=shm_mount_path,
+            ),
+        ],
+        ports=[client.V1ContainerPort(container_port=50051)],
+        volume_mounts=[
+            client.V1VolumeMount(name="shared-tmp", mount_path="/tmp"),
+            client.V1VolumeMount(name=shm_volume_name, mount_path=shm_mount_path),
+        ],
+    )
+    colocated_container.restart_policy = "Always"  # pyrefly: ignore[missing-attribute]
+
+    init_containers = pod_spec.init_containers or []
+    init_containers.append(colocated_container)
+    pod_spec.init_containers = init_containers
+
+    # Add volume mount to pathways-worker.
+    for container in pod_spec.containers:
+      if container.name == "pathways-worker":
+        volume_mounts = container.volume_mounts or []
+        volume_mounts.append(
+            client.V1VolumeMount(
+                name=shm_volume_name, mount_path=shm_mount_path
+            )
+        )
+        container.volume_mounts = volume_mounts
+        # Add env var for shm dir.
+        env = container.env or []
+        env.append(
+            client.V1EnvVar(
+                name="cloud_pathways_sidecar_shm_directory",
+                value=shm_mount_path,
+            )
+        )
+        container.env = env
+        break
+
+    return self
+
   def add_gcsfuse(
       self,
       containers: str | Sequence[str],

pathwaysutils/test/experimental/gke/jobset_test.py

@@ -475,6 +475,150 @@ def test_add_gcsfuse_preserves_existing_metadata(self):
     self.assertEqual(pod_meta.get("annotations", {}).get("existing-pod-anno"), "value")
     self.assertEqual(pod_meta.get("annotations", {}).get("gke-gcsfuse/volumes"), "true")
 
+  def test_add_colocated_python_sidecar(self):
+    pw_jobset = self._create_jobset(topology="2x2", num_slices=1)
+
+    pw_jobset.add_colocated_python(image="gcr.io/my-project/colocated-python:custom")
+    helper = JobSetManifestHelper(pw_jobset.to_dict())
+
+    self.assertIn("colocated-python-sidecar", helper.init_containers["pathways-worker"])
+    sidecar = helper.init_containers["pathways-worker"]["colocated-python-sidecar"]
+    self.assertEqual(sidecar["restartPolicy"], "Always")
+    self.assertEqual(sidecar["image"], "gcr.io/my-project/colocated-python:custom")
+    self.assertTrue(
+        any(
+            m["name"] == "shared-memory" and m["mountPath"] == "/tmp/shared-memory"
+            for m in sidecar["volumeMounts"]
+        )
+    )
+    self.assertTrue(
+        any(
+            e["name"] == "CLOUD_PATHWAYS_SIDECAR_SHM_DIRECTORY"
+            and e["value"] == "/tmp/shared-memory"
+            for e in sidecar["env"]
+        )
+    )
+
+  def test_add_colocated_python_preserves_init_containers(self):
+    pw_jobset = self._create_jobset(topology="2x2", num_slices=1)
+    
+    # Pre-populate init container on worker pod
+    worker_spec = pw_jobset._worker_job_template.spec.template.spec
+    existing_init = client.V1Container(name="existing-init-container", image="ubuntu:latest")
+    worker_spec.init_containers = [existing_init]
+
+    pw_jobset.add_colocated_python(image="gcr.io/my-project/colocated-python:custom")
+    helper = JobSetManifestHelper(pw_jobset.to_dict())
+
+    # Verify both exist
+    self.assertIn("existing-init-container", helper.init_containers["pathways-worker"])
+    self.assertIn("colocated-python-sidecar", helper.init_containers["pathways-worker"])
+
+  def test_add_colocated_python_volume_default(self):
+    pw_jobset = self._create_jobset(topology="2x2", num_slices=1)
+
+    pw_jobset.add_colocated_python(image="gcr.io/my-project/colocated-python:custom")
+    helper = JobSetManifestHelper(pw_jobset.to_dict())
+
+    self.assertIn("shared-memory", helper.volumes["pathways-worker"])
+    shm_vol = helper.volumes["pathways-worker"]["shared-memory"]
+    self.assertNotIn("sizeLimit", shm_vol["emptyDir"])
+
+  def test_add_colocated_python_worker_mount(self):
+    pw_jobset = self._create_jobset(topology="2x2", num_slices=1)
+
+    pw_jobset.add_colocated_python(image="gcr.io/my-project/colocated-python:custom")
+    helper = JobSetManifestHelper(pw_jobset.to_dict())
+
+    self.assertIn("pathways-worker", helper.containers["pathways-worker"])
+    worker_container = helper.containers["pathways-worker"]["pathways-worker"]
+    self.assertTrue(
+        any(
+            m["name"] == "shared-memory" and m["mountPath"] == "/tmp/shared-memory"
+            for m in worker_container["volumeMounts"]
+        )
+    )
+    self.assertTrue(
+        any(
+            e["name"] == "cloud_pathways_sidecar_shm_directory"
+            and e["value"] == "/tmp/shared-memory"
+            for e in worker_container["env"]
+        )
+    )
+
+  def test_add_colocated_python_custom_shm(self):
+    pw_jobset = self._create_jobset(topology="2x2", num_slices=1)
+
+    pw_jobset.add_colocated_python(
+        image="gcr.io/my-project/colocated-python:custom",
+        shm_mount_path="/tmp/custom-shm",
+        shm_size_limit="50Gi",
+    )
+    helper = JobSetManifestHelper(pw_jobset.to_dict())
 
+    self.assertIn("colocated-python-sidecar", helper.init_containers["pathways-worker"])
+    sidecar = helper.init_containers["pathways-worker"]["colocated-python-sidecar"]
+    self.assertTrue(
+        any(
+            m["name"] == "shared-memory" and m["mountPath"] == "/tmp/custom-shm"
+            for m in sidecar["volumeMounts"]
+        )
+    )
+    self.assertTrue(
+        any(
+            e["name"] == "CLOUD_PATHWAYS_SIDECAR_SHM_DIRECTORY"
+            and e["value"] == "/tmp/custom-shm"
+            for e in sidecar["env"]
+        )
+    )
+
+    self.assertIn("shared-memory", helper.volumes["pathways-worker"])
+    shm_vol = helper.volumes["pathways-worker"]["shared-memory"]
+    self.assertEqual(shm_vol["emptyDir"]["sizeLimit"], "50Gi")
+
+    self.assertIn("pathways-worker", helper.containers["pathways-worker"])
+    worker_container = helper.containers["pathways-worker"]["pathways-worker"]
+    self.assertTrue(
+        any(
+            m["name"] == "shared-memory" and m["mountPath"] == "/tmp/custom-shm"
+            for m in worker_container["volumeMounts"]
+        )
+    )
+    self.assertTrue(
+        any(
+            e["name"] == "cloud_pathways_sidecar_shm_directory"
+            and e["value"] == "/tmp/custom-shm"
+            for e in worker_container["env"]
+        )
+    )
+
+  def test_colocated_python_with_jax_command(self):
+    jax_command = "import jax; print(jax.devices());"
+    user_pod_template = {
+        "spec": {
+            "containers": [{
+                "name": "jax-tpu",
+                "image": "gcr.io/my-project/jax-tpu:latest",
+                "command": ["python3", "-c", jax_command],
+            }]
+        }
+    }
+    pw_jobset = self._create_jobset(
+        topology="2x2",
+        num_slices=1,
+        user_pod_template=user_pod_template,
+        main_container_name="jax-tpu",
+    )
+
+    pw_jobset.add_colocated_python(image="gcr.io/my-project/colocated-python:custom")
+    helper = JobSetManifestHelper(pw_jobset.to_dict())
+
+    self.assertIn("pathways-head", helper.jobs)
+    self.assertIn("jax-tpu", helper.containers["pathways-head"])
+    jax_container = helper.containers["pathways-head"]["jax-tpu"]
+    self.assertEqual(jax_container["command"], ["python3", "-c", jax_command])
+
+    self.assertIn("pathways-worker", helper.jobs)
+    self.assertIn("colocated-python-sidecar", helper.init_containers["pathways-worker"])
 if __name__ == "__main__":
   absltest.main()