From a77aab8e1035ecc9a40a16236077580f5f4d67cb Mon Sep 17 00:00:00 2001 From: granthamtaylor Date: Tue, 24 Dec 2024 17:37:01 -0500 Subject: [PATCH 1/5] add template_shm --- flytekit/extras/tasks/template_shm.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 flytekit/extras/tasks/template_shm.py diff --git a/flytekit/extras/tasks/template_shm.py b/flytekit/extras/tasks/template_shm.py new file mode 100644 index 0000000000..bbef57120e --- /dev/null +++ b/flytekit/extras/tasks/template_shm.py @@ -0,0 +1,20 @@ +from flytekit.core.pod_template import PodTemplate + + +def template_shm(name: str, size: str) -> PodTemplate: + + from kubernetes.client.models import ( + V1Container, + V1EmptyDirVolumeSource, + V1PodSpec, + V1Volume, + V1VolumeMount, + ) + + return PodTemplate( + primary_container_name=name, + pod_spec=V1PodSpec( + containers=[V1Container(name=name, volume_mounts=[V1VolumeMount(mount_path="/dev/shm", name="dshm")])], + volumes=[V1Volume(name="dshm", empty_dir=V1EmptyDirVolumeSource(medium="", size_limit=size))], + ), + ) \ No newline at end of file From 6a3dbce3bf84f2e1bc7651648dbb7dcd95748b6f Mon Sep 17 00:00:00 2001 From: granthamtaylor Date: Thu, 26 Dec 2024 11:43:11 -0500 Subject: [PATCH 2/5] rename to attach_shm --- flytekit/extras/pod_templates/__init__.py | 5 +++++ .../{tasks/template_shm.py => pod_templates/attach_shm.py} | 4 ++-- 2 files changed, 7 insertions(+), 2 deletions(-) create mode 100644 flytekit/extras/pod_templates/__init__.py rename flytekit/extras/{tasks/template_shm.py => pod_templates/attach_shm.py} (90%) diff --git a/flytekit/extras/pod_templates/__init__.py b/flytekit/extras/pod_templates/__init__.py new file mode 100644 index 0000000000..787b43591f --- /dev/null +++ b/flytekit/extras/pod_templates/__init__.py @@ -0,0 +1,5 @@ +from flytekit.extras.pod_templates.attach_shm import attach_shm + +__all__ = [ + "attach_shm" +] diff --git a/flytekit/extras/tasks/template_shm.py b/flytekit/extras/pod_templates/attach_shm.py similarity index 90% rename from flytekit/extras/tasks/template_shm.py rename to flytekit/extras/pod_templates/attach_shm.py index bbef57120e..e06f4ab6eb 100644 --- a/flytekit/extras/tasks/template_shm.py +++ b/flytekit/extras/pod_templates/attach_shm.py @@ -1,7 +1,7 @@ from flytekit.core.pod_template import PodTemplate -def template_shm(name: str, size: str) -> PodTemplate: +def attach_shm(name: str, size: str) -> PodTemplate: from kubernetes.client.models import ( V1Container, @@ -17,4 +17,4 @@ def template_shm(name: str, size: str) -> PodTemplate: containers=[V1Container(name=name, volume_mounts=[V1VolumeMount(mount_path="/dev/shm", name="dshm")])], volumes=[V1Volume(name="dshm", empty_dir=V1EmptyDirVolumeSource(medium="", size_limit=size))], ), - ) \ No newline at end of file + ) From 87aa4a7e1a271586b2573048b198e704d791aa0f Mon Sep 17 00:00:00 2001 From: granthamtaylor Date: Mon, 30 Dec 2024 11:11:02 -0500 Subject: [PATCH 3/5] run linters --- flytekit/extras/pod_templates/__init__.py | 4 +--- flytekit/extras/pod_templates/attach_shm.py | 1 - 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/flytekit/extras/pod_templates/__init__.py b/flytekit/extras/pod_templates/__init__.py index 787b43591f..846b4723bf 100644 --- a/flytekit/extras/pod_templates/__init__.py +++ b/flytekit/extras/pod_templates/__init__.py @@ -1,5 +1,3 @@ from flytekit.extras.pod_templates.attach_shm import attach_shm -__all__ = [ - "attach_shm" -] +__all__ = ["attach_shm"] diff --git a/flytekit/extras/pod_templates/attach_shm.py b/flytekit/extras/pod_templates/attach_shm.py index e06f4ab6eb..b8a961bd84 100644 --- a/flytekit/extras/pod_templates/attach_shm.py +++ b/flytekit/extras/pod_templates/attach_shm.py @@ -2,7 +2,6 @@ def attach_shm(name: str, size: str) -> PodTemplate: - from kubernetes.client.models import ( V1Container, V1EmptyDirVolumeSource, From 44e2ea336f074c0c396fb1236d61d15b01528a7c Mon Sep 17 00:00:00 2001 From: granthamtaylor Date: Mon, 30 Dec 2024 11:16:44 -0500 Subject: [PATCH 4/5] add basic unit test --- tests/flytekit/unit/extras/templates/__init__.py | 0 .../unit/extras/templates/test_pod_templates.py | 10 ++++++++++ 2 files changed, 10 insertions(+) create mode 100644 tests/flytekit/unit/extras/templates/__init__.py create mode 100644 tests/flytekit/unit/extras/templates/test_pod_templates.py diff --git a/tests/flytekit/unit/extras/templates/__init__.py b/tests/flytekit/unit/extras/templates/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/flytekit/unit/extras/templates/test_pod_templates.py b/tests/flytekit/unit/extras/templates/test_pod_templates.py new file mode 100644 index 0000000000..4e6b13bee1 --- /dev/null +++ b/tests/flytekit/unit/extras/templates/test_pod_templates.py @@ -0,0 +1,10 @@ +from flytekit.extras.pod_templates import attach_shm +from flytekit.core.task import task + +def test_attach_shm(): + + shm = attach_shm("SHM", "5Gi") + + @task(pod_template=shm) + def my_task(): + pass From 067c902e5f831b9cd0792eacab3fd427e7b11b2f Mon Sep 17 00:00:00 2001 From: Grantham Taylor <54340816+granthamtaylor@users.noreply.github.com> Date: Mon, 30 Dec 2024 13:41:45 -0500 Subject: [PATCH 5/5] Update tests/flytekit/unit/extras/templates/test_pod_templates.py Co-authored-by: Flyte Bot --- .../unit/extras/templates/test_pod_templates.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/flytekit/unit/extras/templates/test_pod_templates.py b/tests/flytekit/unit/extras/templates/test_pod_templates.py index 4e6b13bee1..1d8c51ab75 100644 --- a/tests/flytekit/unit/extras/templates/test_pod_templates.py +++ b/tests/flytekit/unit/extras/templates/test_pod_templates.py @@ -2,9 +2,13 @@ from flytekit.core.task import task def test_attach_shm(): - + shm = attach_shm("SHM", "5Gi") - - @task(pod_template=shm) + assert shm.name == "SHM" + assert shm.size == "5Gi" + def my_task(): pass + + # Verify pod template is attached to task + assert my_task.pod_template == shm