Skip to content

Commit

Permalink
Enable Pydantic I/O types in workflow context
Browse files Browse the repository at this point in the history
Extend experimental Pydantic I/O support to allow passing Pydantic types
into `@script`-decorated functions when inside a `with` workflow context
block, and using fields on the returned Pydantic output as shorthand for
the associated Hera template in subsequent steps.

Signed-off-by: Alice Purcell <alicederyn@gmail.com>
  • Loading branch information
alicederyn committed Sep 4, 2024
1 parent b1933d8 commit 3b97438
Show file tree
Hide file tree
Showing 10 changed files with 628 additions and 19 deletions.
154 changes: 154 additions & 0 deletions docs/examples/workflows/experimental/pydantic_io_in_dag_context.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,154 @@
# Pydantic Io In Dag Context






=== "Hera"

```python linenums="1"
import sys
from typing import List

if sys.version_info >= (3, 9):
from typing import Annotated
else:
from typing_extensions import Annotated


from hera.shared import global_config
from hera.workflows import DAG, Parameter, WorkflowTemplate, script
from hera.workflows.io.v1 import Input, Output

global_config.experimental_features["script_pydantic_io"] = True


class CutInput(Input):
cut_after: Annotated[int, Parameter(name="cut-after")]
strings: List[str]


class CutOutput(Output):
first_strings: Annotated[List[str], Parameter(name="first-strings")]
remainder: List[str]


class JoinInput(Input):
strings: List[str]
joiner: str


class JoinOutput(Output):
joined_string: Annotated[str, Parameter(name="joined-string")]


@script(constructor="runner")
def cut(input: CutInput) -> CutOutput:
return CutOutput(
first_strings=input.strings[: input.cut_after],
remainder=input.strings[input.cut_after :],
exit_code=1 if len(input.strings) <= input.cut_after else 0,
)


@script(constructor="runner")
def join(input: JoinInput) -> JoinOutput:
return JoinOutput(joined_string=input.joiner.join(input.strings))


with WorkflowTemplate(generate_name="pydantic-io-in-steps-context-v1-", entrypoint="d") as w:
with DAG(name="d"):
cut_result = cut(CutInput(strings=["hello", "world", "it's", "hera"], cut_after=1))
join(JoinInput(strings=cut_result.first_strings, joiner=" "))
```

=== "YAML"

```yaml linenums="1"
apiVersion: argoproj.io/v1alpha1
kind: WorkflowTemplate
metadata:
generateName: pydantic-io-in-steps-context-v1-
spec:
entrypoint: d
templates:
- dag:
tasks:
- arguments:
parameters:
- name: cut-after
value: '1'
- name: strings
value: '["hello", "world", "it''s", "hera"]'
name: cut
template: cut
- arguments:
parameters:
- name: strings
value: '{{tasks.cut.outputs.parameters.first-strings}}'
- name: joiner
value: ' '
depends: cut
name: join
template: join
name: d
- inputs:
parameters:
- name: cut-after
- name: strings
name: cut
outputs:
parameters:
- name: first-strings
valueFrom:
path: /tmp/hera-outputs/parameters/first-strings
- name: remainder
valueFrom:
path: /tmp/hera-outputs/parameters/remainder
script:
args:
- -m
- hera.workflows.runner
- -e
- examples.workflows.experimental.pydantic_io_in_dag_context:cut
command:
- python
env:
- name: hera__script_annotations
value: ''
- name: hera__outputs_directory
value: /tmp/hera-outputs
- name: hera__script_pydantic_io
value: ''
image: python:3.8
source: '{{inputs.parameters}}'
- inputs:
parameters:
- name: strings
- name: joiner
name: join
outputs:
parameters:
- name: joined-string
valueFrom:
path: /tmp/hera-outputs/parameters/joined-string
script:
args:
- -m
- hera.workflows.runner
- -e
- examples.workflows.experimental.pydantic_io_in_dag_context:join
command:
- python
env:
- name: hera__script_annotations
value: ''
- name: hera__outputs_directory
value: /tmp/hera-outputs
- name: hera__script_pydantic_io
value: ''
image: python:3.8
source: '{{inputs.parameters}}'
```

152 changes: 152 additions & 0 deletions docs/examples/workflows/experimental/pydantic_io_in_steps_context.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,152 @@
# Pydantic Io In Steps Context






=== "Hera"

