Graph-based Execution with Batch

Batch allows users to submit Python functions, executables, and parallel jobs for distributed execution with automatic parallelization. Simply call function(), process(), or job() to submit tasks, and Batch dispatches them to workers in the background automatically. Users specify data dependencies between tasks via read() and write() calls, and Batch infers an implicit directed acyclic graph (DAG) to maximize parallelism. Results are retrieved by calling .get() on a task’s handle, which blocks until the task completes. Use fence() to wait for all currently submitted tasks to finish before proceeding, and clear_results() to free the memory used to hold task results (implicitly calls fence()).

Batch always creates and owns its internal results DDict. By default it allocates one gibibyte per requested node for that store. For result-heavy workloads or tighter memory budgets, pass results_ddict_mem in bytes when constructing the Batch, for example Batch(results_ddict_mem=8 * 1024**3).

Below is a simple example using Batch to parallelize a list of functions. In this example, file-based dependencies force the functions to execute in order, but it demonstrates the basics of the API.

Listing 21 Basic Batch example with YAML template

# Generate the powers of a matrix and write them to disk
from dragon.workflows.batch import Batch
from pathlib import Path
from some_math_functions import gpu_matmul

import numpy as np

# A base directory, and files in it, will be used for communication of results.
# Optionally size the internal results DDict explicitly in bytes.
batch = Batch(results_ddict_mem=2 * 1024**3)
base_dir = Path("/some/path/to/base_dir")

a = np.array([j for j in range(100)])
m = np.vander(a)

# Import a function to lazily run gpu_matful tasks in parallel. Notice that the
# "ptd" in the file name refers to "parameterized task descriptor", since the yaml
# file describes a collection of tasks that are parameterized by the arguments to
# the returned function, gpu_matmul_func.
get_file = lambda base_dir, i, j: Path(base_dir) / Path(f"file_{i + j}")
gpu_matul_lazy = batch.import_func("gpu_matmul_ptd.yml", gpu_matmul, base_dir, get_file)

results = []
for i in range(1000):
    val = gpu_matmul_lazy(m, i)
    results.append(val)

# If there was an exception while running the task, it will be raised when we call val.get()
for val in results:
    try:
        print(f"{val.get()=}")
    except Exception as e:
        print(f"gpu_matmul failed with the following exception: {e}")

batch.join()
batch.destroy()

Notice that the call to import_func() requires a file called “gpu_matmul_ptd.yml”, which is a Parameterized Task Descriptor, or PTD, file. This type of file describes a collection of tasks, rather than a single task. The collection of tasks are parameterized by the arguments to the function returned by import_func(): each set of arguments generates a new task. Generated tasks are not run immediately, but rather are queued up in the background and run lazily in parallelized batches. Batches of tasks only run when you (1) call <task return value>.get(), try to access either a distributed dictionary or file that’s updated by a task, or call fence(). The PTD file “gpu_matmul_ptd.yml” is shown directly below, and a PTD template (with lots of explanations) is located in the examples directory.

Listing 22 Function YAML example

########################### import-time and call-time arguments specified here ###########################

# these arguments are passed to import_func
import_args:
    # function to run
    - gpu_matmul
    # base directory of files that will be accessed
    - base_dir
    # function to select file
    - get_file

# import_kwargs:

# these arguments are passed to gpu_matmul_lazy
args:
    - matrix
    - i

# kwargs:

######### definition of task executables, resources, dependencies, etc. (in terms of args above) #########

# type of task (should be function, process, or job)
type: function

executables:
    # here we "plug in" the names of the function and its arguments (specified above)
    - target: gpu_matmul
    args:
        - matrix
        - i

reads:
    - files:
    - get_file:
        - base_dir
        - i
        - 0

writes:
    - files:
    - get_file:
        - base_dir
        - i
        - 1

name: gpu_matmul_task

timeout:
    day: 0
    hour: 0
    min: 0
    sec: 30

And here is the same example as above, but done without a PTD file.

Listing 23 Basic Batch example without YAML template

# Generate the powers of a matrix and write them to disk
from dragon.workflows.batch import Batch
from pathlib import Path
from some_math_functions import gpu_matmul

import numpy as np

# A base directory, and files in it, will be used for communication of results.
# Optionally size the internal results DDict explicitly in bytes.
batch = Batch(results_ddict_mem=2 * 1024**3)
base_dir = Path("/some/path/to/base_dir")

