Parallel Output

When running inference on large models with many output fields, writing results to disk can become a bottleneck. The Parallel Output wrapper offloads the I/O work to multiple writer processes, allowing the main inference loop to continue without waiting for disk writes to complete.

Each writer process handles a subset of the output fields and writes them independently to its own file. This distributes both the I/O bandwidth and the encoding/serialisation cost across multiple CPU cores.

How It Works

The parallel output wraps any other output type (e.g. grib, zarr, netcdf) and:

  1. Spawns num_writers writer processes (using fork).

  2. At each forecast step, splits the output fields into chunks according to the configured chunking strategy.

  3. Each chunk is dispatched to a writer process (round-robin when there are more chunks than writers) via a multiprocessing queue, and written using that writer’s own instance of the wrapped output.

  4. On shutdown, all writers are gracefully terminated with a signal via their queues.

Each writer appends a suffix _w<id> to the output file name to avoid conflicts (e.g. output_w0.grib, output_w1.grib).

Configuration

To enable parallel output, wrap your existing output configuration inside a parallel block:

output:
  parallel:
    num_writers: 4
    output:
      grib:
        path: /path/to/output.grib

This produces four output files:

  • /path/to/output_w0.grib

  • /path/to/output_w1.grib

  • /path/to/output_w2.grib

  • /path/to/output_w3.grib

This short syntax, where the inner output is implied, is also supported:

output:
  parallel:
    num_writers: 4
    grib:
      path: /path/to/output.grib

Parameters

ParallelOutput.__init__(context: Context, metadata: Metadata, *, output: Output | Any | None = None, num_writers: int = 1, chunk_strategy: Literal['by_worker', 'by_metadata', 'by_size'] | dict[Literal['by_worker', 'by_metadata', 'by_size'], dict[str, Any]] = 'by_worker', **kwargs: Any)

Initialise the ParallelOutput.

Parameters:

  • context (Context) – The inference context.

  • metadata (Metadata) – Metadata for the dataset.

  • output (Output | Any | None) – The inner output (or its config dict) that will be forked into writer processes.

  • num_writers (int) – Number of writer processes to spawn. Must be >= 1. Defaults to 1 (single output file, asynchronous writes).

  • chunk_strategy (CHUNK_STRATEGIES | dict[CHUNK_STRATEGIES, dict[str, Any]], default="by_worker") –

    The strategy for chunking the output among writer processes. Can be a string (one of “by_size”, “by_metadata”, “by_worker”) or a dictionary with a single key being the strategy name and the value being a dictionary of keyword arguments for that strategy. - by_worker (default), splits the fields into exactly num_writers contiguous

    chunks of roughly equal size, one per writer. No additional arguments are needed.

    • by_size, requires an additional argument fields_per_chunk specifying the number of fields per chunk.

    • by_metadata, requires an additional argument keys specifying a list of metadata keys to keep in a chunk.

  • **kwargs (Any) – Forwarded to the inner output.

Chunking Strategies

At each forecast step, the output fields are divided into chunks that are distributed to the writer processes. The chunk_strategy parameter controls how this split is performed. If there are more chunks than writers, chunks are assigned to writers round-robin (chunk i goes to writer i % num_writers).

by_worker (default)

Splits the fields into exactly num_writers contiguous chunks of roughly equal size, one per writer. This is the simplest strategy and a good default: each writer receives a comparable amount of work.

output:
  parallel:
    num_writers: 4
    chunk_strategy: by_worker
    output:
      grib:
        path: forecast.grib

by_size

Splits the fields into chunks of a fixed size, where “size” is the number of fields in the state (not the number of bytes), given by fields_per_chunk. The number of chunks depends on the number of output fields, and chunks are dispatched to writers round-robin. Use this when you want to control the amount of work per queue message (e.g. to bound memory usage), independently of num_writers.

output:
  parallel:
    num_writers: 4
    chunk_strategy:
      by_size:
        fields_per_chunk: 10
    output:
      grib:
        path: forecast.grib

by_metadata

Groups fields into chunks by their metadata, so that all fields sharing the same value for the given keys are written by the same writer. The keys are read from each variable’s typed metadata (e.g. levtype, param, level). This is useful when the output backend benefits from co-locating related fields in the same file.

output:
  parallel:
    num_writers: 4
    chunk_strategy:
      by_metadata:
        keys: [levtype]
    output:
      grib:
        path: forecast.grib

Note

by_metadata optimises for co-location, not for balanced load. The resulting chunks can be uneven in size: fields whose metadata does not define the requested keys are grouped together under a single None group, so if many fields lack that metadata they all end up in one large chunk on one writer. If balanced work across writers matters more than co-location, prefer by_worker (equal-sized chunks, one per writer) or by_size (fixed number of fields per chunk).

Limiting the number of groups with max_groups

An optional max_groups argument (specific to by_metadata; it is not available for by_worker or by_size) caps the number of chunks produced. Each metadata group becomes one chunk, and chunks are dispatched to writers round-robin (chunk i goes to writer i % num_writers), so a group is not tied one-to-one to a writer: one writer may receive several groups, and with max_groups set, several groups may be merged into a single chunk. If more distinct metadata groups exist than max_groups, groups are merged round-robin so that no more than max_groups chunks are created. This is useful to avoid overwhelming the writers with many small chunks:

output:
  parallel:
    num_writers: 4
    chunk_strategy:
      by_metadata:
        keys: [param]
        max_groups: 8
    output:
      grib:
        path: forecast.grib

Examples

GRIB output with 2 writers

output:
  parallel:
    num_writers: 2
    output:
      grib:
        path: forecast.grib

Zarr output with 4 writers

output:
  parallel:
    num_writers: 4
    output:
      zarr:
        store: forecast.zarr

NetCDF output with 8 writers

output:
  parallel:
    num_writers: 8
    output:
      netcdf:
        path: forecast.nc

Combining with parallel inference

Parallel output can be used together with the parallel runner for distributed model inference. In this case, the CPU process associated with GPU 0 will handle the output writing. Only CPU process 0 will spawn writer processes.

runner: parallel
lead_time: 240h
checkpoint: /path/to/checkpoint.ckpt

input:
  grib: /path/to/input.grib

output:
  parallel:
    num_writers: 4
    output:
      grib:
        path: /path/to/output.grib

Choosing num_writers

For small resolutions or infrequent output steps, parallel writing might be unnecessary.

For larger output sizes, one can set num_writers: 1. This will produce a single output file, but the writing will happen asynchronously in a separate process, allowing the main inference loop to continue without waiting for disk I/O.

If you suspect that writing to disk is a bottleneck, you can experiment with increasing num_writers. 2 or 4 writers is often a good starting point.

Troubleshooting

Writer process crashes

If a writer process crashes (e.g. due to a disk-full error or a bug in the output backend), the main process will detect that the writer is dead and abort inference with a RuntimeError.