Batch and multi-GPU execution

This page covers the execution layer above a single Simulation.run(). Every name here is imported from witwin.maxwell (for example mw.BatchExecution, mw.run_many). You need it when you run many related FDTD simulations and want to reuse GPU state between them, or spread independent simulations over several GPUs.

There are three paths:

To split one large simulation across GPUs, use FDTDParallelConfig on Simulation.fdtd(..., parallel=...) instead; see Simulation. A simulation with a parallel config cannot enter a batch or run_many.

Batch lifecycle#

A batch keeps the ordinary Scene -> Simulation -> Result model. Each variant still produces its own Result.

1. Declare

   Call Simulation.batch(...). It returns a BatchSimulation and only validates the declaration: types, names, shapes and the device policy. It does not touch CUDA.

2. Plan (optional)

   Call batch.plan() for a dry BatchPlanReport. With grouping="strict" it raises ValueError naming the first field where a variant differs from variant 0 (for example frequencies, dt, boundary, grid). The report is provisional: resolved is always False.

3. Prepare

   Call batch.prepare(). It returns a PreparedBatch: it checks the supported feature subset, creates the CUDA stream and memory budget, compiles the first cohort and fixes the cohort size. Use it as a context manager so it closes.

4. Run or iterate

   Call prepared.run() for a BatchResult, or prepared.iter_results() to receive one BatchItem per variant as it finishes. A prepared batch can run again, optionally with new BatchParameters of the same schema.

5. Use results and backward

   Each result owns its tensors and stays valid after the slot is reused or closed. If the batch inputs require gradients and autograd is enabled, the monitor outputs carry a graph: call torch.autograd.grad or .backward() on a loss built from them, even after close().

batch.run() runs prepare, run and close in one call. batch.iter_results() prepares a batch that the iterator owns and closes when the iterator closes.

This example is trimmed from tests/api/public/test_batch_execution.py and tests/fdtd/cuda/test_batch_runtime.py. It solves three source amplitudes of one scene, two variants at a time.

import torch
import witwin.maxwell as mw
 
scene = mw.Scene(domain=mw.Domain(bounds=((-0.2, 0.2),) * 3), grid=mw.GridSpec.uniform(0.025),
                 boundary=mw.BoundarySpec.pml(num_layers=3), device="cuda:0")
scene.add_source(mw.PointDipole(name="drive", position=(-0.05, 0.0, 0.0), polarization="Ez",
                                source_time=mw.GaussianPulse(frequency=1e9, fwidth=1e9)))
scene.add_monitor(mw.PointMonitor(name="probe", position=(0.1, 0.0, 0.0), fields=("Ez",)))
sim = mw.Simulation.fdtd(scene, frequencies=(1e9,), run_time=mw.TimeConfig(time_steps=80, dt=2e-11))
 
scales = mw.BatchParameters(source_scales={"drive": torch.tensor([1.0, 2.0, 3.0], device="cuda:0")})
batch = mw.Simulation.batch(sim, parameters=scales,
                            execution=mw.BatchExecution(devices=("cuda:0",), batch_size=2))
print(batch.plan().to_markdown())
with batch.prepare() as prepared:
    results = prepared.run()
    results.results()  # raises the first DistributedFailure, if any
    ez = results.stack(lambda r: r.raw_monitor("probe")["components"]["Ez"])

For gradients, give the batch a trainable input and build a loss from the results (from tests/gradients/test_batch_adjoint.py):

density = torch.full((3, 2, 2, 2), 0.4, device="cuda:0", requires_grad=True)
batch = mw.Simulation.batch(sim, parameters=mw.BatchParameters(region_densities={"rho": density}),
                            execution=mw.BatchExecution(devices=("cuda:0",), batch_size=2))
results = batch.run()
results.results()
loss = sum(r.raw_monitor("probe")["components"]["Ez"].abs().square().sum() for r in results)
(grad,) = torch.autograd.grad(loss, density)

Here sim.scene must hold a MaterialRegion named "rho" whose density has shape (2, 2, 2).

Supported batch scenes#

