Test that explicitly configured pod template with shm vol is not remo…

…ved if disable in task config

Signed-off-by: Fabio Grätz <fabiogratz@googlemail.com>

plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/pod_template.py

@@ -5,7 +5,9 @@
 
 def add_shared_mem_volume_to_pod_template(pod_template: PodTemplate) -> None:
     """Add shared memory volume and volume mount to the pod template."""
-    shm_volume = V1Volume(name="shm", empty_dir=V1EmptyDirVolumeSource(medium="Memory"))
+    medium = "Memory"
+
+    shm_volume = V1Volume(name="shm", empty_dir=V1EmptyDirVolumeSource(medium=medium))
     shm_volume_mount = V1VolumeMount(name="shm", mount_path="/dev/shm")
 
     if pod_template.pod_spec is None:

plugins/flytekit-kf-pytorch/tests/test_shared.py

@@ -58,6 +58,20 @@
             True,
             True,
         ),
+        # Test that explicitly configured pod template with shared memory volume is not removed if `increase_shared_mem=False`
+        (
+            Elastic(nnodes=2, increase_shared_mem=False),
+            PodTemplate(
+                pod_spec=V1PodSpec(
+                    containers=[
+                        V1Container(name="primary", volume_mounts=[V1VolumeMount(name="shm", mount_path="/dev/shm")]),
+                    ],
+                    volumes=[V1Volume(name="shm", empty_dir=V1EmptyDirVolumeSource(medium="Memory"))],
+                ),
+            ),
+            True,
+            False,
+        ),
     ],
 )
 def test_task_shared_memory(
@@ -90,6 +104,7 @@ def test_task() -> None:
             )
 
         else:
+            # Check that the shared memory volume + volume mount is not added
             no_pod_template = test_task.pod_template is None
             no_pod_spec = no_pod_template or test_task.pod_template.pod_spec is None
             no_volumes = no_pod_spec or test_task.pod_template.pod_spec.volumes is None