Geospatial Orchestration Hub

Branching Workflows Based on Spatial Extent

The extent predicate is the most useful branch in a raster pipeline, and the one most often written as a hard-coded number that stops being true the moment the workers are resized. Compute the output megapixel count from the source headers and the target CRS, derive the threshold from the worker’s actual memory, and record both with the decision. A scene under the threshold takes the single-pass path; one above it is split into windows and merged. The branch then survives a change of instance type without anyone remembering to update a constant.

When to Use This Pattern

  • Scene sizes vary by more than an order of magnitude, which is normal for anything covering both islands and continents.
  • A single-pass warp occasionally exhausts memory, and the failures correlate with geography rather than with time.
  • The tiled path is meaningfully slower — merging has real cost — so it is not acceptable as the universal default.
  • Workers are heterogeneous or resized periodically, so a fixed threshold will drift out of correctness.

Complete Working Example

The prediction runs on headers, and the threshold comes from the machine actually executing the task.

from __future__ import annotations

import math
import os
from dataclasses import dataclass
from enum import Enum
from typing import Optional

import psutil
import rasterio
from rasterio.warp import calculate_default_transform


class Route(str, Enum):
    SINGLE_PASS = "single_pass"
    TILED = "tiled"


@dataclass(frozen=True)
class ExtentDecision:
    route: Route
    output_megapixels: float
    estimated_peak_gb: float
    threshold_gb: float
    reason: str


# A warp holds roughly: the source window, the destination block, and GDAL's warp
# buffers. Four bytes per pixel per band is a reasonable float32 working estimate.
BYTES_PER_PIXEL_BAND = 4
WARP_OVERHEAD = 2.5          # measured: peak is ~2.5x the naive array size


def usable_memory_gb(reserve_gb: float = 2.0, concurrency: int = 1) -> float:
    """What THIS task may use, on THIS machine, right now."""
    total = psutil.virtual_memory().total / 1024**3
    return max(1.0, (total - reserve_gb) / max(1, concurrency))


def predict(source_uri: str, dst_crs: str, band_count: Optional[int] = None) -> tuple[float, float]:
    """(output megapixels, estimated peak GB) — from headers only."""
    with rasterio.open(source_uri) as src:
        transform, width, height = calculate_default_transform(
            src.crs, dst_crs, src.width, src.height, *src.bounds
        )
        bands = band_count or src.count
    megapixels = (width * height * bands) / 1_000_000
    peak_gb = (megapixels * 1_000_000 * BYTES_PER_PIXEL_BAND * WARP_OVERHEAD) / 1024**3
    return megapixels, peak_gb


def choose_route(source_uri: str, dst_crs: str, concurrency: int = 1) -> ExtentDecision:
    threshold = usable_memory_gb(concurrency=concurrency)
    try:
        megapixels, peak = predict(source_uri, dst_crs)
    except rasterio.RasterioIOError as exc:
        # Cannot inspect: take the path that cannot run out of memory.
        return ExtentDecision(Route.TILED, 0.0, 0.0, threshold,
                              f"header unreadable ({exc}) — defaulting to tiled")

    if peak > threshold:
        return ExtentDecision(
            Route.TILED, megapixels, peak, threshold,
            f"{peak:.1f} GB estimated peak exceeds the {threshold:.1f} GB budget",
        )
    return ExtentDecision(
        Route.SINGLE_PASS, megapixels, peak, threshold,
        f"{peak:.1f} GB fits inside the {threshold:.1f} GB budget",
    )

The two paths follow directly, and both end in the same place — a GeoTIFF in the target CRS with the same nodata and the same creation options.

def warp_single_pass(source_uri: str, dest: str, dst_crs: str) -> str:
    run_gdal(["gdalwarp", "-t_srs", dst_crs, "-r", "bilinear",
              "-dstnodata", "-9999", "-co", "COMPRESS=DEFLATE", "-co", "TILED=YES",
              "-wm", "512", source_uri, dest])
    return dest


