From f1e1cd1b31fd7f9e8255c9909178722a81c98be4 Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Fri, 23 Oct 2026 22:31:30 -0700 Subject: [PATCH] Make GKE JobSet container and pod mutators idempotent. PiperOrigin-RevId: 1006105629 --- pathwaysutils/experimental/gke/jobset.py | 53 +++++++++--- .../test/experimental/gke/jobset_test.py | 81 +++++++++++++++++++ 2 files changed, 123 insertions(+), 11 deletions(-) diff --git a/pathwaysutils/experimental/gke/jobset.py b/pathwaysutils/experimental/gke/jobset.py index 37bdf90..5c0f332 100644 --- a/pathwaysutils/experimental/gke/jobset.py +++ b/pathwaysutils/experimental/gke/jobset.py @@ -836,11 +836,15 @@ def _enable_gcsfuse_annotations( def _add_volume_to_pod_spec( self, pod_spec: client.V1PodSpec, volume: client.V1Volume ) -> None: - """Appends volume to pod_spec if not already present.""" - volumes = pod_spec.volumes or [] - if not any(v.name == volume.name for v in volumes): - volumes.append(volume) - pod_spec.volumes = volumes + """Adds or replaces volume by name in pod_spec.""" + volumes = list(pod_spec.volumes or []) + for i, existing in enumerate(volumes): + if existing.name == volume.name: + volumes[i] = volume + pod_spec.volumes = volumes + return + volumes.append(volume) + pod_spec.volumes = volumes def add_user_workload( self, @@ -984,7 +988,10 @@ def add_external_sidecar( # stopped after, pathways-worker. sidecar_container.restart_policy = "Always" # pyrefly: ignore[missing-attribute] - init_containers = list(pod_spec.init_containers or []) + sidecar_name = f"{name}-sidecar" + init_containers = [ + c for c in (pod_spec.init_containers or []) if c.name != sidecar_name + ] init_containers.append(sidecar_container) pod_spec.init_containers = init_containers @@ -1003,7 +1010,17 @@ def add_external_sidecar( ) container.args = args volume_mounts = container.volume_mounts or [] - if not any(vm.name == shm_volume_name for vm in volume_mounts): + existing_shm_mount = next( + (vm for vm in volume_mounts if vm.name == shm_volume_name), None + ) + if existing_shm_mount is not None: + if existing_shm_mount.mount_path != shm_mount_path: + raise ValueError( + f"Conflicting shm_mount_path {shm_mount_path!r};" + f" {shm_volume_name!r} is already mounted at" + f" {existing_shm_mount.mount_path!r}." + ) + else: volume_mounts.append( client.V1VolumeMount( name=shm_volume_name, mount_path=shm_mount_path @@ -1097,7 +1114,11 @@ def add_gcsfuse( self._add_volume_to_pod_spec(pod_spec, volume) for container in matching: - volume_mounts = container.volume_mounts or [] + volume_mounts = [ + vm + for vm in (container.volume_mounts or []) + if vm.mount_path != mount_path + ] volume_mounts.append(volume_mount) container.volume_mounts = volume_mounts @@ -1158,7 +1179,11 @@ def with_custom_libtpu( ) for container in matching: if mount_host_dir: - volume_mounts = container.volume_mounts or [] + volume_mounts = [ + vm + for vm in (container.volume_mounts or []) + if vm.name != CUSTOM_LIBTPU_VOLUME_NAME + ] volume_mounts.append( client.V1VolumeMount( name=CUSTOM_LIBTPU_VOLUME_NAME, @@ -1167,8 +1192,14 @@ def with_custom_libtpu( ) ) container.volume_mounts = volume_mounts - env = container.env or [] - env.append(client.V1EnvVar(name="TPU_LIBRARY_PATH", value=path)) + env = list(container.env or []) + new_env_var = client.V1EnvVar(name="TPU_LIBRARY_PATH", value=path) + for i, existing_env in enumerate(env): + if existing_env.name == "TPU_LIBRARY_PATH": + env[i] = new_env_var + break + else: + env.append(new_env_var) container.env = env return self diff --git a/pathwaysutils/test/experimental/gke/jobset_test.py b/pathwaysutils/test/experimental/gke/jobset_test.py index a11e1ef..c23567f 100644 --- a/pathwaysutils/test/experimental/gke/jobset_test.py +++ b/pathwaysutils/test/experimental/gke/jobset_test.py @@ -1771,6 +1771,87 @@ def test_with_custom_libtpu_rejects_invalid_paths( with self.assertRaisesRegex(ValueError, expected_error): pw_jobset.with_custom_libtpu(path, mount_host_dir=mount_host_dir) + def test_with_custom_libtpu_replaces_existing_env_and_mount(self): + pw_jobset = self._create_jobset(topology="2x2", num_slices=1) + pw_jobset.with_custom_libtpu( + "/var/lib/custom_tpu_v1/libtpu.so", + containers="worker", + mount_host_dir=True, + ) + pw_jobset.with_custom_libtpu( + "/var/lib/custom_tpu_v2/libtpu.so", + containers="worker", + mount_host_dir=True, + ) + helper = JobSetManifestHelper(pw_jobset.to_dict()) + worker = helper.containers["pathways-worker"]["pathways-worker"] + tpu_lib_envs = [ + e for e in worker["env"] if e["name"] == "TPU_LIBRARY_PATH" + ] + self.assertEqual( + tpu_lib_envs, + [{"name": "TPU_LIBRARY_PATH", "value": "/var/lib/custom_tpu_v2/libtpu.so"}], + ) + libtpu_mounts = [ + m for m in worker["volumeMounts"] if m["name"] == "custom-libtpu" + ] + self.assertEqual( + libtpu_mounts, + [{ + "name": "custom-libtpu", + "mountPath": "/var/lib/custom_tpu_v2", + "readOnly": True, + }], + ) + self.assertEqual( + helper.volumes["pathways-worker"]["custom-libtpu"]["hostPath"]["path"], + "/var/lib/custom_tpu_v2", + ) + + def test_add_gcsfuse_replaces_duplicate_mount_path(self): + pw_jobset = self._create_jobset(topology="2x2", num_slices=1) + pw_jobset.add_gcsfuse( + containers="worker", + mount_path="/gcs/data", + bucket="my-bucket", + read_only=False, + ) + pw_jobset.add_gcsfuse( + containers="worker", + mount_path="/gcs/data", + bucket="my-bucket", + read_only=True, + ) + helper = JobSetManifestHelper(pw_jobset.to_dict()) + worker = helper.containers["pathways-worker"]["pathways-worker"] + gcs_mounts = [ + m for m in worker["volumeMounts"] if m["mountPath"] == "/gcs/data" + ] + self.assertLen(gcs_mounts, 1) + self.assertTrue(gcs_mounts[0]["readOnly"]) + + def test_add_external_sidecar_deduplicates_and_rejects_conflicting_shm_path( + self, + ): + pw_jobset = self._create_jobset(topology="2x2", num_slices=1) + pw_jobset.add_external_sidecar( + image="gcr.io/my-project/sidecar:v1", name="custom" + ) + pw_jobset.add_external_sidecar( + image="gcr.io/my-project/sidecar:v2", name="custom" + ) + helper = JobSetManifestHelper(pw_jobset.to_dict()) + init_list = helper.pod_specs["pathways-worker"]["initContainers"] + self.assertLen(init_list, 1) + self.assertEqual(init_list[0]["image"], "gcr.io/my-project/sidecar:v2") + + with self.assertRaisesRegex(ValueError, "Conflicting shm_mount_path"): + pw_jobset.add_external_sidecar( + image="gcr.io/my-project/other:v1", + name="other", + shm_mount_path="/tmp/other-shm", + ) + if __name__ == "__main__": absltest.main()