Build large-scale Argo Workflows with the Python SDK

Introduction

Argo Workflows primarily uses YAML to define workflows, which helps keep configurations clear and concise. However, for data scientists unfamiliar with YAML, its strict indentation and hierarchical structure can make designing a complex workflow difficult.

Hera is a Python SDK framework designed to simplify building and submitting Argo workflows. Hera helps you avoid common YAML syntax errors in complex workflows. The Hera Python SDK also offers the following advantages:

• Simplicity: Hera improves development efficiency with intuitive, easy-to-write code.

• Easy Python ecosystem integration: Each function is a template that integrates seamlessly with various frameworks in the Python ecosystem. This gives you access to a rich set of Python libraries and tools.

• Testability: You can directly use Python's testing frameworks, improving code quality and maintainability.

Prerequisites

• You have installed the Argo components and console and obtained the necessary credentials and Argo Server access IP address. For more information, see Enable batch task orchestration.

• Hera is installed.

  pip install hera-workflows

Scenario 1: Simple DAG diamond

In Argo Workflows, a DAG (directed acyclic graph) is often used to define complex task dependencies. The diamond is a common workflow pattern where a task branches into multiple parallel tasks, which then converge into a single final task. This structure is useful for scenarios that require merging data streams or processing results. This example shows how to use Hera to define a workflow with a diamond structure, where task A runs first, followed by tasks B and C in parallel, which then pass their outputs to task D.

1. Create a file named simpleDAG.py with the following content.

   # Import the required packages.
   from hera.workflows import DAG, Workflow, script
   from hera.shared import global_config
   import urllib3
   urllib3.disable_warnings()
   # Configure the access address and token.
   global_config.host = "https://${IP}:2746"
   global_config.token = "abcdefgxxxxxx"  # Enter the token you obtained previously.
   global_config.verify_ssl = ""
   # The @script decorator is a key feature in Hera that enables near-native Python function orchestration.
   # It allows you to call the decorated function within a Hera context manager, such as a Workflow or Steps context.
   # The function still runs normally outside of any Hera context, which means you can write unit tests for it.
   # This example prints the input message.
   @script()
   def echo(message: str):
       print(message)
   # Build the workflow. A workflow, a primary resource in Argo and a key class in Hera,
   # saves templates, sets an entrypoint, and runs them.
   with Workflow(
       generate_name="dag-diamond-",
       entrypoint="diamond",
       namespace="argo",
   ) as w:
       with DAG(name="diamond"):
           A = echo(name="A", arguments={"message": "A"})  # Build a template.
           B = echo(name="B", arguments={"message": "B"})
           C = echo(name="C", arguments={"message": "C"})
           D = echo(name="D", arguments={"message": "D"})
           A >> [B, C] >> D      # Define dependencies: A runs before B and C, which in turn run before D.
   # Create the workflow.
   w.create()

2. Run the following command to submit the workflow.

   python simpleDAG.py

3. After the workflow runs, view the task DAG and results in the Workflow Console (Argo).

   The DAG shows that the dag-diamond-g9v45 workflow ran successfully. Node A branches into parallel nodes B and C, which then converge into node D. All nodes have a green checkmark, indicating successful completion.

Scenario 2: MapReduce

You can use a DAG template to implement MapReduce-style data processing in Argo Workflows, which helps organize and coordinate multiple tasks that simulate the Map and Reduce phases. This example builds a simple MapReduce workflow for a numerical processing task with Hera. Each step is a Python function, which allows for easy integration with the Python ecosystem.

1. Configure artifacts. For more information, see Configure Artifacts.

