Scaling Batch Inference with Dask and Ray

Run geospatial model inference over terabyte rasters with Dask and Ray: chunk alignment, block-wise prediction, model broadcast, memory budgeting and writing results without a shuffle.

A trained model that takes 40 ms per tile is fast. Applied to a national mosaic of 180,000 tiles it is two hours of pure compute, plus the reading, the feature derivation, and the writing — and none of it fits in memory. Distributed batch inference is the engineering that turns a working model into a production map, and almost all of its difficulty is in memory, chunk geometry and where the model object lives, not in the parallelism itself.

This topic is part of Geospatial MLOps and Model Deployment. It is the compute layer underneath orchestrating geospatial ML pipelines: the orchestrator decides when a run happens and what to do when it fails, while Dask or Ray decides how the tiles inside that run are spread across cores. It usually consumes the portable artifacts produced by ONNX export for geospatial model inference, because a serialisable, thread-safe session is far easier to distribute than a full training framework.

Block-wise inference: read a chunk, predict, write the same chunk A large source raster is divided into chunks whose boundaries line up with the file's internal tiling. Three workers each hold one copy of the model in memory and process chunks independently. Each worker writes its prediction directly into the matching block of the output raster, so the complete prediction array never exists in memory. Chunk in, chunk out — the full array is never assembled source raster chunks = k × internal GeoTIFF block size worker 1 model loaded once worker 2 model loaded once worker 3 model loaded once predict(block) same shape in and out no cross-chunk state output raster written block by block, tiled + compressed Peak memory is set by one chunk per thread, not by the size of the scene.

Problem Framing

Distributed raster inference fails in four characteristic ways, and none of them is “not enough cores”.

Chunk boundaries fight the file layout. A GeoTIFF stores pixels in internal blocks, typically 256×256 or 512×512, each independently compressed. If a Dask chunk is 300×300, every chunk read decompresses parts of four blocks and the same block is decompressed by multiple chunks. Aligning chunks to an exact multiple of the internal block size routinely doubles throughput with no other change.

Peak memory is several times chunk size. A uint16 chunk becomes a float32 feature array, which spawns intermediates during band math, all of which coexist with the model. Multiply by threads per worker. Workers that die with “memory limit exceeded” at 25% of the nominal budget are almost always over-threaded rather than over-chunked.

The model gets serialised into every task. Capturing a model in a closure means Dask pickles it into each of the 180,000 task payloads. The graph becomes gigabytes, the scheduler stalls, and CPU time disappears into deserialisation. The model must be loaded by the worker, not sent to it.

Neighbourhood features silently break at chunk edges. Anything with a kernel — focal statistics, terrain derivatives, a convolutional model — needs a halo. Without it, every chunk boundary gets a band of wrong values, which is invisible at national zoom and obvious the moment somebody inspects a tile edge. The halo mechanics are the same as in DEM and terrain derivative features.

Prerequisites & Environment Setup

# Pinned requirements for distributed raster inference
dask==2024.6.2
distributed==2024.6.2
ray==2.31.0
rioxarray==0.15.7
xarray==2024.6.0
rasterio==1.3.10
onnxruntime==1.18.0
numpy==1.26.4

Install with:

pip install "dask==2024.6.2" "distributed==2024.6.2" "rioxarray==0.15.7" \
            "xarray==2024.6.0" "rasterio==1.3.10" "numpy==1.26.4" \
            "onnxruntime==1.18.0"
pip install "ray==2.31.0"          # only if you use the Ray path

GDAL/PROJ system dependencies (Ubuntu/Debian):

sudo apt-get install -y gdal-bin libgdal-dev libproj-dev

Every worker needs an identical geospatial stack. A PROJ version mismatch between the client and a worker produces silently different reprojections on different machines — the strongest practical argument for shipping one container image to all of them, as covered in containerizing geospatial inference pipelines.

Step-by-Step Implementation

Step 1 — Read the file’s own block size and chunk to a multiple of it

import rasterio
import rioxarray


