dragon.workflows.batch.Batch

Bases: object

Graph-based distributed scheduling for functions, executables, and parallel applications

Create a Batch instance for orchestrating functions, executables, and parallel applications with data dependencies with a directed acyclic graph (DAG).

Parameters:

num_nodes (Optional[int ]) – Number of nodes to use for this Batch instance. Defaults to all nodes in the allocation. Values larger than the allocation are silently clamped.
pool_nodes (Optional[int ]) – Reserved for future worker-pool grouping support. The current implementation overrides this to 1 so each requested node gets its own subnode manager and worker pool.
disable_telem (bool ) – Indicates if telemetry should be disabled for this Batch instance. Defaults to False.
scheduler_workers (Optional[int ]) – Number of workers in the scheduler (manager 0)’s local worker pool. Defaults to the total number of nodes in the allocation (one worker per node). Increase this to allow more concurrent multi-node jobs.
results_ddict_mem (Optional[int ]) – Total memory in bytes to allocate for the Batch-owned results DDict. When omitted, Batch allocates one gibibyte per requested node.
results_ddict_managers_per_pool (Optional[int ]) – Number of results-DDict manager shards to provision per worker pool (one pool per requested node). Must be between 1 and the number of batch workers per node; values outside that range are clamped (and the clamping is logged). When omitted, Batch uses default_results_ddict_managers_per_pool shards per pool (clamped to the worker count). Pass None explicitly to request one results-DDict manager per worker.
managed_lifecycle (bool ) – Controls the lifetime of the Batch instance shared by clients. When False (the default), that Batch instance shuts down automatically after the last client detaches. When True, it stays alive until some client calls Batch.destroy().
stdout (Optional[str | Path]) – Default stdout log file path for tasks created by this client. When omitted and task_logs=True, Batch creates per-task stdout files under runinfo/<batch-run-id>/client-<id>/task_logs/<task-kind>. When omitted and task_logs=False (the default), task stdout is not captured to a file and is forwarded to the client console instead.
stderr (Optional[str | Path]) – Default stderr log file path for tasks created by this client. When omitted and task_logs=True, Batch creates per-task stderr files under runinfo/<batch-run-id>/client-<id>/task_logs/<task-kind>. When omitted and task_logs=False (the default), task stderr is not captured to a file and is forwarded to the client console instead.
task_logs (bool ) – When True, Batch captures each task’s stdout/stderr to per-task files under runinfo/<batch-run-id>/client-<id>/task_logs and maintains a manifest.jsonl index, enabling the experimental log-discovery helpers (Batch.log_dir(), Batch.log_manifest_path(), Batch.iter_log_records(), Batch.find_logs(), and Batch.read_logs()). When False (the default), no runinfo directory or manifest is created and task output is only written to a file when an explicit stdout/stderr path is supplied; otherwise it is forwarded to the client console. These helper APIs may evolve as the logging interface matures. The log-discovery helpers raise RuntimeError while task logging is disabled.

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

import numpy as np

def gpu_matmul(m, base_dir, i):
    # do GPU work with matrix m and data from {base_dir}/file_{i}
    return matrix

