From d69c146a88beba93f48fa9a9715525c0babc9acb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabio=20Gr=C3=A4tz?= Date: Fri, 19 Jul 2024 18:16:37 +0200 Subject: [PATCH] Test that explicitly configured pod template with shm vol is not removed if disable in task config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Fabio Grätz --- .../flytekitplugins/kfpytorch/pod_template.py | 4 +++- plugins/flytekit-kf-pytorch/tests/test_shared.py | 15 +++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/pod_template.py b/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/pod_template.py index 8cc395bb46b..05dfeababf4 100644 --- a/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/pod_template.py +++ b/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: diff --git a/plugins/flytekit-kf-pytorch/tests/test_shared.py b/plugins/flytekit-kf-pytorch/tests/test_shared.py index 796b341acee..9853feda50f 100644 --- a/plugins/flytekit-kf-pytorch/tests/test_shared.py +++ b/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