Parallelising Tile Inference with Ray Actors

Use Ray actors to hold a GPU model across thousands of geospatial tiles: actor pools, fractional GPU allocation, backpressure, failure isolation and object-store hygiene.

A Ray actor is a worker process that keeps state between calls. For geospatial inference that state is the expensive part: a GPU model, an ONNX session, a warmed CUDA context. Create a pool of actors sized to the hardware, stream tiles through it with ActorPool or a ray.wait loop, and have each call return an object key rather than an array.

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

Tasks reload the model; actors keep it On the left, each stateless task spends most of its time loading the model and only a fraction predicting. On the right, four actors load the model once at startup and then spend nearly all of their time predicting, so the same hardware processes far more tiles per minute. stateless tasks actor pool load model 4.2 s 0.4 s load model 4.2 s 0.4 s load model 4.2 s 0.4 s 91% of wall clock spent loading 6 000 tiles ≈ 7.7 hours load once predict · predict · predict · … load once predict · predict · predict · … load once predict · predict · predict · … load cost amortised over 1 500 tiles each 6 000 tiles ≈ 11 minutes The heavier the setup, the larger the gap — and GPU models have very heavy setup.

Why This Fails in Geospatial ML Pipelines

Submitting everything at once is the most common Ray mistake and it kills the driver rather than the workers. [actor.predict.remote(t) for t in tiles] for a hundred thousand tiles creates a hundred thousand pending futures; if each returns an array, the object store fills, spills to disk, and the run grinds to a halt with no error that points at the cause. The fix is backpressure — keep a bounded window of in-flight calls — and returning object keys rather than data.

The second failure is fractional GPU arithmetic that ignores memory. num_gpus=0.25 is a scheduling hint: Ray will place four actors on one device, but it does not partition VRAM. Four actors each holding a 6 GB model on a 16 GB card produces an out-of-memory error at the fourth actor’s first call, long after the pool appeared to start cleanly. The fraction must be derived from measured footprint, not chosen for tidiness.

Third, actor death is silent to the pool. An actor killed by the OOM killer leaves its in-flight call raising RayActorError, and code that does not catch it loses those tiles. With max_restarts and per-call exception handling, the same failure costs one tile; without them it costs whatever was in flight, and the gap only appears at the validation gate — which is exactly what that gate exists for, as described in orchestrating geospatial ML pipelines.

Core Principles

  • Put the expensive state in __init__. That is the whole reason to use an actor.
  • Size the pool from measured VRAM, then declare the matching fraction.
  • Apply backpressure with ray.wait and a bounded in-flight window.
  • Return keys, not arrays. Keep the object store small.
  • Set max_restarts and catch RayActorError so a dead actor costs one tile.
  • Pin intra-op threads to 1 in the model runtime; Ray owns the parallelism.

Production-Ready Code

from __future__ import annotations

import logging
from dataclasses import dataclass

import numpy as np
import ray

logger = logging.getLogger(__name__)


@ray.remote(num_gpus=0.25, max_restarts=2, max_task_retries=1)
class TilePredictor:
    """Holds the model for its lifetime; tiles stream through it.

    Returns the written object key, never the array, so the object store stays small.
    """

    def __init__(self, model_uri: str, mean: np.ndarray, std: np.ndarray,
                 out_prefix: str):
        import onnxruntime as ort                      # imported in the actor process
        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = 1
        self.sess = ort.InferenceSession(model_uri, sess_options=opts,
                                         providers=providers)
        self.mean, self.std = mean, std
        self.out_prefix = out_prefix
        self.n_done = 0

    def predict_tile(self, tile_key: str) -> str:
        from geoml.io import read_tile, write_tile     # noqa: PLC0415

        block = read_tile(tile_key)                    # (bands, y, x) float32
        n_bands, ny, nx = block.shape
        flat = block.reshape(n_bands, -1).T
        out = np.full(flat.shape[0], 255, dtype="uint8")

        valid = np.isfinite(flat).all(axis=1)
        if valid.any():
            x = ((flat[valid] - self.mean) / self.std).astype("float32")
            logits = self.sess.run(None, {self.sess.get_inputs()[0].name: x})[0]
            out[valid] = np.argmax(logits, axis=1).astype("uint8")

        self.n_done += 1
        return write_tile(out.reshape(1, ny, nx), tile_key, self.out_prefix)

    def stats(self) -> dict:
        return {"tiles": self.n_done}


def run_pool(tile_keys: list[str], model_uri: str, mean, std, out_prefix: str,
             n_actors: int = 8, in_flight: int = 32) -> tuple[list[str], list[str]]:
    """Stream tiles through an actor pool with bounded backpressure.

    Returns (written keys, failed tile keys).
    """
    actors = [TilePredictor.remote(model_uri, mean, std, out_prefix)
              for _ in range(n_actors)]

    pending: dict = {}
    queue = list(tile_keys)
    written, failed = [], []
    next_actor = 0

    def submit() -> None:
        nonlocal next_actor
        while queue and len(pending) < in_flight:
            key = queue.pop()
            actor = actors[next_actor % len(actors)]
            next_actor += 1
            pending[actor.predict_tile.remote(key)] = key

    submit()
    while pending:
        done, _ = ray.wait(list(pending), num_returns=1, timeout=None)
        ref = done[0]
        tile = pending.pop(ref)
        try:
            written.append(ray.get(ref))
        except ray.exceptions.RayActorError:
            logger.warning("actor died processing %s — requeueing once", tile)
            queue.append(tile) if tile not in failed else failed.append(tile)
        except Exception as exc:                       # deterministic tile failure
            logger.error("tile %s failed: %s", tile, exc)
            failed.append(tile)
        submit()

    for a in actors:
        logger.info("actor processed %s", ray.get(a.stats.remote()))
    return written, failed