# A base directory, and files in it, will be used for communication of results
batch = Batch()
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 — Batch dispatches them to workers in the background
tasks = [batch.options(reads=[get_read(i)], writes=[get_write(i)], timeout=30)
             .function(gpu_matmul, m, base_dir, 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()

Returns:: Returns None.
Return type:: None

Methods

`__init__`([num_nodes, pool_nodes, ...])	Create a Batch instance for orchestrating functions, executables, and parallel applications with data dependencies with a directed acyclic graph (DAG).
`clear_results`()	Wait for all outstanding tasks to complete then clear the results dict for this batch.
`close`()	Deprecated no-op retained for API compatibility.
`destroy`([timeout, force_timeout])	Gracefully destroy the Batch instance shared by clients.
`fence`([timeout])	Wait for all tasks submitted by this client to complete.
`find_logs`(tuid)	Return the manifest record for tuid, or `None` if it is unknown.
`function`(target, *args[, reads, writes, ...])	Creates a new function task.
`import_func`(ptd_file, *real_import_args, ...)	Loads the PTD dict and creates a `MakeTask` object for the parameterized task group specified by the PTD file and import arguments (`real_import_args` and `real_import_kwargs`).
`iter_log_records`()	Return the current client's task-log records from the manifest.
`job`(process_templates[, reads, writes, ...])	Creates a new job task.
`join`([timeout])	Wait for the completion of all operations started by this client, then detach this client from the Batch instance shared by clients.
`log_dir`()	Return the per-client directory for the experimental log-discovery layer.
`log_manifest_path`()	Return the path to the per-client JSONL manifest of task log files.
`options`([reads, writes, name, timeout, ...])	Return a submission proxy with Batch metadata bound ahead of time.
`poll`([timeout])	Wait up to `timeout` seconds for the next completed task tuid.
`process`(process_template[, reads, writes, ...])	Creates a new process task.
`read`(obj, *channels)	Indicates READ accesses of a specified set of channels on a communication object.
`read_logs`(tuid[, log_type, encoding])	Read stdout/stderr log contents for tuid from the shared filesystem.
`terminate`()	Force the termination of a Batch instance.
`topology`()	Return a `BatchTopology` describing the node placement of managers and worker pools in this Batch instance.
`write`(obj, *channels)	Indicates WRITE accesses of a specified set of channels on a communication object.

Create a Batch instance for orchestrating functions, executables, and parallel applications with data dependencies with a directed acyclic graph (DAG).

Parameters:

num_nodes (Optional[int ]) – Number of nodes to use for this Batch instance. Defaults to all nodes in the allocation. Values larger than the allocation are silently clamped.
pool_nodes (Optional[int ]) – Reserved for future worker-pool grouping support. The current implementation overrides this to 1 so each requested node gets its own subnode manager and worker pool.
disable_telem (bool ) – Indicates if telemetry should be disabled for this Batch instance. Defaults to False.
scheduler_workers (Optional[int ]) – Number of workers in the scheduler (manager 0)’s local worker pool. Defaults to the total number of nodes in the allocation (one worker per node). Increase this to allow more concurrent multi-node jobs.
results_ddict_mem (Optional[int ]) – Total memory in bytes to allocate for the Batch-owned results DDict. When omitted, Batch allocates one gibibyte per requested node.
results_ddict_managers_per_pool (Optional[int ]) – Number of results-DDict manager shards to provision per worker pool (one pool per requested node). Must be between 1 and the number of batch workers per node; values outside that range are clamped (and the clamping is logged). When omitted, Batch uses default_results_ddict_managers_per_pool shards per pool (clamped to the worker count). Pass None explicitly to request one results-DDict manager per worker.
managed_lifecycle (bool ) – Controls the lifetime of the Batch instance shared by clients. When False (the default), that Batch instance shuts down automatically after the last client detaches. When True, it stays alive until some client calls Batch.destroy().
stdout (Optional[str | Path]) – Default stdout log file path for tasks created by this client. When omitted and task_logs=True, Batch creates per-task stdout files under runinfo/<batch-run-id>/client-<id>/task_logs/<task-kind>. When omitted and task_logs=False (the default), task stdout is not captured to a file and is forwarded to the client console instead.
stderr (Optional[str | Path]) – Default stderr log file path for tasks created by this client. When omitted and task_logs=True, Batch creates per-task stderr files under runinfo/<batch-run-id>/client-<id>/task_logs/<task-kind>. When omitted and task_logs=False (the default), task stderr is not captured to a file and is forwarded to the client console instead.
task_logs (bool ) – When True, Batch captures each task’s stdout/stderr to per-task files under runinfo/<batch-run-id>/client-<id>/task_logs and maintains a manifest.jsonl index, enabling the experimental log-discovery helpers (Batch.log_dir(), Batch.log_manifest_path(), Batch.iter_log_records(), Batch.find_logs(), and Batch.read_logs()). When False (the default), no runinfo directory or manifest is created and task output is only written to a file when an explicit stdout/stderr path is supplied; otherwise it is forwarded to the client console. These helper APIs may evolve as the logging interface matures. The log-discovery helpers raise RuntimeError while task logging is disabled.

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

import numpy as np

def gpu_matmul(m, base_dir, i):
    # do GPU work with matrix m and data from {base_dir}/file_{i}
    return matrix

# A base directory, and files in it, will be used for communication of results
batch = Batch()
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 — Batch dispatches them to workers in the background
tasks = [batch.options(reads=[get_read(i)], writes=[get_write(i)], timeout=30)
             .function(gpu_matmul, m, base_dir, 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()

Returns:: Returns None.
Return type:: None

log_dir() → Path [source] 

Return the per-client directory for the experimental log-discovery layer.

This is the root of the current client’s Batch logging area, typically runinfo/<batch-run-id>/client-<id>/task_logs.

Raises:: RuntimeError – If task logging is disabled (Batch(task_logs=False)).

log_manifest_path() → Path [source] 

Return the path to the per-client JSONL manifest of task log files.

The manifest is the index used by the experimental log-discovery helpers. Each record stores the task id, task kind, task/target names, log paths, and any host or completion metadata known to this client.

Raises:: RuntimeError – If task logging is disabled (Batch(task_logs=False)).

iter_log_records() → list [dict [str , Any ]][source] 

Return the current client’s task-log records from the manifest.

This is the primary Python helper in the experimental log-discovery layer for inspecting all known task-log metadata without reading the JSONL file manually.

Raises:: RuntimeError – If task logging is disabled (Batch(task_logs=False)).

find_logs(tuid: str ) → dict [str , Any ] | None [source] 

Return the manifest record for tuid, or None if it is unknown.

Use this experimental helper to locate a task’s stdout/stderr files and any resolved host metadata by task id.

Raises:: RuntimeError – If task logging is disabled (Batch(task_logs=False)).

read_logs(tuid: str , log_type: str | None = None, encoding: str = 'utf-8') → dict [str , str | None ][source] 

Read stdout/stderr log contents for tuid from the shared filesystem.

This is part of Batch’s experimental log-discovery layer.

When log_type is "stdout" or "stderr", only that log file is read. Otherwise both log files are read.

The return value is always a dictionary with stdout and stderr keys. The selected log content is returned as a string, and any unrequested or unavailable side is returned as None.

Raises:: RuntimeError – If task logging is disabled (Batch(task_logs=False)).

poll(timeout: float = 1000000000.0) → str | None [source] : Wait up to timeout seconds for the next completed task tuid.

read(obj, *channels) → DataAccess[source] 

Indicates READ accesses of a specified set of channels on a communication object. These accesses are not yet associated with a given task.

Parameters:

obj – The communication object being accessed.
*channels –
A tuple of channels on the communcation object that will be read from.

Returns:

Returns an descriptor for the data access that can be passed to (in a list) when creating a new task.

Return type:

DataAccess

write(obj, *channels) → DataAccess[source] 

Indicates WRITE accesses of a specified set of channels on a communication object. These accesses are not yet associated with a given task.

Parameters:

obj – The communication object being accessed.
*channels –
A tuple of channels on the communcation object that will be writtent o.

Returns:

Returns an descriptor for the data access that can be passed to (in a list) when creating a new task.

Return type:

DataAccess

fence(timeout: float = 1000000000.0) → None [source] 

Wait for all tasks submitted by this client to complete. Tasks submitted after the FenceRequest is enqueued will be handled after the fence finishes. The client-side compile worker clears per-client compile state after the scheduler acknowledges the fence.

Parameters:: timeout (float ) – Timeout in seconds for each blocking operation. Defaults to 1e9.
Returns:: Returns None.
Return type:: None

close() → None [source] 

Deprecated no-op retained for API compatibility.

Client detachment is now handled by Batch.join(), which flushes pending local submissions, waits for this client’s work to complete, and unregisters the client from the Batch instance shared by clients.

Returns:: Returns None.
Return type:: None

Deprecated since version ``Batch.close()``: no longer changes Batch state. Use Batch.join() when a client is done submitting work.

join(timeout: float = 1000000000.0) → None [source] 

Wait for the completion of all operations started by this client, then detach this client from the Batch instance shared by clients.

After join() returns, this handle can no longer submit additional work. In the default unmanaged mode, the client-shared Batch instance shuts down automatically after the last client detaches. In managed mode, any client may later call Batch.destroy(), even after that handle has already joined. In unmanaged mode, fetch any task results you need before calling join() on the last attached client, because automatic shutdown destroys the shared results DDict. Use managed_lifecycle=True when results must remain retrievable after a handle has joined.

Parameters:: float – A timeout value for waiting on batch completion. Defaults to 1e9.
Returns:: Returns None.
Return type:: None

destroy(timeout: float = 1000000000.0, force_timeout: float | None = None) → None [source] 

Gracefully destroy the Batch instance shared by clients.

Managed-lifecycle runtimes are shut down after all currently attached clients detach. The calling client is detached first if needed. When force_timeout is provided, managers shut down after that many seconds even if other clients remain attached.

Calling destroy() on an unmanaged runtime raises RuntimeError; the default mode shuts down automatically after the last client joins.

Parameters:

timeout (float ) – Timeout in seconds for waiting on manager shutdown.
force_timeout (Optional[float ]) – Optional grace period in seconds before forcing runtime shutdown even if other clients remain attached.

Returns:

Returns None.

Return type:

None

terminate() → None [source] 

Force the termination of a Batch instance.

Returns:: Returns None.
Return type:: None

clear_results() → None [source] 

Wait for all outstanding tasks to complete then clear the results dict for this batch. This can be used to free up memory after tasks have completed and their results have been retrieved. Any task result that has not yet been fetched with Task.get() becomes unavailable after this call.

Returns:: Returns None.
Return type:: None

topology() → BatchTopology[source] 

Return a BatchTopology describing the node placement of managers and worker pools in this Batch instance.

The returned object reports:

the total number of nodes used,
the hostname where the dedicated scheduler runs,
the hostname of the pool node where each subnode manager process runs, and
for each worker pool, the hostnames of the nodes that make up that pool.

Each requested node gets its own worker pool and its own subnode manager. The scheduler manager is an extra process colocated with the first Batch node and is not counted as one of the worker pools. All physical cores on every pool node are available as workers; no core is reserved for the manager.

Example:

batch = Batch(num_nodes=8)
print(batch.topology())
# Batch Topology:
#   Total nodes  : 8
#   Managers     : 9 total (1 scheduler + 8 subnode)
#   Scheduler    : hotlum0001
#   Worker pools : 8 pool(s) (1 dedicated subnode manager per pool)
#     Pool 0 (1 node(s), 32 worker(s)): hotlum0001  [mgr: hotlum0001]
#     Pool 1 (1 node(s), 32 worker(s)): hotlum0002  [mgr: hotlum0002]
#     Pool 2 (1 node(s), 32 worker(s)): hotlum0003  [mgr: hotlum0003]
#     ...

Tip

Install python-hostlist (pip install python-hostlist) to have hostnames in the output compressed into Slurm-style bracket notation, e.g. hotlum[0001-0008] instead of a long comma-separated list. This is especially helpful on large allocations.

Returns:: A BatchTopology object describing the placement.
Return type:: BatchTopology

Return a submission proxy with Batch metadata bound ahead of time.

This is the preferred API for task metadata. It cleanly separates Batch execution options from user payload arguments, following the same broad idea as Ray’s .options(...) API.

For process and job submissions, stdout and stderr supplied via Batch.options() act as task-level defaults. If an individual ProcessTemplate already specifies stdout or stderr, the template setting takes precedence.

Creates a new function task. Arguments for the function that are of type Task will create a dependency for this task on the output of the task specified by the argument. Further, the output of the specified task will be passed in place of the Task argument when the function executes.

Deprecated since version Passing: Batch metadata such as reads, writes, name, timeout, stdout, or stderr directly to Batch.function() is deprecated. Use Batch.options() to separate Batch metadata from user function arguments.

Parameters:

func – The function to associate with the object.
*args –
The arguments for the function.
reads (Optional[list ]) – A list of Read objects created by calling Batch.read().
writes (Optional[list ]) – A list of Write objects created by calling Batch.write().
name (Optional[str ]) – A human-readable name for the task.
stdout (Optional[str | Path]) – Optional stdout log file path for this task. When omitted, Batch uses the client default. If the client was created with task_logs=True and no default is set, Batch auto-generates a per-task file in the client run directory; otherwise the task’s stdout is forwarded to the client console. stdout is a reserved Batch submission keyword and is not passed through to the user function’s **kwargs.
stderr (Optional[str | Path]) – Optional stderr log file path for this task. When omitted, Batch uses the client default. If the client was created with task_logs=True and no default is set, Batch auto-generates a per-task file in the client run directory; otherwise the task’s stderr is forwarded to the client console. stderr is a reserved Batch submission keyword and is not passed through to the user function’s **kwargs.

:raises SubmitAfterCloseError: If this client handle has already: been detached from the Batch runtime by Batch.join(), Batch.destroy(), or Batch.terminate().

Returns:: The new function task.
Return type:: Function

Creates a new process task. Arguments for a process passed using ProcessTemplate.args that are of type Task will create a dependency for this task on the output of the task specified by the argument. Further, the output of the specified task will be passed in place of the Task argument when the process executes.

Deprecated since version Passing: Batch metadata such as reads, writes, name, timeout, stdout, or stderr directly to Batch.process() is deprecated. Use Batch.options(...).process(process_template)`() instead.

Parameters:

process_template (ProcessTemplate) – A Dragon ProcessTemplate to describe the process to be run.
reads (Optional[list ]) – A list of Read objects created by calling Batch.read().
writes (Optional[list ]) – A list of Write objects created by calling Batch.write().
name (Optional[str ]) – A human-readable name for the task.
stdout (Optional[str | Path]) – Optional stdout log file path for this task. When omitted, Batch uses the client default. If the client was created with task_logs=True and no default is set, Batch auto-generates a per-task file in the client run directory; otherwise the task’s stdout is forwarded to the client console.
stderr (Optional[str | Path]) – Optional stderr log file path for this task. When omitted, Batch uses the client default. If the client was created with task_logs=True and no default is set, Batch auto-generates a per-task file in the client run directory; otherwise the task’s stderr is forwarded to the client console.

:raises SubmitAfterCloseError: If this client handle has already: been detached from the Batch runtime by Batch.join(), Batch.destroy(), or Batch.terminate().

Returns:: The new process task.
Return type:: Job

Creates a new job task. Arguments for a process passed using ProcessTemplate.args that are of type Task will create a dependency for this task on the output of the task specified by the argument. Further, the output of the specified task will be passed in place of the Task argument when the job executes.

Deprecated since version Passing: Batch metadata such as reads, writes, name, timeout, pmi, stdout, or stderr directly to Batch.job() is deprecated. Use Batch.options(...).job(process_templates)`() instead.

If any process_template includes an explicit host-name policy (Policy(placement=Policy.Placement.HOST_NAME, host_name=...)), Batch preserves that placement request. For multi-node jobs, Batch also treats those hostnames as allocation constraints and reserves those specific nodes before launch. Leave host_name unset if you want Batch to choose job placement automatically.

Parameters:: process_templates – A list of pairs of the form (num_procs, process_template), where

process_template is template for num_procs processes in the job. The process template is based on Dragon’s ProcessTemplate. :type process_templates: list :param reads: A list of Read objects created by calling Batch.read(). :type reads: Optional[list] :param writes: A list of Write objects created by calling Batch.write(). :type writes: Optional[list] :param name: A human-readable name for the task. :type name: Optional[str] :param pmi: The PMI backend to use for launching MPI jobs. Defaults to PMIBackend.CRAY.

Set to PMIBackend.PMIX for systems using PMIx, or None to disable PMI.

Parameters:

stdout (Optional[str | Path]) – Optional stdout log file path for this task. Multi-rank jobs share one stdout file across all ranks. When omitted, Batch uses the client default. If the client was created with task_logs=True and no default is set, Batch auto-generates a per-task file in the client run directory; otherwise the job’s stdout is forwarded to the client console.
stderr (Optional[str | Path]) – Optional stderr log file path for this task. Multi-rank jobs share one stderr file across all ranks. When omitted, Batch uses the client default. If the client was created with task_logs=True and no default is set, Batch auto-generates a per-task file in the client run directory; otherwise the job’s stderr is forwarded to the client console.

:raises SubmitAfterCloseError: If this client handle has already been detached from the Batch runtime by Batch.join(), Batch.destroy(), or Batch.terminate().

Returns:: The new job task.
Return type:: Job

import_func(ptd_file: str , *real_import_args, **real_import_kwargs) → MakeTask[source] 

Loads the PTD dict and creates a MakeTask object for the parameterized task group specified by the PTD file and import arguments (real_import_args and real_import_kwargs). The group of tasks is parameterized by the arguments passed to the MakeTask object’s MakeTask.__call__() method, with a different task created for each unique collection of arguments. The name of this function comes from the fact that the MakeTask.__call__() method of the MakeTask object is meant to “look and feel” like calling the task and getting a return value without blocking, i.e., writing a serial program that runs locally, even though the tasks are lazily executed in parallel by the remote batch workers.

Parameters:

ptd_file (str ) – Specifies the parameterized task group.
x – Positional arguments to replace the identifiers listed under the “import_args”

key in the PTD file. :param x: Keyword arguments to replace the key/value identifiers specified under the “import_args” key in the PTD file.

Returns:: Returns a MakeTask object representing the parameterized group of tasks

specified by the PTD file and import arguments. :rtype: MakeTask