```python linenums="1"
import sys
from typing import List

if sys.version_info >= (3, 9):
from typing import Annotated
else:
from typing_extensions import Annotated


from hera.shared import global_config
from hera.workflows import Parameter, Steps, WorkflowTemplate, script
from hera.workflows.io.v1 import Input, Output

global_config.experimental_features["script_pydantic_io"] = True


class CutInput(Input):
cut_after: Annotated[int, Parameter(name="cut-after")]
strings: List[str]


class CutOutput(Output):
first_strings: Annotated[List[str], Parameter(name="first-strings")]
remainder: List[str]


class JoinInput(Input):
strings: List[str]
joiner: str


class JoinOutput(Output):
joined_string: Annotated[str, Parameter(name="joined-string")]


@script(constructor="runner")
def cut(input: CutInput) -> CutOutput:
return CutOutput(
first_strings=input.strings[: input.cut_after],
remainder=input.strings[input.cut_after :],
exit_code=1 if len(input.strings) <= input.cut_after else 0,
)


@script(constructor="runner")
def join(input: JoinInput) -> JoinOutput:
return JoinOutput(joined_string=input.joiner.join(input.strings))


with WorkflowTemplate(generate_name="pydantic-io-in-steps-context-v1-", entrypoint="d") as w:
with Steps(name="d"):
cut_result = cut(CutInput(strings=["hello", "world", "it's", "hera"], cut_after=1))
join(JoinInput(strings=cut_result.first_strings, joiner=" "))
```

=== "YAML"

```yaml linenums="1"
apiVersion: argoproj.io/v1alpha1
kind: WorkflowTemplate
metadata:
generateName: pydantic-io-in-steps-context-v1-
spec:
entrypoint: d
templates:
- name: d
steps:
- - arguments:
parameters:
- name: cut-after
value: '1'
- name: strings
value: '["hello", "world", "it''s", "hera"]'
name: cut
template: cut
- - arguments:
parameters:
- name: strings
value: '{{steps.cut.outputs.parameters.first-strings}}'
- name: joiner
value: ' '
name: join
template: join
- inputs:
parameters:
- name: cut-after
- name: strings
name: cut
outputs:
parameters:
- name: first-strings
valueFrom:
path: /tmp/hera-outputs/parameters/first-strings
- name: remainder
valueFrom:
path: /tmp/hera-outputs/parameters/remainder
script:
args:
- -m
- hera.workflows.runner
- -e
- examples.workflows.experimental.pydantic_io_in_steps_context_v1:cut
command:
- python
env:
- name: hera__script_annotations
value: ''
- name: hera__outputs_directory
value: /tmp/hera-outputs
- name: hera__script_pydantic_io
value: ''
image: python:3.8
source: '{{inputs.parameters}}'
- inputs:
parameters:
- name: strings
- name: joiner
name: join
outputs:
parameters:
- name: joined-string
valueFrom:
path: /tmp/hera-outputs/parameters/joined-string
script:
args:
- -m
- hera.workflows.runner
- -e
- examples.workflows.experimental.pydantic_io_in_steps_context_v1:join
command:
- python
env:
- name: hera__script_annotations
value: ''
- name: hera__outputs_directory
value: /tmp/hera-outputs
- name: hera__script_pydantic_io
value: ''
image: python:3.8
source: '{{inputs.parameters}}'
```

84 changes: 84 additions & 0 deletions examples/workflows/experimental/pydantic-io-in-dag-context.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
apiVersion: argoproj.io/v1alpha1
kind: WorkflowTemplate
metadata:
generateName: pydantic-io-in-steps-context-v1-
spec:
entrypoint: d
templates:
- dag:
tasks:
- arguments:
parameters:
- name: cut-after
value: '1'
- name: strings
value: '["hello", "world", "it''s", "hera"]'
name: cut
template: cut
- arguments:
parameters:
- name: strings
value: '{{tasks.cut.outputs.parameters.first-strings}}'
- name: joiner
value: ' '
depends: cut
name: join
template: join
name: d
- inputs:
parameters:
- name: cut-after
- name: strings
name: cut
outputs:
parameters:
- name: first-strings
valueFrom:
path: /tmp/hera-outputs/parameters/first-strings
- name: remainder
valueFrom:
path: /tmp/hera-outputs/parameters/remainder
script:
args:
- -m
- hera.workflows.runner
- -e
- examples.workflows.experimental.pydantic_io_in_dag_context:cut
command:
- python
env:
- name: hera__script_annotations
value: ''
- name: hera__outputs_directory
value: /tmp/hera-outputs
- name: hera__script_pydantic_io
value: ''
image: python:3.8
source: '{{inputs.parameters}}'
- inputs:
parameters:
- name: strings
- name: joiner
name: join
outputs:
parameters:
- name: joined-string
valueFrom:
path: /tmp/hera-outputs/parameters/joined-string
script:
args:
- -m
- hera.workflows.runner
- -e
- examples.workflows.experimental.pydantic_io_in_dag_context:join
command:
- python
env:
- name: hera__script_annotations
value: ''
- name: hera__outputs_directory
value: /tmp/hera-outputs
- name: hera__script_pydantic_io
value: ''
image: python:3.8
source: '{{inputs.parameters}}'
Loading

0 comments on commit 3b97438

Please sign in to comment.