Article

Choosing Chunk Size for Multiprocessing Raster Warps

Pick a chunk size that is an integer multiple of the source raster’s native block edge, typically 512 or 1024 pixels square, then cap the worker count so chunk_bytes * workers stays under your RAM budget. Block alignment removes read amplification, and the memory cap prevents out-of-memory kills. This page is part of the Multiprocessing Geospatial Tasks in Python guide inside the broader Spatial Batch Processing & Async Workflows reference.

The core trade-off is simple: chunks that are too small drown the pool in scheduling overhead and repeated compression setup, while chunks that are too large exhaust memory and leave workers idle at the tail of the job. The helper below reads block_shapes, snaps a requested chunk edge to a block multiple, estimates memory, and derives a safe (chunk_size, workers) pair before dispatching windows.

Prerequisites

  • Python 3.10 or later
  • pip install rasterio (bundles GDAL 3.4+ wheels; no separate GDAL install needed)
  • psutil for reading available RAM: pip install psutil
  • A tiled source raster. Untiled (striped) GeoTIFFs have a block shape of one full row, which defeats chunk alignment. Convert first with gdal_translate -co TILED=YES.

For how the pool itself should be built and isolated, read GDAL Batch Operations with multiprocessing.Pool. For the memory-budget reasoning behind the worker cap, see Memory Management for Large GIS Datasets.

The Sizing Decision

Chunk size is not one number to tune in isolation; it is the vertex where block geometry, per-worker memory, and pool overhead meet. The diagram traces the decision path the helper follows.

Chunk size decision flow A vertical flow: read block_shapes, snap the requested chunk to a block multiple, estimate per-chunk bytes, compare the total against the RAM budget, then either dispatch or shrink and retry. Read source block_shapes e.g. (512, 512) Snap chunk to block multiple 1024 = 2 blocks Estimate per-chunk bytes w x h x bands x dtype bytes x workers under RAM budget? yes Dispatch windows to process pool no Shrink chunk or fewer workers Larger chunks reduce overhead; smaller chunks reduce memory and improve tail balancing.

Trade-offs at a Glance

Chunk size Scheduling overhead Peak worker memory Load balancing at tail COG recompression cost
Too small (< 1 block) High — task setup and IPC dominate Low Good High — many small compressed blocks re-encoded
Aligned (1–2 blocks) Balanced Predictable Good Minimal — reads map onto stored tiles
Large (4–8 blocks) Low High — risk of OOM Poor — stragglers idle the pool Low per chunk, but fewer parallel units
Too large (whole scene) Minimal Exceeds RAM Worst — no parallelism Low

The sweet spot is one to two native blocks per chunk edge. That keeps each read aligned to stored tiles, gives the pool enough independent units to balance the tail, and holds per-worker memory to a predictable figure you can budget against.

Complete Working Implementation

The script reads block_shapes, snaps a target chunk edge to a block multiple, estimates memory, derives a safe worker count from available RAM, and dispatches aligned windows to a ProcessPoolExecutor. Copy it, adjust the CRS and paths, and run directly:

#!/usr/bin/env python3
"""
Choose a block-aligned chunk size for a multiprocessing raster warp and
dispatch the tile windows to a ProcessPoolExecutor.

Usage: python chunked_warp.py source.tif out.tif --crs EPSG:3857 --chunk 1024
"""
import argparse
import sys
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed

import numpy as np
import psutil
import rasterio
from rasterio.windows import Window
from rasterio.warp import calculate_default_transform, reproject, Resampling

# Reserve headroom so GDAL warp scratch buffers do not push the box into swap.
MEMORY_SAFETY_FACTOR = 3.0     # per-chunk arrays + warp working buffers
RAM_BUDGET_FRACTION = 0.7      # never claim more than 70% of free RAM


def align_chunk_to_block(requested: int, block_edge: int) -> int:
    """Round a requested chunk edge to the nearest positive multiple of the
    native block edge. A chunk that is not a block multiple forces GDAL to read
    partial blocks in full, amplifying I/O."""
    if block_edge <= 0:
        return requested
    multiple = max(1, round(requested / block_edge))
    return multiple * block_edge


def estimate_chunk_bytes(edge: int, bands: int, dtype: str) -> int:
    """Resident bytes for one square chunk: width * height * bands * dtype_size."""
    dtype_size = np.dtype(dtype).itemsize
    return edge * edge * bands * dtype_size