2. Create a file named map-reduce.py with the following content.

   Expand to view the code

   from hera.workflows import DAG, Artifact, NoneArchiveStrategy, Parameter, OSSArtifact, Workflow, script
   from hera.shared import global_config
   import urllib3
   urllib3.disable_warnings()
   # Set the access address.
   global_config.host = "https://${IP}:2746"
   global_config.token = "abcdefgxxxxxx"  # Enter the token you obtained previously.
   global_config.verify_ssl = ""
   # When using the @script decorator, pass script parameters to it, such as image, inputs, outputs, and resources.
   @script(
       image="mirrors-ssl.aliyuncs.com/python:alpine",
       inputs=Parameter(name="num_parts"),
       outputs=OSSArtifact(name="parts", path="/mnt/out", archive=NoneArchiveStrategy(), key="{{workflow.name}}/parts"),
   )
   def split(num_parts: int) -> None:  # Creates multiple files based on the input parameter num_parts, writing a "foo" value and part number to each file.
       import json
       import os
       import sys
       os.mkdir("/mnt/out")
       part_ids = list(map(lambda x: str(x), range(num_parts)))
       for i, part_id in enumerate(part_ids, start=1):
           with open("/mnt/out/" + part_id + ".json", "w") as f:
               json.dump({"foo": i}, f)
       json.dump(part_ids, sys.stdout)
   # In the @script decorator, define the image, inputs, and outputs.
   @script(
       image="mirrors-ssl.aliyuncs.com/python:alpine",
       inputs=[Parameter(name="part_id", value="0"), Artifact(name="part", path="/mnt/in/part.json"),],
       outputs=OSSArtifact(
           name="part",
           path="/mnt/out/part.json",
           archive=NoneArchiveStrategy(),
           key="{{workflow.name}}/results/{{inputs.parameters.part_id}}.json",
       ),
   )
   def map_() -> None:  # Reads the "foo" value from an input file, multiplies it by 2, and writes the result to a "bar" value in a new file.
       import json
       import os
       os.mkdir("/mnt/out")
       with open("/mnt/in/part.json") as f:
           part = json.load(f)
       with open("/mnt/out/part.json", "w") as f:
           json.dump({"bar": part["foo"] * 2}, f)
   # In the @script decorator, define the image, inputs, outputs, and resources.
   @script(
       image="mirrors-ssl.aliyuncs.com/python:alpine",
       inputs=OSSArtifact(name="results", path="/mnt/in", key="{{workflow.name}}/results"),
       outputs=OSSArtifact(
           name="total", path="/mnt/out/total.json", archive=NoneArchiveStrategy(), key="{{workflow.name}}/total.json"
       ),
   )
   def reduce() -> None:   # Calculates the sum of all "bar" values from the part files.
       import json
       import os
       os.mkdir("/mnt/out")
       total = 0
       for f in list(map(lambda x: open("/mnt/in/" + x), os.listdir("/mnt/in"))):
           result = json.load(f)
           total = total + result["bar"]
       with open("/mnt/out/total.json", "w") as f:
           json.dump({"total": total}, f)
   # Build the workflow, specifying the name, entrypoint, namespace, and global parameters.
   with Workflow(generate_name="map-reduce-", entrypoint="main", namespace="argo", arguments=Parameter(name="num_parts", value="4")) as w:
       with DAG(name="main"):
           s = split(arguments=Parameter(name="num_parts", value="{{workflow.parameters.num_parts}}")) # Build templates.
           m = map_(
               with_param=s.result,
               arguments=[Parameter(name="part_id", value="{{item}}"), OSSArtifact(name="part", key="{{workflow.name}}/parts/{{item}}.json"),],
           )   # Provide input parameters and build templates.
           s >> m >> reduce()   # Define task dependencies.
   # Create the workflow.
   w.create()
   

3. Run the following command to submit the workflow.

   python map-reduce.py

4. After the workflow runs, you can view the task DAG and results in the Workflow Console (Argo).

Related documentation

• Hera documentation

  • For more details about Hera, see the Hera overview.

  • To learn how to set up and use Hera for LLM training, see Train LLM with Hera.

• YAML deployment examples

  • For the YAML equivalent of the simple-diamond deployment, see dag-diamond.yaml.

  • For the YAML equivalent of the map-reduce deployment, see map-reduce.yaml.

Contact us

If you have any product suggestions or questions, you can contact us by joining the DingTalk group (ID: 35688562).