# Knowledge of reads and writes to files is used by Batch to infer data dependencies
# and automatically parallelize tasks
get_read = lambda i: batch.read(base_dir, Path(f"file_{i}"))
get_write = lambda i: batch.write(base_dir, Path(f"file_{i+1}"))

a = np.array([j for j in range(100)])
m = np.vander(a)

# Submit tasks directly — Batch dispatches them to workers in the background
tasks = [batch.function(gpu_matmul, m, base_dir, i,
                        reads=[get_read(i)], writes=[get_write(i)], timeout=30)
         for i in range(1000)]

# Retrieve results — .get() waits for each task to complete if it hasn't yet
for task in tasks:
    try:
        print(f"result={task.get()}")
    except Exception as e:
        print(f"gpu_matmul failed with the following exception: {e}")

batch.join()
batch.destroy()

Tasks are submitted continuously (although batched in the background for better performance). Calls to function(), process(), and job() return immediately and Batch batches and dispatches submitted tasks to workers in the background. The lifecycle methods are: close(), which is deprecated and retained as a no-op for compatibility; join(), which waits only for work started by the calling client and then detaches that client from the shared runtime; and destroy(), which shuts down the shared Batch runtime after all clients have joined and all in-flight work completes.

Any mix of Python functions, executables, and parallel jobs can be submitted to Batch simultaneously, and dependencies can exist between tasks of any type, e.g., an MPI job can depend on the completion of a Python function if the MPI job reads from a file that the function writes to. MPI jobs are specified using the job() function. Likewise, the process() function submits a task for running a serial executable.

Each task has three handles for obtaining output: result, stdout, and stderr. Calling .get() on the task handle returns the task’s return value, blocking until the task completes if it has not finished yet. If an exception was thrown during the execution of a task, calling .get() on the task’s handle will re-raise that exception.

The initial creation of the Batch object sets up manager and worker processes. Batch objects can be passed between processes to allow multiple clients. Unpickling a Batch object at a destination process will register the new Batch client and allow the user to submit tasks to it. All clients must call join() when they are done. Calling close() is optional and has no effect beyond a deprecation warning. Any client can also call destroy() to shut down the shared runtime; it will block until all active clients have called join() and all work completes.

Data Dependencies

Dependencies between tasks are inferred based on the data being read and written by each task. Data reads and writes are specified by read() and write() (or task read and write, to directly associate reads and writes with a specific task). When a task is created, a list of Read or Write objects, created by read() and write(), can be specified:

task = batch.job(target=mpi_exec, num_procs=256, reads=<list of Reads>, writes=<list of Writes>)

After a task has been create, further Reads and Writes can be specified via task.read and task.write:

task.read(base_dir1, file_path1)
task.write(base_dir2, file_path2)

Batch takes a dataflow approach to parallelizing tasks. Like all dataflow systems, it assumes that tasks have no side effects (think IO) beyond those specified by read() and write() calls. So, for instance, if a process task reads from a file, then a corresponding Read object must be created using read() and appended to the list of Reads when creating the task (or added after task creation using task read). If this read isn’t associated with the task, then the file read could happen out-of-order relative to other operations on that file, e.g., the file read could occur before the data intended to be read is written to the file.

Argument-passing Dependencies

Beyond the type of dependencies described above, there is a second type of dependency: argument-passing dependencies. These dependencies are inferred when a result, stdout, or stderr object is passed as an argument to a Batch Function, Process, or Job. For example:

Listing 24 Example passing arguments between tasks

def foo():
    return "Hi-diddly-ho!!!"

def bar(hello: str):
    print(hello, flush=True)

batch = Batch()

task_foo = batch.function(foo)
task_bar = batch.function(bar, task_foo.result)
task_bar.get()

# prints: "Hi-diddly-ho!!!"

batch.join()
batch.destroy()

In the above example, the function bar will not run until foo completes, and bar will print the string returned by foo.

Synchronization

By default, results can be retrieved in a fine-grained manner with .get(), which blocks only for the specific task being waited on. When you need a hard synchronization point, i.e., before checkpointing state or starting a new phase of work that must not overlap with the previous one, use fence(). It blocks until every task submitted by this client completes, then clears internal dependency state so the next batch of tasks starts from a clean slate. To additionally free the memory used to hold task results, call clear_results(), which calls fence() and then clears the results dictionary.

