Running Block-Wise Raster Inference with Dask Array

Apply a model across a terabyte raster with dask.array.map_blocks: chunk alignment, lazy model loading, halo handling with map_overlap, and streaming results to a tiled GeoTIFF.

Open the raster lazily with chunks that are an exact multiple of its internal block size, map a prediction function over the blocks, and write the result block by block. The prediction function must load the model itself — from a module-level cache keyed by artifact path — rather than closing over a model object, because everything in the closure is pickled into every one of the thousands of tasks.

This page is the Dask path. For the wider comparison with Ray and the memory budgeting, see Scaling Batch Inference with Dask and Ray.

What travels in the graph decides whether the job starts at all Two task graphs are compared. In the first the model object is captured in the closure, so every one of six thousand tasks carries a copy of the model and the graph is gigabytes in size. In the second only a short artifact path travels, so each task payload is tiny and the model is loaded once per worker. model in the closure path in the closure 240 MB240 MB240 MB 240 MB240 MB240 MB × 6 000 tasks graph never finishes serialising dashboard idle, scheduler at 100% CPU 180 B180 B180 B 180 B180 B180 B model loaded once per worker process graph builds in milliseconds workers saturate immediately

Why This Fails in Geospatial ML Pipelines

Closure capture is the failure that looks like a cluster problem. map_blocks(predict, arr, model=model) embeds the serialised model in every task, so a 240 MB gradient-boosting ensemble across 6,000 tasks produces a graph larger than the dataset. The scheduler saturates one core serialising it, the dashboard shows no work, and the natural response — adding workers — makes nothing better. The one-line diagnostic is len(pickle.dumps(fn)).

The second failure is silent chunk misalignment. Dask will happily use whatever chunk size you give it, and a GeoTIFF read of a chunk that straddles four internal blocks decompresses all four. With overlapping chunks, the same compressed block is decompressed several times across different tasks. Nothing errors; the job is simply two to four times slower than it should be, and the task stream shows long red I/O bands.

Third, receptive fields. A per-pixel model is safe under map_blocks, but the moment features include a focal statistic or the model is convolutional, each block needs neighbours it does not have. map_blocks gives it zero-padding instead, producing a faint grid of wrong values along every chunk boundary — the same artefact discussed in computing focal window statistics on rasters.

Core Principles

  • Never capture the model. Pass a path; load lazily inside the worker.
  • Align chunks to the file’s internal blocks.
  • map_overlap whenever there is a receptive field, with depth at least half of it.
  • Single-threaded workers for GDAL-heavy work, more processes instead.
  • Set GDAL configuration on the workers, not on the client.
  • Write while computing, block by block, so the output never materialises in memory.

Production-Ready Code

from __future__ import annotations

import logging
import threading

import dask.array as darr
import numpy as np
import onnxruntime as ort
import rasterio
import rioxarray
from rasterio.windows import Window

logger = logging.getLogger(__name__)

NODATA = 255
_SESSIONS: dict[str, ort.InferenceSession] = {}
_LOCK = threading.Lock()


def get_session(model_uri: str) -> ort.InferenceSession:
    """One session per (process, artifact). Built on first use inside the worker."""
    sess = _SESSIONS.get(model_uri)
    if sess is None:
        with _LOCK:
            sess = _SESSIONS.get(model_uri)
            if sess is None:
                opts = ort.SessionOptions()
                opts.intra_op_num_threads = 1        # Dask owns the parallelism
                sess = ort.InferenceSession(model_uri, sess_options=opts,
                                            providers=["CPUExecutionProvider"])
                _SESSIONS[model_uri] = sess
                logger.info("loaded model %s in this worker", model_uri)
    return sess


