Profiling Memory in a Raster Worker

Run one tile under memray and look at what was allocated at the peak:

python -m memray run -o tile.bin process_one_tile.py 33UVP_0412
python -m memray flamegraph tile.bin        # open memray-flamegraph-tile.html
python -m memray stats tile.bin             # peak memory and top allocators

When a worker is killed for exceeding its memory limit, the process dies before it can say why. Reproducing the task alone, under a memory profiler, answers the question in minutes. This page belongs to monitoring and observability for raster pipelines in Cloud Execution & Orchestration.

Where Raster Memory Usually Goes

The usual suspects at peak In a typical over-budget raster task, peak memory is dominated by a full-resolution read where a window would do, followed by an accidental float64 copy of an integer array, then intermediates kept alive until the function returns, with GDAL's block cache and the interpreter itself as smaller contributors. Anatomy of a peak full-resolution read float64 copy live intermediates GDAL cache runtime The first three are code choices; all three are usually cheap to fix once seen.

Raster tasks have a characteristic memory profile: a few very large arrays and very little else. That makes them easier to profile than general applications — the flame graph usually has one or two wide towers — and also means small code choices have large consequences. Casting a 10,980 × 10,980 int16 band to float64 turns 240 MB into 960 MB, and doing it for four bands at once can be the whole story of an out-of-memory kill.

Environment & Setup

Package Version pin Used for
memray >=1.11 Allocation tracking and flame graphs (Linux, macOS)
psutil >=5.9 Resident memory from inside the process
rasterio >=1.3.0 The code under test
numpy >=1.23 The code under test 
pip install "memray>=1.11" "psutil>=5.9" "rasterio>=1.3.0" "numpy>=1.23"

Complete Working Example

# process_one_tile.py — reproduce a single task outside Dask
import sys
import tracemalloc

import numpy as np
import psutil
import rasterio
from rasterio.windows import from_bounds


def rss_mb() -> float:
    return psutil.Process().memory_info().rss / 1e6


def ndvi_tile_bad(red_path, nir_path):
    with rasterio.open(red_path) as r, rasterio.open(nir_path) as n:
        red = r.read(1).astype("float64")        # full scene, float64
        nir = n.read(1).astype("float64")
    diff = nir - red                              # extra full-size temporaries
    total = nir + red
    return diff / total


def ndvi_tile_good(red_path, nir_path, bounds):
    with rasterio.open(red_path) as r, rasterio.open(nir_path) as n:
        win = from_bounds(*bounds, transform=r.transform)
        red = r.read(1, window=win, out_dtype="float32")   # only the tile, float32
        nir = n.read(1, window=win, out_dtype="float32")
    out = np.subtract(nir, red)
    np.divide(out, nir + red, out=out, where=(nir + red) != 0)
    del red, nir                                  # release before returning
    return out


if __name__ == "__main__":
    tracemalloc.start()
    before = rss_mb()
    res = ndvi_tile_good(sys.argv[1], sys.argv[2], tuple(map(float, sys.argv[3:7])))
    cur, peak = tracemalloc.get_traced_memory()
    print(f"rss {before:.0f} -> {rss_mb():.0f} MB, numpy peak {peak / 1e6:.0f} MB, out {res.nbytes / 1e6:.0f} MB")

The two functions compute the same thing for a tile. The first reads whole scenes, promotes them to float64 and creates two full-size temporaries; the second reads only the window it needs, in float32, and computes into a preallocated output. On a Sentinel-2 scene with a 2 km tile, the difference is roughly two gigabytes against a few tens of megabytes. Choosing what resolution and extent to read is covered in handling pixel resolution and scaling.

memray versus tracemalloc

Two profilers, different reach memray tracks every allocation including those made by native extensions such as GDAL and NumPy's C code, and produces flame graphs and peak reports, at the cost of running under its own launcher. tracemalloc is in the standard library and tracks Python-level allocations, including NumPy arrays, but not GDAL's own buffers and caches. Which profiler to reach for memray native allocations too (GDAL, C code) flame graph at the peak Linux and macOS tracemalloc standard library, no install Python and NumPy allocations misses GDAL's own buffers Start with memray for a kill you cannot explain; use tracemalloc for quick checks inside tests.

