Batch Processing

Dragon’s Batch API lets you submit Python functions, executables, and parallel jobs for distributed execution with automatic scheduling and placement of tasks. You declare what each task reads and writes, and Dragon infers a directed acyclic graph (DAG) to maximize parallel throughput — no manual dependency management needed.

When to Use Batch

Use Batch when you have tasks that should be scheduled automatically from declared read and write dependencies (on the parallel filesystem or Dragon Distributed Dictionary), or from argument passing dependencies. It is especially useful when you need to:

  • Intermix Python functions, external processes, and multi-node MPI jobs in one workflow.

  • Run large numbers of small function tasks with higher throughput.

  • Let Dragon schedule around dependencies efficiently, prioritize the critical path, place tasks near their data, and make intelligent MPI placement decisions.

For simpler parallel-map patterns, Tutorial 1 shows how Pool.map may be sufficient.

A Simple Batch Example

The following example uses Batch to compute squares of numbers in parallel:

Listing 63 batch_squares.py — parallel function tasks with Batch
 1import dragon
 2from multiprocessing import set_start_method
 3from dragon.workflows.batch import Batch
 4
 5
 6def square(x):
 7    return x * x
 8
 9
10if __name__ == "__main__":
11    set_start_method("dragon")
12
13    batch = Batch()
14
15    # Submit 20 tasks — Batch dispatches them in parallel automatically
16    handles = [batch.function(square, i) for i in range(20)]
17
18    # .get() blocks until the task completes and returns its result
19    results = [h.get() for h in handles]
20    print(results)
21
22    batch.join()

Run with:

dragon batch_squares.py

Tasks with Data Dependencies

Batch tracks file dependencies between tasks and automatically orders execution so that a task that reads a file waits until the task that writes it completes:

Listing 64 Chained tasks via file dependencies
 1import dragon
 2from multiprocessing import set_start_method
 3from dragon.workflows.batch import Batch
 4from pathlib import Path
 5import numpy as np
 6
 7
 8def generate(path):
 9    np.save(path, np.arange(1000, dtype=np.float64))
10
11
12def process(in_path, out_path):
13    data = np.load(in_path)
14    np.save(out_path, data ** 2)
15    return float(data.sum())
16
17
18if __name__ == "__main__":
19    set_start_method("dragon")
20
21    batch = Batch()
22    base = Path("/tmp/dragon_batch_demo")
23    base.mkdir(exist_ok=True)
24
25    raw = base / "raw.npy"
26    processed = base / "processed.npy"
27
28    # Task 1: generate data
29    t1 = batch.options(
30        writes=[batch.write(base, Path("raw.npy"))]
31    ).function(generate, raw)
32
33    # Task 2: process data — will not start until t1 finishes
34    t2 = batch.options(
35        reads=[batch.read(base, Path("raw.npy"))],
36        writes=[batch.write(base, Path("processed.npy"))],
37    ).function(process, raw, processed)
38
39    print(f"Sum of raw data: {t2.get()}")
40
41    batch.join()

Submitting Executables and MPI Jobs

Batch is not limited to Python functions. Use process() for executables and job() for MPI/parallel jobs:

Listing 65 Mixed Python + executable tasks in one Batch
 1import dragon
 2from multiprocessing import set_start_method
 3from dragon.workflows.batch import Batch
 4from dragon.native.process import ProcessTemplate, Popen
 5from pathlib import Path
 6
 7
 8def prepare_input(path):
 9    path.write_text("42\n")
10
11
12if __name__ == "__main__":
13    set_start_method("dragon")
14
15    batch = Batch()
16    base = Path("/tmp/dragon_batch_exec")
17    base.mkdir(exist_ok=True)
18    inp = base / "input.txt"
19
20    # Python function writes the input file
21    t1 = batch.options(
22        writes=[batch.write(base, Path("input.txt"))]
23    ).function(prepare_input, inp)
24
25    # External executable reads the file produced by t1. process() takes a
26    # ProcessTemplate; set stdout=Popen.PIPE to capture the process output.
27    t2 = batch.options(
28        reads=[batch.read(base, Path("input.txt"))],
29    ).process(
30        ProcessTemplate(target="cat", args=[str(inp)], stdout=Popen.PIPE),
31    )
32
33    t1.get()
34    # .get() returns the process exit code and prints the captured stdout.
35    exit_code = t2.get()
36    print("cat exit code:", exit_code)
37
38    batch.join()

For process and job tasks, .get() returns the child exit code. A non-zero exit code is returned as a value rather than raised, so you must inspect it yourself to decide whether the task succeeded, and a downstream task that depends on a process or job result receives that exit code as its input. A single process returns a scalar exit code; a multi-rank job returns a list of exit codes, one per launched process.

Lifecycle Methods

  • join() — wait for all tasks submitted by the

    current client to complete, then detach the client from the Batch instance shared by clients.

  • destroy() — in managed_lifecycle=True mode,

    shut down the Batch instance shared by clients after they detach, or after force_timeout

    expires if one was supplied.

  • fence() — wait for all currently submitted tasks to finish before returning (useful as a barrier between phases).

  • clear_results() — free the memory used to hold task results (implicitly calls fence()).