Build large-scale Argo Workflows with the Python SDK

Introduction

Argo Workflows primarily relies on YAML to define a workflow, which ensures clear and concise configurations. However, for data scientists unfamiliar with YAML, its strict indentation requirements and hierarchical structure can make configuring complex workflows difficult.

Hera is a Python SDK designed to simplify building and submitting Argo Workflows. When handling a complex workflow, Hera helps you avoid potential syntax errors that can occur in YAML. Hera also provides the following benefits:

• Hera improves development efficiency with code that is easy to understand and write.

• Each function acts as a template that seamlessly integrates with various frameworks in the Python ecosystem, providing access to a rich set of Python libraries and tools.

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

Prerequisites

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

• You have installed Hera.

  pip install hera-workflows

Scenario 1: Simple DAG Diamond

In Argo Workflows, a DAG is often used to define complex task dependencies. The diamond structure is a common workflow pattern where multiple tasks run in parallel, and their results are then aggregated into a single, subsequent task. This structure is suitable for use cases that require merging different data streams or processing results. This example uses Hera to define a diamond-structured workflow where an initial task (A) is followed by two parallel tasks (B and C), which then converge into a final 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 earlier.
   global_config.verify_ssl = ""
   # The script decorator function is a key Hera feature that enables near-native Python function orchestration.
   # It allows you to call the function within a Hera context manager, such as a Workflow or Steps context.
   # The function also 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 is the primary resource in Argo and a key class in Hera, responsible for saving a template, setting an entrypoint, and running the template.
   with Workflow(
       generate_name="dag-diamond-",
       entrypoint="diamond",
       namespace="argo",
   ) as w:
       with DAG(name="diamond"):
           A = echo(name="A", arguments={"message": "A"})  # Build the 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 task dependencies: Tasks B and C depend on A, and D depends on B and C.
   # Create the workflow.
   w.create()

2. Run the following command to submit the workflow.

   python simpleDAG.py

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

   The DAG for the dag-diamond-g9v45 workflow displays a diamond shape: the main node connects to node A, which branches to nodes B and C. Nodes B and C then converge into the bottom node D. All nodes show a green check mark, indicating successful execution.

Scenario 2: MapReduce

To implement MapReduce-style data processing in Argo Workflows, you can use its DAG template to organize and coordinate multiple tasks, simulating the Map and Reduce phases. The following example shows how to use Hera to build a simple MapReduce workflow for a word count task in text files. Because each step is a Python function, integration with the Python ecosystem is straightforward.

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

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

   View 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 earlier.
   global_config.verify_ssl = ""
   # When you use the script decorator, pass script parameters to the decorator. This includes image, inputs, outputs, resources, and more.
   @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:  # Create multiple files based on num_parts. Write the foo key and the 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)
   # Define image, inputs, and outputs in the script.
   @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:  # Generate a new file based on the count of foo. Multiply the part number by 2 and write it to the bar key.
       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)
   # Define image, inputs, outputs, and resources in the script.
   @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:   # Calculate the sum of bar values across all parts.
       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. Specify 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 the templates.
           m = map_(
               with_param=s.result,
               arguments=[Parameter(name="part_id", value="{{item}}"), OSSArtifact(name="part", key="{{workflow.name}}/parts/{{item}}.json"),],
           )   # Enter the parameters and build the templates.
           s >> m >> reduce()   # Define the 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 workflow's DAG and results in the Workflow Console (Argo). On the WORKFLOW DETAILS page of the Argo Workflows console, you can see that all nodes of the map-reduce workflow (split → 4 parallel map tasks → reduce) have executed successfully. Each node displays a green check mark.

References

• Hera documentation

  • For more information about Hera, see Hera overview.

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

• Sample YAML deployments

  • To learn how to deploy a simple-diamond workflow using YAML, see dag-diamond.yaml.

  • To learn how to deploy a map-reduce workflow using YAML, 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).