prepare() (and therefore run() and iter_results()) checks every simulation and raises for anything outside this subset:

Reverse execution uses native per-lane checkpoint adjoints. Supported trainable inputs are material densities, dielectric geometry and named source scales.

Simulation.batch#

Simulation.batch(simulations, *, execution, parameters=None, variant_ids=None, variation=None, outputs=None) -> BatchSimulation

Static method that declares a batch. It is the same as calling BatchSimulation(...) with the same arguments; see BatchSimulation for each argument.

Returns BatchSimulation.

BatchSimulation#

class BatchSimulation(simulations, *, execution, parameters=None, variant_ids=None, variation=None, outputs=None)

An inspectable batch submission. Construction validates the declaration and does not initialize CUDA.

There are two input forms:

• Template: one Simulation plus parameters. The variant count is parameters.variant_count; variant i is the template with slice i of every parameter tensor applied.
• List: a nonempty sequence of Simulation objects and no parameters. Each simulation is one variant.

Raises TypeError for a wrong argument type. ValueError when a named region, structure, source, monitor or port does not identify exactly one object, when a template parameter's spatial shape differs from the region density, when a simulation sets full_field_dft but outputs.fields is empty, or when outputs.result_device="cpu" is combined with trainable inputs. NotImplementedError when a source_scales target is a port-generated source.

Attributes: execution, outputs, parameters, variation, template (the template simulation or None), simulations (tuple), variant_ids (tuple of str).

batch = mw.Simulation.batch([sim_a, sim_b, sim_c],
                            execution=mw.BatchExecution(devices=("cuda:0",), grouping="compatible"))
report = batch.plan()

BatchSimulation.scene#

BatchSimulation.scene(index) -> Scene

Returns a cloned Scene for variant index. For a template batch the clone carries slice index of each region_densities tensor (gradients still flow back to the batch tensor) and a source amplitude multiplied by source_scales[name][index]. The original scene is not changed.

Raises TypeError if index is not an int; IndexError if it is out of range.

BatchSimulation.plan#

BatchSimulation.plan() -> BatchPlanReport

Builds a dry, provisional plan. It never prepares a scene, allocates solver memory or queries the GPU. Variants are grouped by a digest of their declarations (domain, grid, boundary, symmetry, frequencies, config, sources, monitors, ports, excitations and output plan). Material and structure contents are not compared at this stage. Amplitudes of sources named in parameters.source_scales or variation.sources are ignored when comparing.

Raises ValueError with grouping="strict" when a variant differs from variant 0; the message names the first differing field.

BatchSimulation.prepare#

BatchSimulation.prepare() -> PreparedBatch

Runs plan() and builds a PreparedBatch. With execution.tune=True it first measures candidate policies with complete forward runs and returns a context built from the selected policy.

BatchSimulation.run#

BatchSimulation.run() -> BatchResult

Prepares, runs every variant once and closes the prepared batch.

BatchSimulation.iter_results#

BatchSimulation.iter_results(*, order='completion') -> BatchResultIterator

Prepares a batch and returns an iterator that owns it; closing the iterator also closes the prepared batch.

Raises ValueError for another order value.

BatchSimulation length#

len(batch) -> int

Number of variants.

BatchExecution#

class BatchExecution(*, devices, batch_size='auto', max_batch_size=32, memory_limit_bytes=None, reserve_bytes=536870912, grouping='strict', scheduling='throughput', slots_per_device=1, graph='auto', graph_steps='auto', tune=False, fail_fast=True)

Frozen, keyword-only device and capacity policy. Construction validates values but does not initialize CUDA.

Raises TypeError or ValueError for invalid values.

execution = mw.BatchExecution(devices=("cuda:0", "cuda:1"), batch_size=4, slots_per_device=2,
                              memory_limit_bytes=8 * 1024**3)

BatchParameters#

class BatchParameters(*, region_densities=None, source_scales=None)

Frozen, keyword-only per-variant tensor inputs for a template batch. Every tensor has a leading variant dimension, and all tensors must share the same leading size. Tensors are kept as given (no copy), so views and requires_grad survive.

