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, Python multiprocessing shows how Pool.map may be sufficient.

A Simple Batch Example

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

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

Run with:

dragon batch_squares.py

Context-manager Cleanup

Use Batch as a context manager when a block owns a Batch client for one workflow phase:

with Batch() as batch:
    tasks = [batch.function(square, i) for i in range(20)]
    results = [task.get() for task in tasks]

On normal exit, the context manager calls join(). This waits for work submitted by the current client and detaches that client from the shared Batch runtime. In the default unmanaged mode, the runtime is also shut down automatically when this is the last attached client. Fetch any task results you still need before leaving the block, because shutting down an unmanaged runtime destroys its internal results DDict.

If code inside the block raises an exception, the context manager calls terminate() and then re-raises the original exception. This is the appropriate cleanup path when work may still be running and should not be allowed to continue.

For a runtime shared by multiple clients, use managed lifecycle mode explicitly:

batch = Batch(managed_lifecycle=True)
with batch:
    task = batch.function(square, 7)
    print(task.get())

# The context detached this client, but the shared runtime remains alive.
batch.destroy()

Batch.close() is deprecated and remains a no-op for compatibility. Use a context manager for scoped work, or call join() when a client is done and destroy() when a managed runtime should be shut down.

Worker Resource Configuration

By default, each Batch worker pool uses the physical cores available on its node. Pass the keyword-only workers_per_node option to limit the number of task worker processes in each pool, for example Batch(workers_per_node=4). This controls task workers and does not change the dedicated scheduler worker count configured by scheduler_workers.

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 70 Chained tasks via file dependencies
 1import dragon
 2from dragon import Batch
 3from multiprocessing import set_start_method
 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 71 Mixed Python + executable tasks in one Batch
 1import dragon
 2from dragon import Batch, ProcessTemplate
 3from multiprocessing import set_start_method
 4from dragon.native.process import 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.

  • terminate() — force shutdown when work should not be

    allowed to continue; the context manager uses this path when the block raises an exception.

  • 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()).