Fix for volumes in WorkflowTemplate Spec (#629)

**Description of PR** Currently, there is bug with the `WorkflowTemplate` class where volumes are not being correctly built. Example: ```python from hera.workflows import WorkflowTemplate, DownwardAPIVolume from hera.workflows import models as m from hera.workflows import Script with WorkflowTemplate( name="test", volumes=[DownwardAPIVolume(name="podinfo", mount_path="/etc/podinfo", items=[m.DownwardAPIVolumeFile(field_ref=m.ObjectFieldSelector(field_path="metadata.annotations"), path="annotations")])] ) as w: Script(name="test", source="test", volume_mounts=[m.VolumeMount(name="podinfo", mount_path="/etc/podinfo")]) print(w.to_yaml()) ``` This generates: ```yaml ... volumes: name: podinfo ``` It should generate: ```yaml ... volumes: - downwardAPI: items: - fieldRef: fieldPath: metadata.annotations path: annotations name: podinfo ``` This works as expected for `Workflow` but not for `WorkflowTemplate`, this because for `WorkflowTemplate` `self._build_volumes()` is not being called and volume claims are not being built. This PR fixes this by generalising the logic used to build templates such that the same code is used for `Workflow` and `WorkflowTemplate`. Signed-off-by: Kurt Degiorgio <[email protected]> Co-authored-by: Kurt Degiorgio <[email protected]>
argoproj-labs · May 17, 2023 · e358864 · e358864
1 parent acfeb05
commit e358864
Show file tree

Hide file tree

Showing 5 changed files with 255 additions and 22 deletions.
diff --git a/docs/examples/workflows/volume_mounts_wt.md b/docs/examples/workflows/volume_mounts_wt.md
@@ -0,0 +1,129 @@
+# Volume Mounts Wt
+
+
+
+
+
+
+=== "Hera"
+
+    ```python linenums="1"
+    from hera.workflows import (
+        DAG,
+        Parameter,
+        Volume,
+        WorkflowTemplate,
+        models as m,
+        script,
+    )
+
+
+    @script(
+        inputs=Parameter(name="vol"),
+        volume_mounts=[m.VolumeMount(name="{{inputs.parameters.vol}}", mount_path="/mnt/vol")],
+    )
+    def foo():
+        import os
+        import subprocess
+
+        print(os.listdir("/mnt"))
+        print(subprocess.run("cd /mnt && df -h", shell=True, capture_output=True).stdout.decode())
+
+
+    with WorkflowTemplate(
+        generate_name="volumes-",
+        entrypoint="d",
+        volumes=[
+            Volume(name="v1", mount_path="/mnt/v1", size="1Gi"),
+            Volume(name="v2", mount_path="/mnt/v2", size="3Gi"),
+            Volume(name="v3", mount_path="/mnt/v3", size="5Gi"),
+        ],
+    ) as w:
+        with DAG(name="d"):
+            foo(name="v1", arguments=Parameter(name="vol", value="v1"))
+            foo(name="v2", arguments=Parameter(name="vol", value="v2"))
+            foo(name="v3", arguments=Parameter(name="vol", value="v3"))
+    ```
+
+=== "YAML"
+
+    ```yaml linenums="1"
+    apiVersion: argoproj.io/v1alpha1
+    kind: WorkflowTemplate
+    metadata:
+      generateName: volumes-
+    spec:
+      entrypoint: d
+      templates:
+      - dag:
+          tasks:
+          - arguments:
+              parameters:
+              - name: vol
+                value: v1
+            name: v1
+            template: foo
+          - arguments:
+              parameters:
+              - name: vol
+                value: v2
+            name: v2
+            template: foo
+          - arguments:
+              parameters:
+              - name: vol
+                value: v3
+            name: v3
+            template: foo
+        name: d
+      - inputs:
+          parameters:
+          - name: vol
+        name: foo
+        script:
+          command:
+          - python
+          image: python:3.8
+          source: 'import os
+
+            import sys
+
+            sys.path.append(os.getcwd())
+
+            import os
+
+            import subprocess
+
+            print(os.listdir(''/mnt''))
+
+            print(subprocess.run(''cd /mnt && df -h'', shell=True, capture_output=True).stdout.decode())'
+          volumeMounts:
+          - mountPath: /mnt/vol
+            name: '{{inputs.parameters.vol}}'
+      volumeClaimTemplates:
+      - metadata:
+          name: v1
+        spec:
+          accessModes:
+          - ReadWriteOnce
+          resources:
+            requests:
+              storage: 1Gi
+      - metadata:
+          name: v2
+        spec:
+          accessModes:
+          - ReadWriteOnce
+          resources:
+            requests:
+              storage: 3Gi
+      - metadata:
+          name: v3
+        spec:
+          accessModes:
+          - ReadWriteOnce
+          resources:
+            requests:
+              storage: 5Gi
+    ```
+
diff --git a/examples/workflows/volume-mounts-wt.yaml b/examples/workflows/volume-mounts-wt.yaml
@@ -0,0 +1,77 @@
+apiVersion: argoproj.io/v1alpha1
+kind: WorkflowTemplate
+metadata:
+  generateName: volumes-
+spec:
+  entrypoint: d
+  templates:
+  - dag:
+      tasks:
+      - arguments:
+          parameters:
+          - name: vol
+            value: v1
+        name: v1
+        template: foo
+      - arguments:
+          parameters:
+          - name: vol
+            value: v2
+        name: v2
+        template: foo
+      - arguments:
+          parameters:
+          - name: vol
+            value: v3
+        name: v3
+        template: foo
+    name: d
+  - inputs:
+      parameters:
+      - name: vol
+    name: foo
+    script:
+      command:
+      - python
+      image: python:3.8
+      source: 'import os
+
+        import sys
+
+        sys.path.append(os.getcwd())
+
+        import os
+
+        import subprocess
+
+        print(os.listdir(''/mnt''))
+
+        print(subprocess.run(''cd /mnt && df -h'', shell=True, capture_output=True).stdout.decode())'
+      volumeMounts:
+      - mountPath: /mnt/vol
+        name: '{{inputs.parameters.vol}}'
+  volumeClaimTemplates:
+  - metadata:
+      name: v1
+    spec:
+      accessModes:
+      - ReadWriteOnce
+      resources:
+        requests:
+          storage: 1Gi
+  - metadata:
+      name: v2
+    spec:
+      accessModes:
+      - ReadWriteOnce
+      resources:
+        requests:
+          storage: 3Gi
+  - metadata:
+      name: v3
+    spec:
+      accessModes:
+      - ReadWriteOnce
+      resources:
+        requests:
+          storage: 5Gi
diff --git a/examples/workflows/volume_mounts_wt.py b/examples/workflows/volume_mounts_wt.py
@@ -0,0 +1,35 @@
+from hera.workflows import (
+    DAG,
+    Parameter,
+    Volume,
+    WorkflowTemplate,
+    models as m,
+    script,
+)
+
+
+@script(
+    inputs=Parameter(name="vol"),
+    volume_mounts=[m.VolumeMount(name="{{inputs.parameters.vol}}", mount_path="/mnt/vol")],
+)
+def foo():
+    import os
+    import subprocess
+
+    print(os.listdir("/mnt"))
+    print(subprocess.run("cd /mnt && df -h", shell=True, capture_output=True).stdout.decode())
+
+
+with WorkflowTemplate(
+    generate_name="volumes-",
+    entrypoint="d",
+    volumes=[
+        Volume(name="v1", mount_path="/mnt/v1", size="1Gi"),
+        Volume(name="v2", mount_path="/mnt/v2", size="3Gi"),
+        Volume(name="v3", mount_path="/mnt/v3", size="5Gi"),
+    ],
+) as w:
+    with DAG(name="d"):
+        foo(name="v1", arguments=Parameter(name="vol", value="v1"))
+        foo(name="v2", arguments=Parameter(name="vol", value="v2"))
+        foo(name="v3", arguments=Parameter(name="vol", value="v3"))
diff --git a/src/hera/workflows/workflow.py b/src/hera/workflows/workflow.py
@@ -198,10 +198,8 @@ def _set_image_pull_secrets(cls, v):
                 result.append(secret)
         return result
 
-    def build(self) -> TWorkflow:
-        """Builds the Workflow and its components into an Argo schema Workflow object."""
-        self = self._dispatch_hooks()
-
+    def _build_templates(self) -> List[TTemplate]:
+        """Builds the templates into an Argo schema."""
         templates = []
         for template in self.templates:
             if isinstance(template, HookMixin):
@@ -242,7 +240,13 @@ def build(self) -> TWorkflow:
                     for claim_name, claim in new_volume_claims_map.items():
                         if claim_name not in current_volume_claims_map:
                             self.volume_claim_templates.append(claim)
+        return templates
+
+    def build(self) -> TWorkflow:
+        """Builds the Workflow and its components into an Argo schema Workflow object."""
+        self = self._dispatch_hooks()
 
+        templates = self._build_templates()
         workflow_claims = self._build_persistent_volume_claims()
         volume_claim_templates = (self.volume_claim_templates or []) + (workflow_claims or [])
         return _ModelWorkflow(

diff --git a/src/hera/workflows/workflow_template.py b/src/hera/workflows/workflow_template.py
@@ -4,18 +4,15 @@
 for more on WorkflowTemplates.
 """
 from pydantic import validator
-from typing_extensions import get_args
 
-from hera.workflows._mixins import HookMixin
-from hera.workflows.exceptions import InvalidType
 from hera.workflows.models import (
     ObjectMeta,
     WorkflowSpec as _ModelWorkflowSpec,
     WorkflowTemplate as _ModelWorkflowTemplate,
     WorkflowTemplateCreateRequest,
     WorkflowTemplateLintRequest,
 )
-from hera.workflows.protocol import Templatable, TTemplate, TWorkflow
+from hera.workflows.protocol import TWorkflow
 from hera.workflows.workflow import Workflow
 
 
@@ -52,18 +49,9 @@ def build(self) -> TWorkflow:
         """Builds the WorkflowTemplate and its components into an Argo schema WorkflowTemplate object."""
         self = self._dispatch_hooks()
 
-        templates = []
-        for template in self.templates:
-            if isinstance(template, HookMixin):
-                template = template._dispatch_hooks()
-
-            if isinstance(template, Templatable):
-                templates.append(template._build_template())
-            elif isinstance(template, get_args(TTemplate)):
-                templates.append(template)
-            else:
-                raise InvalidType(f"{type(template)} is not a valid template type")
-
+        templates = self._build_templates()
+        workflow_claims = self._build_persistent_volume_claims()
+        volume_claim_templates = (self.volume_claim_templates or []) + (workflow_claims or [])
         return _ModelWorkflowTemplate(
             api_version=self.api_version,
             kind=self.kind,
@@ -124,8 +112,8 @@ def build(self) -> TWorkflow:
                 tolerations=self.tolerations,
                 ttl_strategy=self.ttl_strategy,
                 volume_claim_gc=self.volume_claim_gc,
-                volume_claim_templates=self.volume_claim_templates,
-                volumes=self.volumes,
+                volume_claim_templates=volume_claim_templates or None,
+                volumes=self._build_volumes(),
                 workflow_metadata=self.workflow_metadata,
                 workflow_template_ref=self.workflow_template_ref,
             ),