Distributed Dictionary

The Distributed Dictionary provides a scalable, in-memory distributed key-value store with semantics that are generally similar to a standard Python dictionary. The Distributed Dictionary uses shared memory and RDMA to handle communication of keys and values, and avoids central coordination so there are no bottlenecks to scaling. Revisiting an example above for the Batch service, we will replace the file system as a means of inter-task communication with a Distributed Dictionary. The only change needed is in the task definitions - the rest of the workflow (submitting tasks, retrieving results, closing) is identical.

Listing 25 Example using Batch and DDict

# Generate the powers of a matrix and write them to a distributed dictionary
from dragon.data import DDict
from dragon.workflows.batch import Batch
from some_math_functions import gpu_matmul

import numpy as np

def gpu_matmul(original_matrix: np.ndarray, ddict: DDict, i: int):
   # read the current matrix stored in the Distributed Dictionary at key=i
   current_matrix = ddict[i]

   # actual matrix multiplication happens here
   next_matrix = do_dgemm(original_matrix, current_matrix)

   # write the next power of the matrix to the Distributed Dictionary at key=i+1
   ddict[i + 1] = new_matrix

# The Distributed Dictionary service will be used for communication of results
batch = Batch()
ddict = DDict()

# Knowledge of reads and writes to the Distributed Dictionary will also be used by
# the Batch service to determine data dependencies and how to parallelize tasks
get_read = lambda i: batch.read(ddict, i)
get_write = lambda i: batch.write(ddict, i + 1)

a = np.array([i for i in range(100)])
m = np.vander(a)

# Submit tasks — Batch dispatches them in the background
tasks = [batch.function(gpu_matmul, m, ddict, i,
                        reads=[get_read(i)], writes=[get_write(i)])
         for i in range(1000)]

# Retrieve results — .get() waits for each task to complete if needed
for task in tasks:
    try:
        print(f"result={task.get()}")
    except Exception as e:
        print(f"gpu_matmul failed with the following exception: {e}")

batch.join()
batch.destroy()

Inspecting the Topology

topology() returns a BatchTopology object that describes how Batch has mapped managers and worker pools onto the nodes of the allocation. The reported topology includes the dedicated scheduler host separately from the pool-backed subnode managers. This is useful for understanding how work will be distributed before submitting tasks.

The example below creates several Batch instances with different configurations, prints a summary of each resulting topology, and shuts down cleanly without submitting any real work.

Listing 26 topology.py: Batch topology API walkthrough

"""
batch_topology.py — Batch topology API walkthrough
===================================================

This example demonstrates how ``Batch.topology()`` exposes the node layout
chosen by Dragon when you create a ``Batch`` instance. It walks through the
current manager model: one dedicated scheduler colocated with the client plus
one subnode manager and one worker pool per requested node.

Concepts
--------
Worker pools
    Each requested node contributes one worker pool. Every pool contributes
    ``num_cpus // 2`` workers (one per physical core, assuming hyperthreading
    doubles the logical count).

Manager co-location
    The scheduler runs on the client host. Each subnode manager runs
    on the nodes backing its worker pool. No core is reserved for the
    manager — the full physical-core count of every pool node is available to
    workers.

Useful relationships
    ``total_managers`` is always ``total_nodes + 1`` in the current runtime:
    one dedicated scheduler plus one subnode manager per requested node.
    ``manager_hostnames`` and ``pool_hostnames`` therefore have the same
    length, and each subnode manager host lines up with exactly one worker
    pool.

Usage
-----
Run on an allocation of at least a few nodes::

    dragon topology.py

The script does not submit any real work — it just creates ``Batch``
objects, prints their ``topology()``, and tears them down cleanly.
"""

from dragon.workflows.batch import Batch, BatchTopology
from dragon.native.machine import System


def section(title: str) -> None:
    width = 72
    print()
    print("=" * width)
    print(f"  {title}")
    print("=" * width)


def show(label: str, batch: Batch) -> None:
    """Print a labelled summary of a Batch topology."""
    t: BatchTopology = batch.topology()
    print(f"\n[{label}]")
    print(t)
    print(f"  repr  : {repr(t)}")
    print(f"  total workers : {sum(t.workers_per_pool)}")