def warp_tiled(source_uri: str, dest: str, dst_crs: str, window_px: int = 8192) -> str:
    """Warp in windows, then build a VRT and translate once.

    A VRT costs nothing to build and defers the merge to a single sequential
    pass, which is far cheaper in memory than merging arrays in Python.
    """
    parts = []
    for i, window in enumerate(windows_for(source_uri, window_px)):
        part = f"{dest}.{i:04d}.tif"
        run_gdal(["gdalwarp", "-t_srs", dst_crs, "-r", "bilinear",
                  "-dstnodata", "-9999", "-te", *map(str, window),
                  "-co", "COMPRESS=DEFLATE", "-wm", "512", source_uri, part])
        parts.append(part)

    run_gdal(["gdalbuildvrt", f"{dest}.vrt", *parts])
    run_gdal(["gdal_translate", "-co", "COMPRESS=DEFLATE", "-co", "TILED=YES",
              f"{dest}.vrt", dest])
    return dest
The threshold is a property of the worker, not of the dataEstimated peak memory rises linearly with output megapixels. Three horizontal threshold lines at eight, sixteen and sixty-four gigabytes intersect it at very different scene sizes.64 GB32 GB0output megapixels8 GB worker — tiles almost everything16 GB worker64 GB worker — single-pass for nearly all scenesA hard-coded megapixel threshold picks one of these lines and freezes it, whatever machine the task lands on.
The blue line is the data; the dashed lines are the machine. Only one of them changes when someone resizes the worker pool, and it is not the one usually written into the code.

The VRT step in the tiled path deserves a note, because it is what keeps that path from re-creating the problem it exists to solve. A naive merge reads every window into memory and concatenates, which peaks at the size of the whole output — exactly the number that sent the scene down this path. A VRT is a small XML file describing where each window’s pixels live; gdal_translate then walks it block by block, holding one block at a time. The whole merge runs in a few hundred megabytes regardless of how large the mosaic is, and it costs one sequential read and write.

Merging without undoing the tilingConcatenating windows in memory peaks at the full output size, defeating the purpose of tiling. Building a VRT and translating holds one block at a time.CONCATENATE IN MEMORYpeak = the whole output, 26 GBBUILD A VRT, THEN TRANSLATEVRT · 4 KBpeak 340 MBSame windows, same output bytes. The upper row simply reassembles the problem the tiling removed.
The VRT is the cheapest object in the pipeline and the reason the tiled path is bounded rather than merely deferred.

Parameter & Option Reference

Parameter Type Default Spatial notes
BYTES_PER_PIXEL_BAND int 4 Float32 working estimate. Raise to 8 for float64 elevation, lower to 1 for byte imagery that stays byte throughout.
WARP_OVERHEAD float 2.5 Measured, not derived: source window, destination block and GDAL’s own buffers. Re-measure after a GDAL upgrade.
reserve_gb float 2.0 Left for the OS, the Python interpreter and anything else on the box.
concurrency int 1 Divides the budget. Four concurrent warps on one worker each get a quarter.
window_px int 8192 Window edge on the tiled path. Large enough that per-window overhead is amortised, small enough to fit comfortably.
-wm MB 512 GDAL’s own warp memory. Included in the overhead factor above, so changing one means re-measuring the other.

Verification & Testing

The estimate does not need to be accurate; it needs to be ordered and safe. Assert that bigger scenes predict bigger peaks, and that the tiled path is chosen before memory actually runs out.

def test_prediction_is_monotonic(small_cog, large_cog) -> None:
    small_mp, small_peak = predict(small_cog, "EPSG:3857")
    large_mp, large_peak = predict(large_cog, "EPSG:3857")
    assert large_mp > small_mp
    assert large_peak > small_peak


