From 6f32dce9ffa0b17747c7a86d92c6741eed90b525 Mon Sep 17 00:00:00 2001 From: Alice Purcell Date: Mon, 14 Oct 2024 14:11:26 +0100 Subject: [PATCH] Move with_param examples into script_annotations tests --- .../script-annotations-dynamic-fanout.yaml | 60 ----------------- .../script-runner-io-dynamic-fanout.yaml | 64 ------------------- .../pydantic_io_with_param.py | 9 +-- .../script_annotations/with_param.py | 9 +-- tests/test_script_annotations.py | 64 +++++++++++++++++++ 5 files changed, 66 insertions(+), 140 deletions(-) delete mode 100644 examples/workflows/experimental/script-annotations-dynamic-fanout.yaml delete mode 100644 examples/workflows/experimental/script-runner-io-dynamic-fanout.yaml rename examples/workflows/experimental/script_runner_io_dynamic_fanout.py => tests/script_annotations/pydantic_io_with_param.py (69%) rename examples/workflows/experimental/script_annotations_dynamic_fanout.py => tests/script_annotations/with_param.py (64%) diff --git a/examples/workflows/experimental/script-annotations-dynamic-fanout.yaml b/examples/workflows/experimental/script-annotations-dynamic-fanout.yaml deleted file mode 100644 index c3df0b42d..000000000 --- a/examples/workflows/experimental/script-annotations-dynamic-fanout.yaml +++ /dev/null @@ -1,60 +0,0 @@ -apiVersion: argoproj.io/v1alpha1 -kind: Workflow -metadata: - generateName: dynamic-fanout- -spec: - entrypoint: d - templates: - - dag: - tasks: - - name: generate - template: generate - - arguments: - parameters: - - description: this is some value - name: some-value - value: '{{item}}' - depends: generate - name: consume - template: consume - withParam: '{{tasks.generate.outputs.parameters.some-values}}' - name: d - - name: generate - outputs: - parameters: - - name: some-values - valueFrom: - path: /tmp/hera-outputs/parameters/some-values - script: - args: - - -m - - hera.workflows.runner - - -e - - examples.workflows.experimental.script_annotations_dynamic_fanout:generate - command: - - python - env: - - name: hera__script_annotations - value: '' - - name: hera__outputs_directory - value: /tmp/hera-outputs - image: python:3.9 - source: '{{inputs.parameters}}' - - inputs: - parameters: - - description: this is some value - name: some-value - name: consume - script: - args: - - -m - - hera.workflows.runner - - -e - - examples.workflows.experimental.script_annotations_dynamic_fanout:consume - command: - - python - env: - - name: hera__script_annotations - value: '' - image: python:3.9 - source: '{{inputs.parameters}}' diff --git a/examples/workflows/experimental/script-runner-io-dynamic-fanout.yaml b/examples/workflows/experimental/script-runner-io-dynamic-fanout.yaml deleted file mode 100644 index 0d8237cc1..000000000 --- a/examples/workflows/experimental/script-runner-io-dynamic-fanout.yaml +++ /dev/null @@ -1,64 +0,0 @@ -apiVersion: argoproj.io/v1alpha1 -kind: Workflow -metadata: - generateName: dynamic-fanout- -spec: - entrypoint: d - templates: - - dag: - tasks: - - name: generate - template: generate - - arguments: - parameters: - - description: this is some value - name: some-value - value: '{{item}}' - depends: generate - name: consume - template: consume - withParam: '{{tasks.generate.outputs.parameters.some-values}}' - name: d - - name: generate - outputs: - parameters: - - name: some-values - valueFrom: - path: /tmp/hera-outputs/parameters/some-values - script: - args: - - -m - - hera.workflows.runner - - -e - - examples.workflows.experimental.script_runner_io_dynamic_fanout:generate - 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.9 - source: '{{inputs.parameters}}' - - inputs: - parameters: - - description: this is some value - name: some-value - name: consume - script: - args: - - -m - - hera.workflows.runner - - -e - - examples.workflows.experimental.script_runner_io_dynamic_fanout:consume - command: - - python - env: - - name: hera__script_annotations - value: '' - - name: hera__script_pydantic_io - value: '' - image: python:3.9 - source: '{{inputs.parameters}}' diff --git a/examples/workflows/experimental/script_runner_io_dynamic_fanout.py b/tests/script_annotations/pydantic_io_with_param.py similarity index 69% rename from examples/workflows/experimental/script_runner_io_dynamic_fanout.py rename to tests/script_annotations/pydantic_io_with_param.py index 2f73d803d..4d509844b 100644 --- a/examples/workflows/experimental/script_runner_io_dynamic_fanout.py +++ b/tests/script_annotations/pydantic_io_with_param.py @@ -1,9 +1,3 @@ -""" -This example showcases how clients can use Hera to dynamically generate tasks that process outputs from one task in -parallel. This is useful for batch jobs and instances where clients do not know ahead of time how many tasks/entities -they may need to process. -""" - from typing import Annotated, List from hera.shared import global_config @@ -30,9 +24,8 @@ def consume(input: ConsumeInput) -> None: print("Received value: {value}!".format(value=input.some_value)) -# assumes you used `hera.set_global_token` and `hera.set_global_host` so that the workflow can be submitted with Workflow(generate_name="dynamic-fanout-", entrypoint="d") as w: - with DAG(name="d"): + with DAG(name="dag"): g = generate(arguments={}) c = consume(with_param=g.get_parameter("some-values")) g >> c diff --git a/examples/workflows/experimental/script_annotations_dynamic_fanout.py b/tests/script_annotations/with_param.py similarity index 64% rename from examples/workflows/experimental/script_annotations_dynamic_fanout.py rename to tests/script_annotations/with_param.py index 0b5e08d19..eb9ab9567 100644 --- a/examples/workflows/experimental/script_annotations_dynamic_fanout.py +++ b/tests/script_annotations/with_param.py @@ -1,9 +1,3 @@ -""" -This example showcases how clients can use Hera to dynamically generate tasks that process outputs from one task in -parallel. This is useful for batch jobs and instances where clients do not know ahead of time how many tasks/entities -they may need to process. -""" - from typing import Annotated, List from hera.shared import global_config @@ -22,9 +16,8 @@ def consume(some_value: Annotated[int, Parameter(name="some-value", description= print("Received value: {value}!".format(value=some_value)) -# assumes you used `hera.set_global_token` and `hera.set_global_host` so that the workflow can be submitted with Workflow(generate_name="dynamic-fanout-", entrypoint="d") as w: - with DAG(name="d"): + with DAG(name="dag"): g = generate(arguments={}) c = consume(with_param=g.get_parameter("some-values")) g >> c diff --git a/tests/test_script_annotations.py b/tests/test_script_annotations.py index 0c7e994ea..8efb5139e 100644 --- a/tests/test_script_annotations.py +++ b/tests/test_script_annotations.py @@ -438,3 +438,67 @@ def test_script_pydantic_without_experimental_flag(global_config_fixture): "Unable to instantiate since it is an experimental feature." in str(e.value) ) + + +def test_script_annotated_with_param(global_config_fixture): + """Test that with_param works correctly with annotated types.""" + # GIVEN + global_config_fixture.experimental_features["script_annotations"] = True + global_config_fixture.experimental_features["script_pydantic_io"] = True + # Force a reload of the test module, as the runner performs "importlib.import_module", which + # may fetch a cached version + module_name = "tests.script_annotations.with_param" + + module = importlib.import_module(module_name) + importlib.reload(module) + workflow = importlib.import_module(module.__name__).w + + # WHEN + workflow_dict = workflow.to_dict() + assert workflow == Workflow.from_dict(workflow_dict) + assert workflow == Workflow.from_yaml(workflow.to_yaml()) + + # THEN + (dag,) = (t for t in workflow_dict["spec"]["templates"] if t["name"] == "dag") + (consume_task,) = (t for t in dag["dag"]["tasks"] if t["name"] == "consume") + + assert consume_task["arguments"]["parameters"] == [ + { + "name": "some-value", + "value": "{{item}}", + "description": "this is some value", + } + ] + assert consume_task["withParam"] == "{{tasks.generate.outputs.parameters.some-values}}" + + +def test_script_pydantic_io_with_param(global_config_fixture): + """Test that with_param works correctly with Pydantic IO types.""" + # GIVEN + global_config_fixture.experimental_features["script_annotations"] = True + global_config_fixture.experimental_features["script_pydantic_io"] = True + # Force a reload of the test module, as the runner performs "importlib.import_module", which + # may fetch a cached version + module_name = "tests.script_annotations.pydantic_io_with_param" + + module = importlib.import_module(module_name) + importlib.reload(module) + workflow = importlib.import_module(module.__name__).w + + # WHEN + workflow_dict = workflow.to_dict() + assert workflow == Workflow.from_dict(workflow_dict) + assert workflow == Workflow.from_yaml(workflow.to_yaml()) + + # THEN + (dag,) = (t for t in workflow_dict["spec"]["templates"] if t["name"] == "dag") + (consume_task,) = (t for t in dag["dag"]["tasks"] if t["name"] == "consume") + + assert consume_task["arguments"]["parameters"] == [ + { + "name": "some-value", + "value": "{{item}}", + "description": "this is some value", + } + ] + assert consume_task["withParam"] == "{{tasks.generate.outputs.parameters.some-values}}"