Tensors must be strided, real floating point, on CPU or CUDA, and have no empty dimension. At least one tensor is required.

Raises TypeError or ValueError for a wrong dtype, rank, device, empty input or mismatched leading size.

BatchParameters.variant_count#

BatchParameters.variant_count -> int

Shared leading size V. Set at construction.

BatchParameters.requires_grad#

BatchParameters.requires_grad -> bool

True if any tensor requires gradients.

BatchParameters.schema#

BatchParameters.schema -> tuple

Tuple of (kind, name, shape, dtype) entries. PreparedBatch.run(parameters=...) requires new parameters with the same schema.

BatchVariation#

class BatchVariation(*, material_regions=(), structures=None, sources=())

Frozen declaration of which named scene objects change between variants. For a list batch, named material regions and structures let the runtime reuse compiled materials outside the changed region.

Raises TypeError if a name list is a string or a structures value is not a geometry; ValueError for empty or duplicate names.

BatchOutput#

class BatchOutput(*, monitors='all', ports='all', fields=(), materials=False, result_device='local', max_pending_results=2)

Frozen selection of what each variant's result keeps and where it lives.

BatchPlanReport#

class BatchPlanReport(_encoded)

Immutable, JSON-backed planning snapshot. You get it from BatchSimulation.plan() or PreparedBatch.report; you do not build it yourself. Each top-level key is also readable as an attribute (report.groups, report.variant_count), and each read returns an independent copy.

Keys from plan():

PreparedBatch.report adds execution_tier, the chosen batch_size, graph, slots_per_device, memory (budget statistics) and capacity_retries. A multi-worker report adds total_capacity and per-worker reports under workers.

BatchPlanReport.to_dict#

BatchPlanReport.to_dict() -> dict

Returns a new dict decoded from the stored JSON.

BatchPlanReport.to_markdown#

BatchPlanReport.to_markdown() -> str

Returns a short Markdown summary: variant count, a group table and the warnings.

BatchPlanReport.from_dict#

BatchPlanReport.from_dict(payload) -> BatchPlanReport

Class method that encodes a JSON-serializable dict. NaN values are rejected.

PreparedBatch#

class PreparedBatch(batch)

A reusable CUDA context for one BatchSimulation. Create it with BatchSimulation.prepare(). With more than one device, or slots_per_device=2, you get a subclass that runs one worker per slot on separate threads; the public methods are the same.

Construction compiles the first cohort, measures capacity when batch_size='auto', and builds report. Only one run() or iter_results() may be active at a time. Do not mutate the input Scene objects while the prepared batch is open.

Attributes: batch, report (BatchPlanReport), capacity (resolved cohort size; summed over workers in the multi-worker case), device (first execution device).

Raises RuntimeError without CUDA; NotImplementedError or ValueError for unsupported scenes (see Supported batch scenes).

PreparedBatch.run#

PreparedBatch.run(*, parameters=None) -> BatchResult

Runs every variant and returns the results in input order. Failures are stored as entries, not raised.

Raises RuntimeError if the batch is closed or another run or iteration is active; ValueError if parameters change the schema.

with batch.prepare() as prepared:
    first = prepared.run()
    second = prepared.run(parameters=mw.BatchParameters(source_scales={"drive": new_scales}))

PreparedBatch.iter_results#

PreparedBatch.iter_results(*, parameters=None, order='completion') -> BatchResultIterator

Starts a demand-driven run. With one slot, cohorts are solved as you pull items; with several workers, at most outputs.max_pending_results finished results wait in the queue.

Raises as for run(), plus ValueError for another order.

PreparedBatch.request_cancel#

PreparedBatch.request_cancel() -> None

Asks the active run to stop. Variants not yet finished are reported with FailureKind.CANCELLED. Safe to call from another thread.

PreparedBatch.close#

PreparedBatch.close() -> None

Releases slots, caches, graphs and the memory budget. Results already returned stay valid. Calling it twice is a no-op. with batch.prepare() as prepared: calls it on exit.

