Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Grantham/add-attach_shm-template #3020

Open
wants to merge 5 commits into
base: master
Choose a base branch
from
+36 −0
Open

Grantham/add-attach_shm-template #3020

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a77aab8
add template_shm
granthamtaylor Dec 24, 2024
6a3dbce
rename to attach_shm
granthamtaylor Dec 26, 2024
87aa4a7
run linters
granthamtaylor Dec 30, 2024
44e2ea3
add basic unit test
granthamtaylor Dec 30, 2024
067c902
Update tests/flytekit/unit/extras/templates/test_pod_templates.py
granthamtaylor Dec 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions flytekit/extras/pod_templates/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
from flytekit.extras.pod_templates.attach_shm import attach_shm

__all__ = ["attach_shm"]
19 changes: 19 additions & 0 deletions flytekit/extras/pod_templates/attach_shm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
from flytekit.core.pod_template import PodTemplate


def attach_shm(name: str, size: str) -> PodTemplate:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider adding size parameter validation

The attach_shm function parameters lack type validation. Consider validating that size follows Kubernetes resource quantity format (e.g., '1Gi', '500Mi') to prevent runtime errors.

Code suggestion
Check the AI-generated fix before applying
Suggested change
def attach_shm(name: str, size: str) -> PodTemplate:
def attach_shm(name: str, size: str) -> PodTemplate:
import re
size_pattern = r'^[0-9]+(Gi|Mi|Ki|G|M|K)?$'
if not re.match(size_pattern, size):
raise ValueError(f"Invalid size format: {size}. Expected format like '1Gi', '500Mi'")

Code Review Run #dd1f99

Is this a valid issue, or was it incorrectly flagged by the Agent?

  • it was incorrectly flagged

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should we make name -> primary_container_name? i thought it was the name of the mount

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

re thomas's comment, i was also thinking you were mutating a pod template. if it's just generating a default pod template with an shm, should we rename the function?

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")])],
Copy link
Contributor

@flyte-bot flyte-bot Dec 30, 2024 •
edited by kumare3
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Potential insecure temp file usage

The hardcoded path '/dev/shm' could potentially be insecure. Consider making this path configurable or validating it before use.

Code suggestion
Check the AI-generated fix before applying
 @@ -4,6 +4,6 @@
 def attach_shm(name: str, size: str) -> PodTemplate:
 -            containers=[V1Container(name=name, volume_mounts=[V1VolumeMount(mount_path="/dev/shm", name="dshm")])],
 +            containers=[V1Container(name=name, volume_mounts=[V1VolumeMount(mount_path="/dev/shm", name="dshm", read_only=True)])],

Code Review Run #dd1f99

Is this a valid issue, or was it incorrectly flagged by the Agent?

  • it was incorrectly flagged

volumes=[V1Volume(name="dshm", empty_dir=V1EmptyDirVolumeSource(medium="", size_limit=size))],
),
)
Empty file added tests/flytekit/unit/extras/templates/__init__.py
View file
Open in desktop
Empty file.
14 changes: 14 additions & 0 deletions tests/flytekit/unit/extras/templates/test_pod_templates.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
from flytekit.extras.pod_templates import attach_shm
from flytekit.core.task import task

def test_attach_shm():

shm = attach_shm("SHM", "5Gi")
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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Missing task decorator for pod template

Consider adding a decorator to attach the pod template to my_task before asserting. Currently, the test may fail as pod_template might not be properly attached.

Code suggestion
Check the AI-generated fix before applying
 @@ -10,5 +10,6 @@
      def my_task():
          pass
 
 +    my_task = task(pod_template=shm)(my_task)
      # Verify pod template is attached to task
      assert my_task.pod_template == shm

Code Review Run #8c9989

Is this a valid issue, or was it incorrectly flagged by the Agent?

  • it was incorrectly flagged

Loading