def test_threshold_follows_the_machine(monkeypatch, large_cog) -> None:
    monkeypatch.setattr(psutil, "virtual_memory", lambda: FakeMem(total=8 * 1024**3))
    assert choose_route(large_cog, "EPSG:3857").route is Route.TILED

    monkeypatch.setattr(psutil, "virtual_memory", lambda: FakeMem(total=128 * 1024**3))
    assert choose_route(large_cog, "EPSG:3857").route is Route.SINGLE_PASS


def test_unreadable_source_defaults_to_tiled() -> None:
    decision = choose_route("/vsis3/missing/nope.tif", "EPSG:3857")
    assert decision.route is Route.TILED
    assert "unreadable" in decision.reason


def test_prediction_is_conservative(sample_cogs) -> None:
    """Predicted peak must never be LOWER than the observed peak."""
    for cog in sample_cogs:
        _, predicted = predict(cog, "EPSG:3857")
        observed = measure_peak_gb(warp_single_pass, cog)
        assert predicted >= observed * 0.95, f"prediction under-estimates for {cog}"

That last test is the one that keeps the branch honest. An estimate that runs low is worse than no estimate at all, because it produces confident single-pass decisions that OOM — so when the measurement drifts, the correct response is to raise WARP_OVERHEAD rather than to widen the threshold.

The estimate must sit above the truthTwelve scenes plotted with predicted peak on one axis and observed peak on the other. Every point sits on or above the parity line, meaning the estimate never under-predicts.predictedobserved peak (GB)below this line is where the OOMs liveA point below the dashed line means a scene that predicted “fits” and did not. One is enough to change the constant.
Deliberately conservative: every point sits above parity, which costs a few unnecessary tiled runs and buys a branch that never chooses the path that fails.

Common Pitfalls

  • Predicting from the source size rather than the output. A reprojection into a much finer grid produces far more pixels than it consumes. calculate_default_transform gives the output shape for the cost of a header read, which is the number that matters.
  • A hard-coded megapixel threshold. It encodes a worker size that will change. Derive from psutil.virtual_memory() and divide by the task concurrency, and the branch stays correct through an instance-type change.
  • Ignoring concurrency. A threshold computed for one task on a 32 GB worker is wrong by a factor of four when four tasks share the machine. The divisor belongs in the calculation, not in a comment.
  • Merging tiled output in Python. Reading every window into numpy and concatenating rebuilds the exact memory peak the tiled path was avoiding. Build a VRT and let gdal_translate do one sequential pass.
  • Different creation options per path. If the single-pass output is tiled and compressed and the tiled output is not, downstream reads behave differently depending on a decision they cannot see. Keep the creation options identical.

Frequently Asked Questions

Why not always take the tiled path?

Because it is meaningfully slower — per-window process start-up, a VRT build, and a full re-encode in the translate step. On a pipeline where ninety per cent of scenes fit comfortably, always tiling adds thirty to forty per cent to the total runtime to protect against the ten per cent. The branch exists to pay that cost only where it buys something.

How do I measure `WARP_OVERHEAD` for my data?

Run a representative scene under psutil sampling, take the peak RSS of the GDAL process, and divide by the naive array size. Repeat for the largest, the most-banded and the highest-bit-depth scenes you have. Take the maximum, not the mean — the constant exists to protect the worst case, and averaging it away defeats it.

Does the same branch work for vector data?

The shape does; the predicate changes. For vector the useful measures are feature count and total vertex count, both of which ogrinfo -so reports without reading geometry. The threshold then bounds the memory a GeoDataFrame will occupy, which is far harder to estimate than a raster’s — so vector thresholds are usually set from measurement rather than from a formula.

What if the estimate says tiled but the tiled path also fails?

Then window_px is too large for this scene, and the same arithmetic applies one level down: the window’s output megapixels must fit the same budget. Deriving the window size from the threshold rather than fixing it at 8192 makes the tiled path self-adjusting, which is worth doing once the pipeline meets sources large enough to need it. See choosing tile sizes for raster partitioning.

← Conditional Branching in Geospatial DAGs

Related in this section