Raises RuntimeError while an iterator is still open.

PreparedBatch.records#

PreparedBatch.records -> tuple[ExecutionRecord, ...]

Records of the latest run, in input order.

BatchResultIterator#

class BatchResultIterator

Iterator returned by iter_results(). It yields one BatchItem per variant and keeps no unbounded queue of finished results. Use it in a with block or call close() when you stop early; variants not yet delivered are then recorded as cancelled.

with prepared.iter_results(order="input") as items:
    for item in items:
        item.raise_for_error()
        save(item.variant_id, item.result)

BatchResultIterator.close#

BatchResultIterator.close() -> None

Stops the run and releases the prepared batch for the next run. If the iterator came from BatchSimulation.iter_results(), the prepared batch is closed too.

Raises RuntimeError if called while next() is running in another thread; call PreparedBatch.request_cancel() first.

BatchResultIterator.records#

BatchResultIterator.records -> tuple[ExecutionRecord, ...]

Records so far; after close(), one record per variant.

BatchItem#

class BatchItem(index, variant_id, result, failure, record)

Frozen record of one completed, failed or cancelled variant. Exactly one of result and failure is set.

BatchItem.raise_for_error#

BatchItem.raise_for_error() -> None

Raises failure if it is set.

BatchResult#

class BatchResult(entries, records, variant_ids)

A ResultSequence of ordinary Result objects with stable variant IDs. You get it from run(). Each result's metadata holds batch_index and variant_id, and its solver_stats hold batch keys such as batch_execution_tier. Results keep no solver or prepared scene. Indexing or iterating raises a failed variant's DistributedFailure; use entries for the non-raising view.

Attribute: variant_ids (tuple of str).

BatchResult.stack#

BatchResult.stack(selector, *, device=None) -> torch.Tensor

Applies selector to each result and stacks the tensors along a new first axis.

Raises the first DistributedFailure; TypeError if the selector returns a non-tensor; ValueError for an empty batch, mismatched shapes or dtypes, or mixed devices without device.

BatchResult.save#

BatchResult.save(directory) -> None

Writes result-<index>.pt for each completed variant and a manifest.json with IDs, devices, statuses, timings and failures.

Raises FileExistsError if the directory is not empty.

BatchResult.load#

BatchResult.load(directory, *, scenes, map_location='cpu') -> BatchResult

Class method that reads a directory written by save().

Raises ValueError for an unknown manifest version, non-contiguous indices or a wrong number of scenes.

result.save(tmp_path)
loaded = mw.BatchResult.load(tmp_path, scenes=[s.scene for s in sims])

MultiGPUExecution#

class MultiGPUExecution(devices, max_concurrency, placement='memory_aware', fail_fast=False, per_device_concurrency=1)

Frozen descriptor for running independent simulations over several GPUs with run_many. Construction validates structure only; CUDA availability and memory are checked when the run starts.

Raises TypeError or ValueError for invalid values.

MultiGPUExecution.ensemble#

MultiGPUExecution.ensemble(*, devices, max_concurrency=None, placement='memory_aware', fail_fast=False, per_device_concurrency=1) -> MultiGPUExecution

Class method and the usual constructor. max_concurrency=None becomes len(devices) * per_device_concurrency.

execution = mw.MultiGPUExecution.ensemble(devices=("cuda:0", "cuda:1"))

MultiGPUExecution.build_pool#

MultiGPUExecution.build_pool(*, require_cuda=True) -> DevicePool

Builds the device pool that run_many uses. With require_cuda=True it snapshots each device's free and total memory. You do not need to call it yourself.

run_many#

run_many(simulations, *, execution) -> ResultSequence

Runs independent simulations concurrently over the devices in execution and returns their outcomes in submission order: entry i belongs to simulations[i]. Each task leases one device, runs Simulation.run() with its scene placed on that device, and releases the device. The caller's scene object is not changed. When more than one task can run at a time, CUDA Graph capture is suspended for the whole run.

Returns ResultSequence of Result or DistributedFailure entries.

