Right-Sizing Workers for Raster Mosaic Jobs

Worker memory is not a property of the data; it is a product of two numbers you choose — the per-unit peak and how many units run at once — plus a headroom factor. Fleets are almost always sized for the largest input ever seen, at every level of concurrency, which means paying four times over for a case that occurs once in a hundred. Measuring the distribution and routing its tail to a second pool typically halves the compute bill without touching the code that does the work.

When to Use This Pattern#

• Workers are provisioned by guess and nobody can say what the per-tile peak actually is.
• Out-of-memory kills happen occasionally, and the response so far has been to double the memory.
• Input sizes vary by an order of magnitude, which is normal for scene-based sources.
• Compute is a visible line on the bill and the fleet runs well below its provisioned memory.

Complete Working Example#

Measure per unit, estimate before allocating, and route by the estimate.

from __future__ import annotations

import resource
from dataclasses import dataclass

import rasterio
from prefect import flow, task, get_run_logger
from prometheus_client import Histogram

PEAK_MB = Histogram(
    "tile_peak_memory_mb", "Marginal peak RSS per tile task",
    buckets=(64, 128, 256, 512, 1024, 2048, 4096, 8192), labelnames=("layer",),
)


@dataclass(frozen=True)
class SizeEstimate:
    megapixels: float
    bands: int
    bytes_per_px: int

    @property
    def working_set_mb(self) -> float:
        # Source window, warped destination, and one working copy: three times the
        # naive figure is the estimate that has matched measurement most closely.
        return 3 * self.megapixels * self.bands * self.bytes_per_px


@task(timeout_seconds=120)
def estimate(src_uri: str, tile: Tile) -> SizeEstimate:
    """Header read only. Costs one range request and decides which pool to use."""
    with rasterio.open(src_uri) as src:
        w = src.window(*tile.bounds_in(src.crs))
        return SizeEstimate(
            megapixels=(w.width * w.height) / 1e6,
            bands=src.count,
            bytes_per_px=rasterio.dtypes.dtype_ranges[src.dtypes[0]] and src.dtypes[0].itemsize,
        )


@task(retries=2, timeout_seconds=900, tags=["tile"])
def build_tile(tile: Tile, layer: str, src_uri: str, dest: str) -> str:
    start = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    warp_window(src_uri, tile, dest)
    peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    # Marginal, not absolute: in a reused process the absolute figure only rises.
    PEAK_MB.labels(layer=layer).observe(max(peak - start, 0) / 1024)
    return dest


@flow(name="mosaic")
def mosaic(tiles: list[Tile], layer: str, src_uri: str, dest: str, big_mb: float = 1500):
    sizes = estimate.map(src_uri=src_uri, tile=tiles)
    small = [t for t, s in zip(tiles, sizes) if s.result().working_set_mb < big_mb]
    large = [t for t, s in zip(tiles, sizes) if s.result().working_set_mb >= big_mb]
    get_run_logger().info("routing: %d small, %d large", len(small), len(large))

    build_tile.with_options(task_run_name="small-{tile.path}").map(
        small, layer=layer, src_uri=src_uri, dest=dest)
    build_tile.with_options(tags=["tile", "bigmem"]).map(
        large, layer=layer, src_uri=src_uri, dest=dest)

The worker size itself is then arithmetic rather than judgement:

def worker_memory_gb(p99_peak_mb: float, concurrency: int, headroom: float = 1.3) -> float:
    """The only three numbers that matter, and two of them are settings."""
    return (p99_peak_mb * concurrency * headroom) / 1024


assert worker_memory_gb(p99_peak_mb=420, concurrency=8, headroom=1.3) == pytest.approx(4.3)