GDAL allocates its block cache and read buffers in C, outside Python’s allocator, so a profiler that only sees Python allocations can report a modest peak while the process is actually far larger. memray intercepts native allocations and shows them attributed to the Python line that triggered them, which is what makes it the right tool for raster work. Comparing tracemalloc’s peak with resident memory from psutil is a quick way to see how much of the footprint is native.

Controlling GDAL’s Share

GDAL’s block cache defaults to 5% of physical memory per process. On a worker with 64 GB and eight processes, that is up to 25 GB of cache across the node — often the unexplained part of a memory budget. Setting GDAL_CACHEMAX explicitly, for example to 256 MB per process, makes the footprint predictable. The cache still helps when a task reads the same blocks repeatedly, as in overlapping windows or multi-pass algorithms, so shrinking it to almost nothing can slow tasks down; measure both runtime and peak when changing it. Setting it in the worker environment rather than in code ensures it takes effect before GDAL initialises, and is covered with other settings in caching PROJ data and GDAL config in containers.

From One Tile to the Cluster

Once a single tile’s peak is known, the worker budget follows: peak per task times tasks running concurrently per worker, plus GDAL cache and a margin for the runtime. If a worker runs four threads and each task peaks at 1.5 GB, a 4 GB memory limit will fail regardless of tuning elsewhere. Either reduce the per-task peak, as above, or run fewer threads per worker. Recording the peak memory of every tile in the structured events, as in logging structured events from raster tasks, shows which tiles approach the limit before any of them cross it.

Verification

Measure the fix, do not assume it Peak memory for the same tile before and after changes: reading a window instead of the full scene and keeping float32 instead of float64 reduce the peak from around 2.1 GB to under 100 MB. The output values are identical within float32 precision. Peak memory, same tile before — 2.1 GB after — 90 MB Also assert the outputs match, so the saving is not a bug. 
import numpy as np
import tracemalloc

for fn, args in ((ndvi_tile_bad, (RED, NIR)), (ndvi_tile_good, (RED, NIR, BOUNDS))):
    tracemalloc.start(); out = fn(*args); _, peak = tracemalloc.get_traced_memory(); tracemalloc.stop()
    print(fn.__name__, f"{peak / 1e6:.0f} MB")

a = ndvi_tile_bad(RED, NIR)[ROW_SLICE, COL_SLICE]
b = ndvi_tile_good(RED, NIR, BOUNDS)
assert np.allclose(a, b, atol=1e-6, equal_nan=True)

Common Errors

The profiler shows a small peak but the worker is killed

Native allocations are invisible to tracemalloc. Use memray, and check GDAL_CACHEMAX.

Memory keeps growing across tiles

References to previous tiles’ arrays are kept, often in a list or a cache. Release them, or run each tile in a fresh task.

The local reproduction does not use as much memory

The worker runs several tasks concurrently. Multiply the per-task peak by threads per worker.

memray is not available on Windows

Profile inside a Linux container that matches the worker image.

Frequently Asked Questions

Q: Why is my raster worker killed for memory? Usually because a task reads more than it needs, promotes integers to float64, or holds intermediates while several tasks run concurrently. Profiling one tile shows which.

Q: Which memory profiler works with GDAL? memray, because it tracks native allocations made by GDAL and NumPy’s C code and attributes them to Python lines. tracemalloc sees only Python-level allocations.

Q: How much memory does GDAL’s cache use? By default 5% of physical memory per process, which adds up across processes. Set GDAL_CACHEMAX explicitly to make it predictable.

Q: How do I size worker memory from a profile? Multiply the per-task peak by the number of concurrent tasks per worker, add GDAL cache and a margin for the Python runtime.