From 74e6ebcd5fa868af1d649833fb9fa73169838bf2 Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Tue, 14 Jul 2026 23:00:39 -0700 Subject: [PATCH] Add Colocated Python support to PathwaysJobSet. PiperOrigin-RevId: 948118787 --- pathwaysutils/experimental/gke/jobset.py | 71 +++++++++ .../test/experimental/gke/jobset_test.py | 144 ++++++++++++++++++ 2 files changed, 215 insertions(+) diff --git a/pathwaysutils/experimental/gke/jobset.py b/pathwaysutils/experimental/gke/jobset.py index 0d23858..0fc89c4 100644 --- a/pathwaysutils/experimental/gke/jobset.py +++ b/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], diff --git a/pathwaysutils/test/experimental/gke/jobset_test.py b/pathwaysutils/test/experimental/gke/jobset_test.py index 10116c6..9372a71 100644 --- a/pathwaysutils/test/experimental/gke/jobset_test.py +++ b/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()