def native_chunks(path: str, target_mb: int = 128, dtype_bytes: int = 4,
                  n_bands: int = 6) -> dict:
    """Chunk sizes that are exact multiples of the GeoTIFF's internal blocks."""
    with rasterio.open(path) as src:
        bx = src.block_shapes[0][1] or 256
        by = src.block_shapes[0][0] or 256

    px_budget = (target_mb * 1024 ** 2) / (dtype_bytes * n_bands)
    side = int(px_budget ** 0.5)
    cx = max(bx, (side // bx) * bx)
    cy = max(by, (side // by) * by)
    return {"x": cx, "y": cy}


chunks = native_chunks("features.tif")
da = rioxarray.open_rasterio("features.tif", chunks=chunks, lock=False)
assert da.chunks is not None, "array is not chunked — the whole scene would load eagerly"
print("chunk grid:", [len(c) for c in da.chunks])

Step 2 — Load the model once per worker

import threading
import numpy as np
import onnxruntime as ort

_SESSION = None
_LOCK = threading.Lock()


def get_session(model_uri: str) -> ort.InferenceSession:
    """Module-level cache: one session per worker process, created on first use."""
    global _SESSION
    if _SESSION is None:
        with _LOCK:
            if _SESSION is None:
                opts = ort.SessionOptions()
                opts.intra_op_num_threads = 1     # let Dask own the parallelism
                _SESSION = ort.InferenceSession(model_uri, sess_options=opts,
                                                providers=["CPUExecutionProvider"])
    return _SESSION

The double-checked lock matters: threaded Dask workers will call this concurrently, and building two sessions doubles peak memory at exactly the wrong moment.

What actually occupies a worker's memory during inference A stacked bar shows the components of one task's resident memory: the raw uint16 chunk, its float32 copy, band-math intermediates and the model session. A second bar shows the same stack multiplied by four threads, exceeding the worker memory limit line, illustrating why thread count is the first thing to reduce. Peak memory = one chunk’s stack × threads per worker 1 thread chunk 0.4 GB float32 0.8 GB temps 0.3 GB model 0.2 GB 1.7 GB — fits 4 threads 4 × the same stack over the limit memory_limit 6 GB 6.8 GB — worker restarts More single-threaded processes beat fewer multi-threaded ones for GDAL-heavy work.

Step 3 — Map the prediction over blocks

import dask.array as darr

NODATA = 255


def predict_block(block: np.ndarray, model_uri: str, mean, std) -> np.ndarray:
    """Predict one chunk. Shape (bands, y, x) in, (1, y, x) out."""
    sess = get_session(model_uri)
    bands, ny, nx = block.shape

    flat = block.reshape(bands, -1).T.astype("float32")
    valid = np.isfinite(flat).all(axis=1)
    out = np.full(flat.shape[0], NODATA, dtype="uint8")

    if valid.any():
        x = (flat[valid] - mean) / std          # stats from the TRAINING artifact
        logits = sess.run(None, {sess.get_inputs()[0].name: x})[0]
        out[valid] = logits.argmax(axis=1).astype("uint8")

    return out.reshape(1, ny, nx)


preds = darr.map_blocks(
    predict_block,
    da.data,
    model_uri="s3://models/landcover/v7/model.onnx",
    mean=mean, std=std,
    dtype="uint8",
    chunks=((1,), da.data.chunks[1], da.data.chunks[2]),
)
assert preds.dtype == np.uint8

map_blocks keeps the graph embarrassingly parallel: no chunk depends on any other, so there is no shuffle and no scheduler bottleneck.

Step 4 — Write block by block

from dask.distributed import Client
import rasterio
from rasterio.windows import Window


def write_blocks(preds, template_path: str, out_path: str) -> None:
    """Stream chunks to a tiled GeoTIFF; the full array is never materialised."""
    with rasterio.open(template_path) as src:
        profile = src.profile.copy()
    profile.update(count=1, dtype="uint8", nodata=NODATA, compress="deflate",
                   tiled=True, blockxsize=512, blockysize=512, predictor=2)

    ychunks, xchunks = preds.chunks[1], preds.chunks[2]
    with rasterio.open(out_path, "w", **profile) as dst:
        y0 = 0
        for iy, ylen in enumerate(ychunks):
            x0 = 0
            for ix, xlen in enumerate(xchunks):
                tile = preds.blocks[0, iy, ix].compute()
                dst.write(tile[0], 1, window=Window(x0, y0, xlen, ylen))
                x0 += xlen
            y0 += ylen


client = Client(n_workers=8, threads_per_worker=1, memory_limit="6GB")
write_blocks(preds, "features.tif", "landcover_v7.tif")

threads_per_worker=1 with more processes is the right default for GDAL-heavy work: GDAL’s internal caches are per-process, and single-threaded workers make memory accounting predictable.

Step 5 — The Ray variant, for heavyweight per-worker state

When the model is a GPU network, an actor that owns the device and stays alive across thousands of tiles beats re-creating it per task.

import ray


@ray.remote(num_gpus=0.25)
class TilePredictor:
    """One long-lived model per GPU slice; tiles are streamed through it."""

    def __init__(self, model_uri: str, mean, std):
        import onnxruntime as ort
        self.sess = ort.InferenceSession(
            model_uri, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
        self.mean, self.std = mean, std

    def predict(self, tile_key: str) -> str:
        block = read_tile(tile_key)
        out = predict_block(block, self.sess, self.mean, self.std)
        return write_tile(out, tile_key.replace("/features/", "/predictions/"))


ray.init(address="auto")
actors = [TilePredictor.remote(MODEL_URI, mean, std) for _ in range(8)]
futures = [actors[i % len(actors)].predict.remote(k) for i, k in enumerate(tile_keys)]
written = ray.get(futures)
assert len(written) == len(tile_keys)

Verification & Testing

The one property distributed inference must preserve is that it changes nothing.

import numpy as np


def test_distributed_matches_single_threaded(small_scene, model_uri, mean, std):
    """A small scene must produce byte-identical predictions either way."""
    arr = rioxarray.open_rasterio(small_scene).values
    reference = predict_block(arr, model_uri, mean, std)

    lazy = rioxarray.open_rasterio(small_scene, chunks={"x": 256, "y": 256}).data
    distributed = darr.map_blocks(
        predict_block, lazy, model_uri=model_uri, mean=mean, std=std,
        dtype="uint8", chunks=((1,), lazy.chunks[1], lazy.chunks[2]),
    ).compute()

    assert np.array_equal(reference, distributed), "chunking changed the predictions"


def test_chunks_align_to_internal_blocks():
    ch = native_chunks("features.tif")
    with rasterio.open("features.tif") as src:
        by, bx = src.block_shapes[0]
    assert ch["x"] % bx == 0 and ch["y"] % by == 0

Beyond correctness, watch the dashboard. In a healthy run the task stream is a dense band of one colour; large white gaps mean workers are waiting on I/O, and a growing red memory bar means chunks are too large for the thread count. Both are diagnosable in a five-minute run on 1% of the tiles, which is always worth doing before launching the full job.

Chunks that straddle internal blocks decompress the same bytes repeatedly Two panels show the same file whose internal compression blocks are drawn as a fine grid. In the first panel the chunk boundaries fall between block boundaries, so each chunk overlaps four blocks and every block is decompressed more than once. In the second panel the chunk boundaries coincide with block boundaries, so each block is read and decompressed exactly once. chunk 300 × 300 on 256 × 256 blocks chunk 512 × 512 on 256 × 256 blocks Each chunk overlaps 9 blocks; every block is decompressed up to 4 times Each chunk is exactly 4 blocks; every block is decompressed exactly once

Troubleshooting & Common Errors

distributed.worker - WARNING - Memory use is high then workers restart — too many threads per worker, or chunks too large for the float32 expansion. Drop to one thread per worker and halve the chunk area; the wall-clock cost is far lower than the cost of restarts.

The graph takes minutes to build and the dashboard is idle — a large object is captured in the mapped function. Check len(pickle.dumps(fn)); anything above a few kilobytes means the model or a big array is riding along in the closure.

CPLE_OpenFailed on some workers only — GDAL configuration (credentials, AWS_REGION, GDAL_DISABLE_READDIR_ON_OPEN) is set on the client but not on the workers. Set it in the worker environment or with a WorkerPlugin, never with a client-side rasterio.Env that does not travel.

Predictions differ between a local run and the cluster — normalisation statistics are being computed per chunk instead of loaded from the training artifact. Load them from the file written during training, exactly as in scaling features consistently between training and inference.

Bands of wrong values along every chunk boundary — a neighbourhood operation without a halo. Use map_overlap with a depth of at least half the kernel width, and trim the halo before writing.

Ray actors sit idle while tasks queue — the actor pool is smaller than the available parallelism, or num_gpus fractions do not divide evenly into the devices. Size the pool from device count and per-actor memory, then submit with ray.util.ActorPool so work is balanced rather than round-robined.

Performance Optimisation

Fix the I/O before touching the compute. For cloud-optimized GeoTIFFs, set GDAL_DISABLE_READDIR_ON_OPEN=EMPTY_DIR, GDAL_HTTP_MULTIPLEX=YES and a GDAL_CACHEMAX of a few hundred megabytes per worker. On a remote-read pipeline this frequently matters more than doubling the worker count.

Convert once to a format built for chunked reads. If the same scene is read repeatedly, convert it to Zarr or to a tiled, overview-carrying COG first. Reading a striped GeoTIFF in blocks is pathological, and the conversion pays for itself on the second pass.

Batch rows into the model, not pixels. ONNX Runtime and most tree libraries amortise call overhead over batch size. Reshape a chunk into an (n_pixels, n_features) matrix and make one call per chunk — the difference against per-pixel calls is one to two orders of magnitude.

Cap the model’s internal threading. intra_op_num_threads=1 for ONNX Runtime, n_jobs=1 for scikit-learn, OMP_NUM_THREADS=1 in the worker environment. Otherwise every worker spawns a full thread pool and the machine thrashes.

Profile with a 1% sample. Run the pipeline over a hundred tiles, read the task stream, and only then launch the full job. The failure modes above all show up in that sample, and a two-minute diagnostic beats a six-hour job that dies at 80%.

FAQ

Dask or Ray for raster inference?

Dask fits array-shaped work — the raster is already a chunked array, and map_blocks preserves lazy evaluation and geometry. Ray fits task-shaped work with heavy per-worker state, such as a GPU model in a long-lived actor. Using Dask for feature computation and Ray for the model call is a common and reasonable split.

Why do workers run out of memory even though each chunk is small?

Because peak memory is the chunk, plus its float32 copy, plus intermediates, plus the model — multiplied by threads per worker. A 256 MB chunk on a four-threaded worker can peak past 4 GB. Reduce threads before reducing chunk size.

How do I keep the model from being serialised for every task?

Never capture it in the closure. Load it lazily from a module-level cache inside the worker (Step 2), or hold it in a Ray actor. Capturing it pickles the model into every task payload and stalls the scheduler.

Does distributed inference change the predictions?

It must not. Any difference means state is leaking across chunk boundaries — normalisation computed per chunk, or a neighbourhood feature without a halo. The equality test in the verification section should be part of your test suite.

Part of: Geospatial MLOps and Model Deployment