The three-times factor in the estimate is empirical and worth explaining, because a naive calculation of width times height times bands times bytes is invariably too low. A warp holds the source window and the destination array simultaneously, resampling needs working space proportional to the kernel, and the library will typically hold a copy during any type conversion. Three times the naive figure has matched measured peaks closely enough across bilinear and cubic resampling to be a usable routing threshold; it is not a substitute for measuring, but it is a good enough estimate to decide which pool a unit belongs in before allocating anything.

Parameter & Option Reference#

Verification & Testing#

def test_estimate_is_conservative_against_measurement(sample_tiles) -> None:
    for tile in sample_tiles:
        est = estimate.fn(SRC, tile).working_set_mb
        actual = measure_peak(lambda: warp_window(SRC, tile, "/dev/null"))
        assert est >= actual * 0.8, f"estimate {est:.0f} well below measured {actual:.0f}"


def test_routing_sends_big_tiles_to_the_big_pool(monkeypatch) -> None:
    routed: list[str] = []
    monkeypatch.setattr("mymodule.submit", lambda t, tags: routed.append(tags[-1]))
    mosaic(mixed_tiles(), layer="ortho", src_uri=SRC, dest=DEST)
    assert routed.count("bigmem") == 3, "the tail is not being routed"


def test_worker_memory_formula() -> None:
    assert worker_memory_gb(420, 8) == pytest.approx(4.27, abs=0.05)
    assert worker_memory_gb(420, 16) == pytest.approx(8.53, abs=0.05)


def test_gdal_cache_is_bounded() -> None:
    # The default can be a large fraction of the machine, per process.
    assert int(os.environ["GDAL_CACHEMAX"]) <= 512

The GDAL_CACHEMAX assertion catches a genuinely confusing failure. The cache is per process and its default on some builds is a percentage of physical memory, so a worker running eight tiles in eight processes can allocate eight large caches on top of eight working sets — and the resulting out-of-memory kill looks like the tiles being bigger than measured, because nothing in the task’s own memory accounting includes the driver’s cache. Pinning it explicitly makes the arithmetic in worker_memory_gb true rather than approximately true.

There is one more reason to prefer two pools over one large fleet, and it has nothing to do with money. A worker sized for the outlier will happily run the outlier at concurrency eight, which means eight outliers at once — and that configuration has never been tested, because outliers are rare enough that they almost never coincide by chance. The first night they do, the fleet discovers a memory profile nobody has measured. Routing the tail to a pool with a concurrency of one removes the coincidence entirely: the large tiles are serialised by construction, which is both cheaper and more predictable than trusting that they will stay rare.

The same logic argues for keeping the routing threshold well below the large pool’s actual capacity. A threshold set at the exact point where the small pool fails leaves no margin for an estimate that came in low, and estimates do come in low — a source with an unexpected alpha band, a nodata mask that materialises as an extra array. Setting the threshold at roughly half the small worker’s per-slot budget sends slightly too much work to the large pool and costs very little, because the large pool is one machine and the alternative is an out-of-memory kill in the middle of a nightly window.

Common Pitfalls#

• Sizing for the largest input. It multiplies every worker’s memory to accommodate a case that happens once in a hundred.
• Using absolute ru_maxrss in a reused process. It is a high-water mark and will size the fleet for the biggest task the worker ever ran.
• Leaving the driver cache at its default. Per-process caches on a concurrent worker can exceed the working sets they accompany.
• Raising concurrency to raise throughput. Memory scales with it linearly; the limit is usually reached long before the CPU is.
• Estimating from width times height times bands. A warp holds several arrays at once; a factor of about three matches measurement.
• Routing after allocation. The decision must come from a header read, or the first thing the outlier does is exhaust the small worker.

Frequently Asked Questions#

Related#

• Cost optimization for spatial compute — where this lever sits among the others
• Using spot instances for interruptible raster jobs — the price lever
• Limiting DAG fan-out with concurrency groups — where the concurrency number is applied
• Offloading GDAL work to a process pool — per-process memory in practice
• Containerizing GDAL workers with Docker — where the cache setting belongs

← Cost Optimization for Spatial Compute