Raises TypeError if execution is not a MultiGPUExecution or an item is not a Simulation; ValueError for a SceneModule input, trainable parameters, a parallel config or a shared Scene. Task failures are not raised here; they are stored in the sequence.

execution = mw.MultiGPUExecution.ensemble(devices=("cuda:0", "cuda:1"), max_concurrency=2)
seq = mw.run_many([sim_a, sim_b], execution=execution)
if not seq.failed:
    peaks = [float(r.fields["EZ"].abs().max()) for r in seq]

ResultSequence#

class ResultSequence(entries, records)

Ordered container of task outcomes. Indexing and iteration return the Result and raise the slot's DistributedFailure for a failed task, so a failure is never silently skipped. len(seq) is the number of tasks.

ResultSequence.entries#

ResultSequence.entries -> tuple

Non-raising view: each entry is a Result or a DistributedFailure.

ResultSequence.records#

ResultSequence.records -> tuple[ExecutionRecord, ...]

One record per task, in submission order.

ResultSequence.failures#

ResultSequence.failures -> tuple[DistributedFailure, ...]

All failures in order.

ResultSequence.failed#

ResultSequence.failed -> bool

True if any task failed or was cancelled.

ResultSequence.results#

ResultSequence.results() -> tuple

Returns every result in order, or raises the first failure.

ExecutionRecord#

class ExecutionRecord(index, device, status, estimated_bytes=None, wall_time_s=None, device_time_s=None, failure=None)

Frozen bookkeeping for one task or batch variant.

The timing fields are measurements only. Comparing them as a speedup needs exclusive use of the GPUs.

ExecutionRecord.completed#

ExecutionRecord.completed -> bool

True when status == "completed".

DistributedFailure#

class DistributedFailure(*, index, kind, device=None, exception=None, message=None)

Exception describing one failed task. It keeps the original exception as exception and as __cause__.

Attributes: index, kind, device, exception, detail.

try:
    results[0]
except mw.DistributedFailure as failure:
    print(failure.index, failure.kind, failure.device, failure.detail)

FailureKind#

class FailureKind(str, Enum)

Failure category. str(kind) returns the value.

AggregatedSceneGradient#

class AggregatedSceneGradient(gradient, reduction_order, reduction_device, column_devices, port_names)

Frozen result of an array scene-gradient reduction.

aggregate_scene_gradient_vjp#

aggregate_scene_gradient_vjp(basis, *, columns, parameters, weights, objective=None, field_cotangents=None, reduction_order=None, reduction_device=None) -> AggregatedSceneGradient

Reduces per-column vector-Jacobian products (VJPs) for an array objective. It combines the columns into E = sum_n w_n e_n on the reduction device, takes the cotangent of the objective with respect to E, seeds column n with conj(w_n) times that cotangent, back-propagates each column on its own device, and sums the gradients in a fixed order. The result does not depend on which GPU ran which column.

Raises ValueError if both or neither of objective and field_cotangents are given, if counts or shapes do not match, if columns carry no autograd graph, or if no column depends on the parameters. TypeError for wrong dtypes or a non-scalar objective.

ensemble_scene_gradient_vjp#

ensemble_scene_gradient_vjp(basis, *, column_forward, weights, execution, objective=None, field_cotangents=None, reduction_order=None, reduction_device=None) -> AggregatedSceneGradient

Runs each array column's forward as an independent task over the execution device pool (round-robin placement, fail-fast), then calls aggregate_scene_gradient_vjp. The backward and the sum run on the calling thread after every device has synchronized. Use this instead of run_many, which rejects trainable simulations.

Raises TypeError if execution is not a MultiGPUExecution; RuntimeError if a column forward fails; otherwise as for aggregate_scene_gradient_vjp.

grad = mw.ensemble_scene_gradient_vjp(
    basis, column_forward=column_forward, weights=weights,
    execution=mw.MultiGPUExecution.ensemble(devices=("cuda:0", "cuda:1")),
    objective=objective, reduction_device="cuda:0",
)
print(grad.gradient, grad.column_devices)