Skip to content

Commit

Permalink
Fix volume mounts on user container (#931)
Browse files Browse the repository at this point in the history
**Pull Request Checklist**
- [x] Fixes #930
- [x] Tests added
- [x] Documentation/examples added
- [x] [Good commit messages](https://cbea.ms/git-commit/) and/or PR
title

See #930

CC: @abschm

---------

Signed-off-by: Flaviu Vadan <flaviuvadan@gmail.com>
  • Loading branch information
flaviuvadan committed Jan 17, 2024
1 parent 7c9de71 commit 7982f90
Show file tree
Hide file tree
Showing 6 changed files with 191 additions and 2 deletions.
67 changes: 67 additions & 0 deletions docs/examples/workflows/user_container.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
# User Container



This example showcases the user of a user container with a volume mount.


=== "Hera"

```python linenums="1"
from hera.workflows import (
DAG,
UserContainer,
Workflow,
models as m,
script,
)


@script(
sidecars=[
UserContainer(name="sidecar-name", volume_mounts=[m.VolumeMount(mount_path="/whatever", name="something")])
]
)
def foo():
print("hi")


with Workflow(generate_name="sidecar-volume-mount-", entrypoint="d") as w:
with DAG(name="d"):
foo()
```

=== "YAML"

```yaml linenums="1"
apiVersion: argoproj.io/v1alpha1
kind: Workflow
metadata:
generateName: sidecar-volume-mount-
spec:
entrypoint: d
templates:
- dag:
tasks:
- name: foo
template: foo
name: d
- name: foo
script:
command:
- python
image: python:3.8
source: 'import os

import sys

sys.path.append(os.getcwd())

print(''hi'')'
sidecars:
- name: sidecar-name
volumeMounts:
- mountPath: /whatever
name: something
```

Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
apiVersion: argoproj.io/v1alpha1
kind: Workflow
metadata:
generateName: title-and-description-with-markdown-
labels:
workflows.argoproj.io/archive-strategy: "false"
annotations:
workflows.argoproj.io/title: "**Test Title**"
workflows.argoproj.io/description: |
`This is a simple hello world example.`
You can also run it in Python: https://couler-proj.github.io/couler/examples/#hello-world
spec:
entrypoint: whalesay
templates:
- name: whalesay
container:
image: docker/whalesay:latest
command: [cowsay]
args: ["hello world"]
29 changes: 29 additions & 0 deletions examples/workflows/user-container.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
apiVersion: argoproj.io/v1alpha1
kind: Workflow
metadata:
generateName: sidecar-volume-mount-
spec:
entrypoint: d
templates:
- dag:
tasks:
- name: foo
template: foo
name: d
- name: foo
script:
command:
- python
image: python:3.8
source: 'import os
import sys
sys.path.append(os.getcwd())
print(''hi'')'
sidecars:
- name: sidecar-name
volumeMounts:
- mountPath: /whatever
name: something
23 changes: 23 additions & 0 deletions examples/workflows/user_container.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
"""This example showcases the user of a user container with a volume mount."""

from hera.workflows import (
DAG,
UserContainer,
Workflow,
models as m,
script,
)


@script(
sidecars=[
UserContainer(name="sidecar-name", volume_mounts=[m.VolumeMount(mount_path="/whatever", name="something")])
]
)
def foo():
print("hi")


with Workflow(generate_name="sidecar-volume-mount-", entrypoint="d") as w:
with DAG(name="d"):
foo()
16 changes: 15 additions & 1 deletion src/hera/workflows/user_container.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
ImagePullPolicy,
ResourceRequirements,
UserContainer as _ModelUserContainer,
VolumeMount as _ModelVolumeMount,
)
from hera.workflows.resources import Resources
from hera.workflows.volume import _BaseVolume
Expand Down Expand Up @@ -49,6 +50,18 @@ def _build_image_pull_policy(self) -> Optional[str]:
"Use one of {ImagePullPolicy.__members__}"
) from e

def _build_volume_mounts(self) -> Optional[List[_ModelVolumeMount]]:
"""Processes the volume mounts field and returns a generated `VolumeMount`."""
volume_mounts = self.volume_mounts
# extra volume mounts stem from using custom Hera volumes, which dynamically provision PVs + PVCs
extra_volume_mounts = None if self.volumes is None else [v._build_volume_mount() for v in self.volumes]
if volume_mounts is None:
volume_mounts = extra_volume_mounts
elif extra_volume_mounts is not None:
volume_mounts.extend(extra_volume_mounts)

return volume_mounts

def build(self) -> _ModelUserContainer:
"""Builds the Hera auto-generated model of the user container."""
env: List[EnvVar] = [
Expand All @@ -57,6 +70,7 @@ def build(self) -> _ModelUserContainer:
env_from: List[EnvFromSource] = [
var if isinstance(var, EnvFromSource) else cast(_BaseEnvFrom, var).build() for var in (self.env_from or [])
]

return _ModelUserContainer(
args=self.args,
command=self.command,
Expand All @@ -83,7 +97,7 @@ def build(self) -> _ModelUserContainer:
termination_message_policy=self.termination_message_policy,
tty=self.tty,
volume_devices=self.volume_devices,
volume_mounts=None if self.volumes is None else [v._build_volume_mount() for v in self.volumes],
volume_mounts=self._build_volume_mounts(),
working_dir=self.working_dir,
)

Expand Down
39 changes: 38 additions & 1 deletion tests/test_unit/test_user_container.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,10 @@
from hera.workflows.models import ImagePullPolicy
from hera.workflows.models import (
ImagePullPolicy,
UserContainer as ModelUserContainer,
VolumeMount,
)
from hera.workflows.user_container import UserContainer
from hera.workflows.volume import Volume


class TestUserContainer:
Expand All @@ -9,3 +14,35 @@ def test_build_image_pull_policy(self) -> None:
UserContainer(name="test", image_pull_policy=ImagePullPolicy.always)._build_image_pull_policy() == "Always"
)
assert UserContainer(name="test")._build_image_pull_policy() is None

def test_builds_volume_mounts(self) -> None:
uc: ModelUserContainer = UserContainer(
name="test", volume_mounts=[VolumeMount(name="test", mount_path="/test")]
).build()
assert isinstance(uc, ModelUserContainer)
assert uc.volume_mounts is not None
assert len(uc.volume_mounts) == 1
assert uc.volume_mounts[0].name == "test"
assert uc.volume_mounts[0].mount_path == "/test"

uc: ModelUserContainer = UserContainer(
name="test", volumes=[Volume(name="test", mount_path="/test", size="1Gi")]
).build()
assert isinstance(uc, ModelUserContainer)
assert uc.volume_mounts is not None
assert len(uc.volume_mounts) == 1
assert uc.volume_mounts[0].name == "test"
assert uc.volume_mounts[0].mount_path == "/test"

uc: ModelUserContainer = UserContainer(
name="test",
volume_mounts=[VolumeMount(name="test1", mount_path="/test1")],
volumes=[Volume(name="test2", mount_path="/test2", size="1Gi")],
).build()
assert isinstance(uc, ModelUserContainer)
assert uc.volume_mounts is not None
assert len(uc.volume_mounts) == 2
assert uc.volume_mounts[0].name == "test1"
assert uc.volume_mounts[0].mount_path == "/test1"
assert uc.volume_mounts[1].name == "test2"
assert uc.volume_mounts[1].mount_path == "/test2"

0 comments on commit 7982f90

Please sign in to comment.