def explain_fields(t: BatchTopology) -> None:
    """Walk through the key relationships exposed by BatchTopology."""
    print(f"\n  t.total_nodes        = {t.total_nodes}   (requested worker-pool nodes)")
    print(f"  t.scheduler_hostname = {t.scheduler_hostname}   (dedicated scheduler host)")
    print(f"  t.total_managers     = {t.total_managers}   (1 scheduler + one subnode manager per pool)")
    print(
        f"  t.manager_hostnames  = {t.manager_hostnames}"
        "\n                           (subnode manager hostnames, one per worker pool)"
    )
    print(f"  t.pool_hostnames     = (list of {len(t.pool_hostnames)} pool(s))")
    for i, (mgr_host, hosts, wpp) in enumerate(zip(t.manager_hostnames, t.pool_hostnames, t.workers_per_pool)):
        print(f"    pool {i}: {wpp} worker(s) on {hosts}  [mgr: {mgr_host}]")
    print(
        f"  t.workers_per_pool   = {t.workers_per_pool}"
        "\n                           (physical cores x nodes for each pool)"
    )


def explain_invariants(t: BatchTopology) -> None:
    """Print the current topology invariants that users are most likely to care about."""
    print("\n  Current invariants")
    print(f"    total_managers = total_nodes + 1  ->  {t.total_managers} = {t.total_nodes} + 1")
    print(
        "    len(manager_hostnames) == len(pool_hostnames)"
        f"  ->  {len(t.manager_hostnames)} == {len(t.pool_hostnames)}"
    )
    print(
        "    one subnode manager backs each worker pool"
        f"  ->  {all(mgr == hosts[0] for mgr, hosts in zip(t.manager_hostnames, t.pool_hostnames))}"
    )
    print(f"    total workers = sum(workers_per_pool)  ->  {sum(t.workers_per_pool)}")


def main() -> None:
    alloc = System()
    avail = len(alloc._node_objs)
    cpus_per_node = max(1, alloc._node_objs[0].num_cpus // 2)

    print(f"Allocation: {avail} node(s), {cpus_per_node} physical core(s) per node")

    # 1. Default: all nodes, one node per pool
    section("1  Default — scheduler plus one pool per node")

    print(
        "\n  The default topology creates one dedicated scheduler manager on the\n"
        "  client host plus one subnode manager and one worker pool for every\n"
        "  requested node. That matches the intuitive model of requesting N\n"
        "  nodes and getting N node-local worker pools.\n"
    )

    # Keep b_default open; we reuse it for the field-by-field walkthrough
    # (section 5) to avoid an extra Batch lifecycle.
    b_default = Batch()
    show("Batch()", b_default)

    # 2. Explicit node count smaller than the allocation
    if avail >= 2:
        section("2  Explicit node count — use part of the allocation")

        subset = max(1, avail // 2)
        if subset >= avail:
            subset = avail - 1

        print(
            f"\n  num_nodes={subset} restricts the Batch to {subset} node(s) even though\n"
            f"  {avail} are available.  This is useful when you want to reserve\n"
            f"  nodes for other work running in parallel.\n"
        )
        b = Batch(num_nodes=subset)
        show(f"Batch(num_nodes={subset})", b)
        b.close()
        b.join()

    # 3. Manager count tracks requested nodes
    if avail >= 2:
        candidate = min(avail, max(2, avail // 2))
        section("3  Manager count follows requested nodes")
        print(
            f"\n  num_nodes={candidate} creates {candidate} worker pool(s), {candidate}\n"
            f"  subnode manager(s), and one extra scheduler manager. In other words,\n"
            f"  the topology has {candidate + 1} total managers for {candidate}\n"
            f"  requested nodes.\n"
        )
        b = Batch(num_nodes=candidate)
        show(f"Batch(num_nodes={candidate})", b)
        b.close()
        b.join()

    # 4. BatchTopology field-by-field walkthrough
    #
    # Reuses b_default from section 1 — no extra Batch creation needed.
    section("4  BatchTopology field-by-field walkthrough")

    t = b_default.topology()
    explain_fields(t)
    explain_invariants(t)

    b_default.close()
    b_default.join()


if __name__ == "__main__":
    main()