Step-by-Step Walkthrough

Step 1 — Measure the model’s footprint before sizing the pool

import subprocess
import ray

ray.init(address="auto")
print(ray.cluster_resources())      # {'GPU': 2.0, 'CPU': 32.0, ...}

# Start ONE actor and read the device memory it actually uses.
probe = TilePredictor.remote(MODEL_URI, mean, std, "s3://out/probe/")
ray.get(probe.predict_tile.remote(tile_keys[0]))
print(subprocess.run(["nvidia-smi", "--query-gpu=memory.used",
                      "--format=csv"], capture_output=True, text=True).stdout)
ray.kill(probe)

If one actor uses 3.4 GB on a 16 GB card, four per device is safe and num_gpus=0.25 is honest. Guessing here is what produces the out-of-memory at actor four.

Step 2 — Run with backpressure

written, failed = run_pool(tile_keys, MODEL_URI, mean, std,
                           out_prefix="s3://predictions/_staging/2026-08-04/",
                           n_actors=8, in_flight=32)
print(f"{len(written)} tiles written, {len(failed)} failed")

in_flight=32 for eight actors keeps four calls queued per actor — enough that no actor idles, small enough that the object store never grows.

Step 3 — Watch the right dashboard panels

ray dashboard        # actor view: state, restarts, per-actor task count

Two signals matter: actors sitting in IDLE while tiles remain (the in-flight window is too small, or reads are the bottleneck), and a rising restart count (memory pressure).

Step 4 — Hand the failures to the gate

if failed:
    quarantine_tiles(failed, run_id=spec.run_id)
report = validate_with_tolerance(spec, written, max_missing_frac=0.002)
if not report["ok"]:
    raise RuntimeError(report["failures"])

The pool’s job is to process what it can and report honestly; the decision about whether the run publishes belongs to the validation gate.

Bounded in-flight calls keep the object store flat Two object store memory traces are compared. With unbounded submission the memory rises steadily until the store spills to disk and throughput collapses. With a bounded in-flight window the memory stays flat and throughput remains constant for the whole run. Object store memory over a 6 000-tile run RAM spill threshold unbounded submission in_flight = 32 Once the store spills, every subsequent get() pays a disk round trip. Returning keys instead of arrays keeps the flat line flat.

Verification

import numpy as np
import pytest
import ray


def test_actor_loads_the_model_once():
    a = TilePredictor.remote(MODEL_URI, mean, std, "s3://out/test/")
    ray.get([a.predict_tile.remote(k) for k in tile_keys[:5]])
    assert ray.get(a.stats.remote())["tiles"] == 5


def test_pool_returns_keys_not_arrays():
    written, _ = run_pool(tile_keys[:20], MODEL_URI, mean, std, "s3://out/test/",
                          n_actors=2, in_flight=4)
    assert all(isinstance(w, str) and w.startswith("s3://") for w in written)


def test_in_flight_window_is_respected(monkeypatch):
    seen = []
    real_wait = ray.wait
    def spy(refs, **kw):
        seen.append(len(refs))
        return real_wait(refs, **kw)
    monkeypatch.setattr(ray, "wait", spy)
    run_pool(tile_keys[:50], MODEL_URI, mean, std, "s3://out/test/",
             n_actors=2, in_flight=8)
    assert max(seen) <= 8


def test_dead_actor_does_not_lose_the_run():
    written, failed = run_pool(tile_keys[:30], MODEL_URI, mean, std, "s3://out/test/",
                               n_actors=2, in_flight=4)
    assert len(written) + len(failed) == 30


def test_predictions_match_single_process(small_tiles):
    ref = [predict_single_process(k) for k in small_tiles]
    got, _ = run_pool(small_tiles, MODEL_URI, mean, std, "s3://out/test/",
                      n_actors=2, in_flight=2)
    assert sorted(read_all(got)) == sorted(ref)

The accounting test — written plus failed equals submitted — is the one that catches silently dropped tiles, which is the failure mode most likely to reach production unnoticed.

Task or actor: decide on setup cost and reusable state A two column comparison. Ray tasks suit cheap setup, no reusable state and per call isolation. Ray actors suit expensive setup such as a GPU model, state that is reused across many calls, and require explicit restart handling. Geospatial tile inference with a neural model falls clearly on the actor side. Which one fits the work? ray.remote task setup is cheap (< 100 ms) nothing worth keeping between calls failure isolated per call scheduler places it anywhere e.g. reprojecting a vector chunk ray.remote actor setup is expensive (seconds, GPU) state reused across thousands of calls needs max_restarts and error handling pinned to a device fraction e.g. segmentation over 6 000 tiles Mixed pipelines often use both: tasks for I/O and feature prep, actors for the model.

FAQ

When is an actor better than a plain Ray task?

When per-call setup is expensive and reusable — a GPU model, a warmed session, an open connection. Tasks rebuild that state every call. For stateless work, tasks are simpler and schedule better.

How do I share one GPU between several actors?

Declare a fraction such as num_gpus=0.25. Ray schedules accordingly but does not partition VRAM, so derive the fraction from measured memory use, as in Step 1.

Why does my driver run out of memory submitting the work?

Because every future and every returned object is being held. Bound the in-flight window with ray.wait and return object keys instead of arrays.

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