diff --git a/pathwaysutils/experimental/gke/jobset.py b/pathwaysutils/experimental/gke/jobset.py index 8089851..3524733 100644 --- a/pathwaysutils/experimental/gke/jobset.py +++ b/pathwaysutils/experimental/gke/jobset.py @@ -859,6 +859,10 @@ def add_user_workload( ) -> "PathwaysJobSet": """Adds a user workload container to the head pod. + By default, mounts a memory-backed `emptyDir` (`medium: Memory`) volume + named `dshm` at `/dev/shm` unless `volume_mounts` already defines a + `/dev/shm` mount. + Args: image: Docker image for the user workload. command: Command to execute in the container (string or sequence of @@ -890,6 +894,25 @@ def add_user_workload( ) _apply_container_kwargs(user_container, container_kwargs) + volume_mounts = list(user_container.volume_mounts or []) + if not any( + getattr(vm, "mount_path", None) == "/dev/shm" + or (isinstance(vm, Mapping) and vm.get("mountPath") == "/dev/shm") + for vm in volume_mounts + ): + shm_volume_name = "dshm" + self._add_volume_to_pod_spec( + pod_spec, + client.V1Volume( + name=shm_volume_name, + empty_dir=client.V1EmptyDirVolumeSource(medium="Memory"), + ), + ) + volume_mounts.append( + client.V1VolumeMount(name=shm_volume_name, mount_path="/dev/shm") + ) + user_container.volume_mounts = volume_mounts + init_containers = list(pod_spec.init_containers or []) containers = [] for c in pod_spec.containers: diff --git a/pathwaysutils/test/experimental/gke/jobset_test.py b/pathwaysutils/test/experimental/gke/jobset_test.py index 596ee75..2e1c5f8 100644 --- a/pathwaysutils/test/experimental/gke/jobset_test.py +++ b/pathwaysutils/test/experimental/gke/jobset_test.py @@ -1130,6 +1130,34 @@ def test_add_user_workload(self): self.assertNotIn("pathways-proxy", helper.containers["pathways-head"]) self.assertIn("user-workload", helper.containers["pathways-head"]) + # Verify memory-backed /dev/shm emptyDir volume and mount on user-workload + head_volumes = helper.pod_specs["pathways-head"]["volumes"] + self.assertIn( + {"name": "dshm", "emptyDir": {"medium": "Memory"}}, + head_volumes, + ) + self.assertIn( + {"name": "dshm", "mountPath": "/dev/shm"}, + user_c["volumeMounts"], + ) + + def test_add_user_workload_custom_shm_mount_skips_default_dshm(self): + pw_jobset = self._create_jobset(topology="2x2", num_slices=1) + pw_jobset.add_user_workload( + image="us-docker.pkg.dev/my-project/test:v1", + command=["python3", "test.py"], + volume_mounts=[ + client.V1VolumeMount(name="custom-shm", mount_path="/dev/shm") + ], + ) + helper = JobSetManifestHelper(pw_jobset.to_dict()) + self.assertNotIn("dshm", helper.volumes["pathways-head"]) + user_c = helper.containers["pathways-head"]["user-workload"] + self.assertEqual( + user_c["volumeMounts"], + [{"name": "custom-shm", "mountPath": "/dev/shm"}], + ) + def test_add_user_workload_roundtrip(self): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) pw_jobset.add_user_workload(