def plan_chunking(src_path: Path, requested_chunk: int) -> tuple[int, int]:
    """Return a safe (chunk_edge, worker_count) pair for this raster."""
    with rasterio.open(src_path) as src:
        # block_shapes is one (rows, cols) tuple per band; band 1 is representative.
        block_rows, block_cols = src.block_shapes[0]
        block_edge = min(block_rows, block_cols)
        chunk_edge = align_chunk_to_block(requested_chunk, block_edge)

        per_chunk = estimate_chunk_bytes(chunk_edge, src.count, src.dtypes[0])

    budget = int(psutil.virtual_memory().available * RAM_BUDGET_FRACTION)
    footprint = int(per_chunk * MEMORY_SAFETY_FACTOR)
    safe_workers = max(1, budget // max(footprint, 1))
    workers = min(safe_workers, psutil.cpu_count(logical=False) or 1)

    print(f"native block edge : {block_edge}px")
    print(f"chunk edge        : {chunk_edge}px "
          f"({chunk_edge // block_edge}x block)")
    print(f"per-chunk estimate : {per_chunk / 1e6:.1f} MB "
          f"(x{MEMORY_SAFETY_FACTOR} = {footprint / 1e6:.1f} MB budgeted)")
    print(f"RAM budget        : {budget / 1e9:.2f} GB free-derived")
    print(f"safe workers      : {workers}")
    return chunk_edge, int(workers)


def iter_windows(width: int, height: int, edge: int):
    """Yield block-aligned windows, clipping the last row/col to raster bounds."""
    for row_off in range(0, height, edge):
        for col_off in range(0, width, edge):
            w = min(edge, width - col_off)     # clip remainder column
            h = min(edge, height - row_off)    # clip remainder row
            yield Window(col_off, row_off, w, h)


def warp_window(src_path: str, dst_path: str, dst_crs: str,
                window: Window) -> tuple[int, int]:
    """Reproject one window and write it into the output at the same offset."""
    with rasterio.open(src_path) as src:
        transform = src.window_transform(window)
        data = src.read(window=window)
        dst_transform, dst_w, dst_h = calculate_default_transform(
            src.crs, dst_crs, window.width, window.height,
            *rasterio.windows.bounds(window, src.transform),
        )
        out = np.empty((src.count, dst_h, dst_w), dtype=src.dtypes[0])
        reproject(
            source=data,
            destination=out,
            src_transform=transform,
            src_crs=src.crs,
            dst_transform=dst_transform,
            dst_crs=dst_crs,
            resampling=Resampling.bilinear,
            num_threads=1,                     # match one-task-per-process model
        )
    with rasterio.open(dst_path, "r+") as dst:
        dst.write(out, window=Window(window.col_off, window.row_off, dst_w, dst_h))
    return (window.row_off, window.col_off)


def build_output(src_path: Path, dst_path: Path, dst_crs: str, edge: int) -> None:
    """Create the destination raster with a tiled, block-aligned layout."""
    with rasterio.open(src_path) as src:
        transform, width, height = calculate_default_transform(
            src.crs, dst_crs, src.width, src.height, *src.bounds)
        profile = src.profile.copy()
    profile.update(
        crs=dst_crs, transform=transform, width=width, height=height,
        tiled=True, blockxsize=edge, blockysize=edge,
        compress="deflate", BIGTIFF="YES",
    )
    with rasterio.open(dst_path, "w", **profile):
        pass  # allocate the file; workers fill it window by window


def main() -> None:
    parser = argparse.ArgumentParser(description="Block-aligned chunked raster warp")
    parser.add_argument("src", type=Path)
    parser.add_argument("dst", type=Path)
    parser.add_argument("--crs", default="EPSG:3857", help="Target CRS")
    parser.add_argument("--chunk", type=int, default=1024,
                        help="Requested chunk edge in pixels (snapped to block)")
    args = parser.parse_args()

    if not args.src.exists():
        print(f"source not found: {args.src}", file=sys.stderr)
        sys.exit(2)

    chunk_edge, workers = plan_chunking(args.src, args.chunk)
    build_output(args.src, args.dst, args.crs, chunk_edge)

    with rasterio.open(args.src) as src:
        windows = list(iter_windows(src.width, src.height, chunk_edge))

    failed = 0
    with ProcessPoolExecutor(max_workers=workers) as pool:
        futures = {
            pool.submit(warp_window, str(args.src), str(args.dst),
                        args.crs, w): w
            for w in windows
        }
        for fut in as_completed(futures):
            try:
                fut.result()
            except Exception as exc:                # noqa: BLE001
                failed += 1
                print(f"window {futures[fut]} failed: {exc}", file=sys.stderr)

    print(f"completed {len(windows) - failed}/{len(windows)} windows")
    sys.exit(0 if failed == 0 else 12)             # 12 = partial batch failure


if __name__ == "__main__":
    main()

Step Annotations

  1. align_chunk_to_block() — This is the load-bearing function. It rounds the requested edge to the nearest whole multiple of the native block edge (round(requested / block_edge)), never below one block. A requested 1000 against a 512 block becomes 1024, not 1000, so every read lands on a stored tile boundary.

  2. src.block_shapes[0] — Returns the (rows, cols) internal tiling for band 1. For a properly tiled GeoTIFF this is something like (512, 512); for a striped file it is (1, width), which is the signal to re-tile the source before chunking.

  3. estimate_chunk_bytes() — The footprint is edge * edge * bands * dtype_size. A 1024 chunk of a 4-band uint16 raster is 1024 * 1024 * 4 * 2 = 8 MB of pixel data per chunk, before warp buffers.

  4. MEMORY_SAFETY_FACTOR = 3.0 — The reproject call holds the source array, the destination array, and internal GDAL scratch simultaneously. Budgeting three times the raw pixel bytes keeps the box off swap. See Memory Management for Large GIS Datasets for how to tune this factor per dtype.

  5. safe_workers = budget // footprint — The worker count is derived, not guessed. It is the RAM budget divided by the padded per-chunk footprint, then capped at the physical core count so CPU-bound warps do not oversubscribe.

  6. iter_windows() clipping — min(edge, width - col_off) shrinks the final column and row to the true remainder so no window reads past the raster edge. This is the fix for the remainder gotcha below.

  7. num_threads=1 — Each window is warped by exactly one thread because parallelism already lives at the process level. Letting GDAL spin its own thread pool per worker reproduces the oversubscription problem covered in GDAL Batch Operations with multiprocessing.Pool.

Named Gotcha: Non-Block-Aligned and Remainder Chunks

The most common performance and correctness failure comes from two related mistakes. First, choosing a chunk edge that is not a multiple of the source block edge. A tiled GeoTIFF stores pixels in fixed blocks, commonly 256 or 512 square. When a chunk window straddles those boundaries, GDAL reads every overlapping block in full and throws away the parts outside the window. A 700-pixel chunk against a 512 block reads two full block rows and two full block columns per window — roughly 4x the bytes you actually use. Throughput collapses and workers spend their time waiting on I/O rather than warping.

Second, assuming every window is a full chunk. Raster dimensions are almost never an exact multiple of the chunk size, so the last row and column of windows are smaller remainders. Code that hard-codes the chunk edge into Window(col_off, row_off, edge, edge) reads past the raster bounds, which either raises a RasterioIOError or silently pads with nodata and corrupts the output mosaic edge.

The fix for both is in the implementation: snap the chunk edge with align_chunk_to_block() so reads are aligned, and clip each window with min(edge, width - col_off) and min(edge, height - row_off) so the remainder tiles read only real pixels.

Verification

Run a throughput sweep across candidate chunk sizes to confirm the block-aligned choice actually wins on your storage and dtype:

# Sweep chunk sizes and time each run. Aligned sizes should dominate.
for chunk in 256 512 700 1024 2048; do
  printf "chunk=%-5s " "$chunk"
  /usr/bin/time -f "%e s  %M KB" \
    python chunked_warp.py source.tif "/tmp/out_${chunk}.tif" \
    --crs EPSG:3857 --chunk "$chunk" 2>&1 | tail -n1
done

# Confirm the output tiling matches the chosen chunk edge.
python3 - <<'EOF'
import rasterio
with rasterio.open("/tmp/out_1024.tif") as ds:
    print("block_shapes:", ds.block_shapes[0])   # expect (1024, 1024)
    print("size:", ds.width, "x", ds.height)
    print("crs:", ds.crs)                          # expect EPSG:3857
EOF

The 700 run should be visibly slower and read more bytes than either 512 or 1024 despite sitting between them, confirming that alignment beats raw chunk size. An exit code of 0 and a block_shapes of (1024, 1024) confirm a clean, block-aligned output.

FAQ

What chunk size should I start with for a raster warp?

Start at an integer multiple of the source block edge, usually 512 or 1024 pixels square. Read block_shapes first, then round your target chunk to the nearest multiple so each read maps cleanly onto stored tiles. Sweep 512, 1024, and 2048 to find the throughput peak for your specific storage and dtype.

Why do non-block-aligned chunks slow the warp down?

A tiled GeoTIFF is stored as fixed blocks, often 256 or 512 pixels square. When a chunk window straddles block boundaries, GDAL must read every overlapping block in full and discard the parts outside the window. This read amplification can double or triple the bytes pulled from disk, so snapping chunk edges to a multiple of the block edge removes the waste.

How do I stop chunk workers from running out of memory?

Estimate the per-chunk footprint as width * height * bands * dtype_size, then multiply by the worker count and a safety factor of two to three for GDAL warp scratch buffers. Keep that product under your RAM budget. If it exceeds the budget, shrink the chunk edge to the next smaller block multiple or reduce the worker count.

Should I handle remainder chunks at the raster edge?

Yes. Raster dimensions are rarely an exact multiple of the chunk size, so the last row and column of windows are smaller remainders. Clip each window to the raster bounds rather than assuming a full chunk. A window that reads past the raster edge raises a read error or pads with nodata, corrupting the output mosaic.