Make GKE JobSet container and pod mutators idempotent.

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

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()