def native_chunks(path: str, target_mb: int = 128, n_bands: int = 6) -> dict:
    """Chunks that are exact multiples of the GeoTIFF's internal block size."""
    with rasterio.open(path) as src:
        by, bx = src.block_shapes[0]
    side = int(((target_mb * 1024 ** 2) / (4 * n_bands)) ** 0.5)
    return {"y": max(by, (side // by) * by), "x": max(bx, (side // bx) * bx)}


def predict_block(block: np.ndarray, model_uri: str, mean: np.ndarray,
                  std: np.ndarray) -> np.ndarray:
    """(bands, y, x) -> (1, y, x) uint8 class map. Pure: no captured state."""
    sess = get_session(model_uri)
    n_bands, ny, nx = block.shape
    flat = block.reshape(n_bands, -1).T.astype("float32")

    out = np.full(flat.shape[0], NODATA, dtype="uint8")
    valid = np.isfinite(flat).all(axis=1)
    if valid.any():
        x = (flat[valid] - mean) / std
        logits = sess.run(None, {sess.get_inputs()[0].name: x})[0]
        out[valid] = np.argmax(logits, axis=1).astype("uint8")
    return out.reshape(1, ny, nx)


def build_prediction_array(raster_path: str, model_uri: str, mean, std,
                           depth: int = 0) -> darr.Array:
    """Lazy prediction array. Uses map_overlap when the model has a receptive field."""
    chunks = native_chunks(raster_path)
    da = rioxarray.open_rasterio(raster_path, chunks=chunks, lock=False).data

    kwargs = dict(model_uri=model_uri, mean=mean, std=std)
    out_chunks = ((1,), da.chunks[1], da.chunks[2])

    if depth == 0:
        return darr.map_blocks(predict_block, da, dtype="uint8",
                               chunks=out_chunks, **kwargs)

    return darr.map_overlap(
        predict_block, da,
        depth={0: 0, 1: depth, 2: depth}, boundary="nearest",
        trim=True, dtype="uint8", chunks=out_chunks, **kwargs,
    )


def write_streaming(preds: darr.Array, template: str, out_path: str) -> None:
    """Compute and write one block at a time; the full array never exists."""
    with rasterio.open(template) as src:
        profile = src.profile.copy()
    profile.update(count=1, dtype="uint8", nodata=NODATA, compress="deflate",
                   predictor=2, tiled=True, blockxsize=512, blockysize=512)

    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
            logger.info("wrote row %d/%d of blocks", iy + 1, len(ychunks))

Step-by-Step Walkthrough

Step 1 — Start a cluster shaped for GDAL

from dask.distributed import Client, WorkerPlugin

class GdalEnv(WorkerPlugin):
    def setup(self, worker):
        import os
        os.environ.update({
            "GDAL_DISABLE_READDIR_ON_OPEN": "EMPTY_DIR",
            "GDAL_HTTP_MULTIPLEX": "YES",
            "GDAL_CACHEMAX": "256",
            "OMP_NUM_THREADS": "1",
        })

client = Client(n_workers=8, threads_per_worker=1, memory_limit="6GB")
client.register_plugin(GdalEnv())
print(client.dashboard_link)

Setting the environment through a plugin is what makes it apply on the workers. A client-side rasterio.Env does not travel.

Step 2 — Check the graph is light before computing

import pickle
from functools import partial

stats = np.load("artifacts/norm_stats.npz")
fn = partial(predict_block, model_uri="s3://models/landcover/v7/model.onnx",
             mean=stats["mean"], std=stats["std"])
print(f"closure size: {len(pickle.dumps(fn)) / 1024:.1f} KB")
assert len(pickle.dumps(fn)) < 100_000, "something heavy is captured in the closure"

Step 3 — Build and inspect the lazy array

preds = build_prediction_array("features.tif",
                               "s3://models/landcover/v7/model.onnx",
                               stats["mean"], stats["std"], depth=0)
print(preds)
print("blocks:", [len(c) for c in preds.chunks])

Nothing has been read yet. If the chunk grid has thousands of tiny blocks, the chunk size is too small and per-task overhead will dominate.

Step 4 — Sample before committing

sample = preds[:, :2048, :2048].compute()
vals, counts = np.unique(sample, return_counts=True)
print(dict(zip(vals.tolist(), counts.tolist())))

A 2048-pixel corner runs in seconds and exposes every configuration error — missing credentials, wrong band count, an all-nodata result — before a six-hour job.

Step 5 — Write

write_streaming(preds, "features.tif", "landcover_v7.tif")
map_overlap expands, computes, then trims A block is drawn with a surrounding halo taken from its neighbouring blocks. The prediction function runs over the expanded array so every pixel has a full neighbourhood. The halo is then trimmed from the result so the written block matches the original chunk exactly. depth must be at least half the model’s receptive field 1. chunk grid the block to compute expand 2. with halo depth pixels from each neighbour trim 3. result exactly the original chunk With depth=0 and a convolutional model, every chunk edge is wrong.

Verification

import numpy as np
import pickle
import pytest


def test_closure_is_light():
    from functools import partial
    fn = partial(predict_block, model_uri="s3://models/v7/model.onnx",
                 mean=np.zeros(6, "float32"), std=np.ones(6, "float32"))
    assert len(pickle.dumps(fn)) < 100_000


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["y"] % by == 0 and ch["x"] % bx == 0


def test_distributed_equals_single_threaded(small_scene, model_uri, mean, std):
    ref = predict_block(rioxarray.open_rasterio(small_scene).values,
                        model_uri, mean, std)
    lazy = build_prediction_array(small_scene, model_uri, mean, std, depth=0).compute()
    assert np.array_equal(ref, lazy), "chunking changed the predictions"


def test_overlap_removes_edge_artifacts(focal_model_uri, mean, std):
    """With a receptive field, depth=0 and depth=8 must differ at chunk edges."""
    a = build_prediction_array("features.tif", focal_model_uri, mean, std, depth=0)
    b = build_prediction_array("features.tif", focal_model_uri, mean, std, depth=8)
    edge = slice(508, 516)
    assert not np.array_equal(a[0, edge, :].compute(), b[0, edge, :].compute())


def test_output_grid_matches_input(tmp_path):
    out = tmp_path / "p.tif"
    write_streaming(preds, "features.tif", str(out))
    with rasterio.open("features.tif") as a, rasterio.open(str(out)) as b:
        assert a.transform.almost_equals(b.transform, precision=1e-6)

The equality test is the one that matters most: distributing the computation must not change a single pixel. Any difference means state is leaking across block boundaries.

Read the task stream before adding workers Two task stream traces are compared. The healthy trace is a dense continuous band of compute tasks across all workers. The unhealthy trace has large gaps between short compute bursts, indicating workers are waiting on input and output rather than being short of CPU. Dask task stream, four workers healthy dense compute — adding workers helps I/O bound gaps mean chunk misalignment or untuned GDAL — more workers will not help

FAQ

Why is my Dask graph gigabytes in size?

Because the model is captured in the closure and pickled into every task. Pass an artifact path and load lazily inside the worker; len(pickle.dumps(fn)) tells you in one line whether the closure is heavy.

map_blocks or map_overlap?

map_blocks for a per-pixel model. map_overlap whenever the model or its features have a spatial receptive field, with depth at least half of it and trim=True.

Why is the cluster idle while the job is slow?

It is I/O, not compute. Align chunks to the internal block size and set the GDAL environment on the workers via a WorkerPlugin — a client-side rasterio.Env never reaches them.

Part of: Scaling Batch Inference with Dask and Ray Part of: Geospatial MLOps and Model Deployment