Orchestrating Geospatial ML Pipelines

Schedule and orchestrate geospatial ML pipelines with Airflow and Prefect: tile-level task granularity, idempotent writes, backfills, sensors and retry policies that survive partial failure.

A geospatial inference pipeline is a scheduled batch job with an unusually hostile failure profile. It waits on satellite scenes that arrive late or partially, processes gigabytes per tile, calls a model, and writes rasters to object storage — all across thousands of tiles, any one of which can fail on a transient network error. The orchestration question is not “how do I run these steps in order”; it is “what happens when tile 4,812 of 6,000 fails at 3 a.m.”

This topic is part of Geospatial MLOps and Model Deployment. Orchestration is the layer that ties together the containers from containerizing geospatial inference pipelines, the exported artifacts from ONNX export for geospatial model inference, and the monitoring described in model drift detection for geospatial inference. Where scaling batch inference with Dask and Ray answers how fast, orchestration answers when, in what order, and what happens on failure.

The shape of a tile-granular geospatial pipeline A schedule triggers a sensor that waits for the source scene to be complete. A tiling task then expands the graph into many mapped per-tile tasks, each of which extracts features, runs inference and writes to a temporary key. A validation gate inspects the produced tiles for coverage and distribution problems, and only then does a publish task atomically promote the mosaic and register the run. Fan out per tile · write to a staging key · validate · publish once sensor scene ready? tile plan 6 000 tiles tile 0001 features → infer → stage tile 0002 features → infer → stage tile 4812 — failed retry only this task validate coverage · CRS · stats publish atomic move + register Nothing downstream reads a staged tile, so a partial run is invisible to consumers. Publish is the only step that changes what the outside world sees.

Problem Framing

Four properties separate a geospatial pipeline from a generic ETL job.

Inputs arrive asynchronously and incompletely. A Sentinel-2 tile for a given date may appear hours after the nominal acquisition time, and a partially uploaded scene is readable but wrong. A scheduler that fires at a fixed time and assumes data exists will silently produce predictions from yesterday’s imagery.

Work is naturally partitioned by geography, and that partitioning is the failure boundary. If a task is “process the country”, one bad tile means the whole task fails and a retry redoes six hours of work. If a task is “process one tile”, the same failure costs ninety seconds.

Writes are large and non-atomic. Object stores make a PUT atomic per object, but a 400 MB GeoTIFF written incrementally through GDAL’s virtual filesystem is not. A worker killed at 60% leaves a file that opens without error and reads garbage past the truncation point.

Every artifact is versioned along two axes at once: model and data. A prediction mosaic is only meaningful as “model v7 applied to scene 2026-08-04”. Resolving the model version once per run and propagating it — rather than letting each task look up “latest” — is the difference between a reproducible mosaic and a patchwork.

Prerequisites & Environment Setup

# Pinned requirements for orchestration
apache-airflow==2.9.2
prefect==2.19.4
rasterio==1.3.10
boto3==1.34.131
pystac-client==0.8.2
pydantic==2.7.4

Install the orchestrator you actually intend to run — both are shown here for comparison, and installing both in one environment invites dependency conflicts:

# Airflow (use the official constraints file for a reproducible install)
pip install "apache-airflow==2.9.2" \
  --constraint "https://raw.githubusercontent.com/apache/airflow/constraints-2.9.2/constraints-3.11.txt"

# or Prefect
pip install "prefect==2.19.4"

pip install "rasterio==1.3.10" "boto3==1.34.131" "pystac-client==0.8.2" "pydantic==2.7.4"

GDAL/PROJ system dependencies (Ubuntu/Debian):

sudo apt-get install -y gdal-bin libgdal-dev libproj-dev

Keep the orchestrator and the geospatial runtime in separate images. The scheduler needs neither GDAL nor a model; the worker needs both. Mixing them produces an image where an Airflow upgrade can break PROJ — the exact coupling the container guidance in pinning GDAL and PROJ versions in a container build exists to prevent.

Step-by-Step Implementation

Step 1 — Model the run as data, not as a schedule

Everything downstream depends on a single immutable description of what this run is processing.

from dataclasses import dataclass, asdict
from datetime import date


@dataclass(frozen=True)
class RunSpec:
    """Immutable description of one pipeline run. Passed to every task."""
    scene_date: date
    collection: str            # e.g. "sentinel-2-l2a"
    model_uri: str             # resolved ONCE, e.g. "s3://models/landcover/v7/model.onnx"
    model_version: str         # "v7" — recorded in the output metadata
    grid_id: str               # the tiling scheme the outputs use
    run_id: str                # unique per attempt of this logical run

    def staging_prefix(self) -> str:
        return f"s3://predictions/_staging/{self.grid_id}/{self.scene_date:%Y-%m-%d}/{self.run_id}/"

    def published_prefix(self) -> str:
        return f"s3://predictions/{self.grid_id}/{self.scene_date:%Y-%m-%d}/"

Resolving model_uri once, at the top of the run, is the single highest-value line in this whole design.

Step 2 — Wait for the scene, do not assume it

from pystac_client import Client


def scene_is_ready(spec: RunSpec, bbox, min_assets: int = 12) -> bool:
    """True only when the catalogue lists a complete item for this date and area."""
    catalog = Client.open("https://earth-search.aws.element84.com/v1")
    search = catalog.search(
        collections=[spec.collection],
        bbox=bbox,
        datetime=f"{spec.scene_date}/{spec.scene_date}",
    )
    items = list(search.items())
    if not items:
        return False
    # A partially published item lists fewer assets than a complete one.
    return all(len(item.assets) >= min_assets for item in items)

In Airflow this becomes a PythonSensor with mode="reschedule" so it releases its worker slot between polls; in Prefect it is a task in a while loop with a timeout. Either way, the pipeline blocks rather than producing output from missing data.

A fixed schedule fires on the clock; a sensor fires on the data Two timelines run from four to seven in the morning. On the first, a fixed four o'clock trigger starts processing while the source scene is still uploading, so the run consumes an incomplete scene. On the second, a sensor polls until the catalogue reports the scene complete at five twenty, then starts the tile tasks against complete data. The scene decides when the run can start — not the clock fixed schedule scene uploading 04:00 trigger tiles run on partial data wrong output sensor gated scene uploading sensor polls, slot released between checks 05:20 complete tiles run on full scene correct output The cost is 80 minutes of latency. The alternative is a mosaic nobody can trust.

Step 3 — Make every write idempotent

import os
import tempfile
import boto3
import rasterio

s3 = boto3.client("s3")


def write_tile_atomically(array, profile, bucket: str, key: str) -> None:
    """Write locally, upload to a temp key, then copy into place.

    A retry re-runs all three steps and lands on exactly the same final object,
    and no reader ever observes a partial file at `key`.
    """
    tmp_key = f"{key}.{os.getpid()}.tmp"
    with tempfile.NamedTemporaryFile(suffix=".tif", delete=True) as fh:
        with rasterio.open(fh.name, "w", **profile) as dst:
            dst.write(array, 1)
        s3.upload_file(fh.name, bucket, tmp_key)

    s3.copy_object(Bucket=bucket, Key=key,
                   CopySource={"Bucket": bucket, "Key": tmp_key})
    s3.delete_object(Bucket=bucket, Key=tmp_key)


def tile_already_done(bucket: str, key: str) -> bool:
    """Cheap skip check so a retried run does not recompute finished tiles."""
    try:
        s3.head_object(Bucket=bucket, Key=key)
        return True
    except s3.exceptions.ClientError:
        return False

Step 4 — Fan out over tiles

The Airflow form, using dynamic task mapping:

from airflow.decorators import dag, task
from datetime import datetime


@dag(schedule="0 4 * * *", start_date=datetime(2026, 1, 1),
     catchup=True, max_active_runs=2, default_args={"retries": 3})
def landcover_inference():

    @task
    def build_spec(**context) -> dict:
        spec = RunSpec(
            scene_date=context["logical_date"].date(),
            collection="sentinel-2-l2a",
            model_uri=resolve_model_uri("landcover", stage="production"),
            model_version=resolve_model_version("landcover", stage="production"),
            grid_id="utm33n_10m",
            run_id=context["run_id"],
        )
        return asdict(spec)

    @task
    def list_tiles(spec: dict) -> list[str]:
        return tiles_for_area(spec["grid_id"])          # -> ["T33UUU_0001", ...]

    @task(retries=3, retry_exponential_backoff=True, max_active_tis_per_dag=32)
    def process_tile(spec: dict, tile_id: str) -> str:
        return run_one_tile(RunSpec(**spec), tile_id)   # returns the staged key

    @task
    def validate(spec: dict, keys: list[str]) -> dict:
        return validate_staged_tiles(RunSpec(**spec), keys)

    @task
    def publish(spec: dict, report: dict) -> None:
        if not report["ok"]:
            raise ValueError(f"validation failed: {report['failures']}")
        promote_staging_to_published(RunSpec(**spec))

    spec = build_spec()
    keys = process_tile.partial(spec=spec).expand(tile_id=list_tiles(spec))
    publish(spec, validate(spec, keys))


landcover_inference()

The Prefect form of the same graph is a plain function, which is easier when the tile list depends on run-time state:

from prefect import flow, task
from prefect.futures import wait


@task(retries=3, retry_delay_seconds=[10, 60, 300])
def process_tile(spec: RunSpec, tile_id: str) -> str:
    return run_one_tile(spec, tile_id)


@flow(name="landcover-inference")
def landcover_inference(scene_date: date):
    spec = build_spec(scene_date)
    futures = [process_tile.submit(spec, t) for t in tiles_for_area(spec.grid_id)]
    wait(futures)
    keys = [f.result() for f in futures]
    report = validate_staged_tiles(spec, keys)
    if not report["ok"]:
        raise ValueError(report["failures"])
    promote_staging_to_published(spec)

Step 5 — Gate publication on validation

def validate_staged_tiles(spec: RunSpec, keys: list[str]) -> dict:
    """Coverage, geometry and distribution checks over the staged outputs."""
    failures = []
    expected = set(tiles_for_area(spec.grid_id))
    produced = {k.rsplit("/", 1)[-1].removesuffix(".tif") for k in keys}
    missing = expected - produced
    if missing:
        failures.append(f"{len(missing)} tiles missing, e.g. {sorted(missing)[:3]}")

    for key in keys[:50]:                       # sample, not exhaustive
        with rasterio.open(key) as src:
            if src.crs is None:
                failures.append(f"{key}: no CRS")
            if src.nodata is None:
                failures.append(f"{key}: no nodata value set")
    return {"ok": not failures, "failures": failures, "n_tiles": len(keys)}

Verification & Testing

Orchestration bugs are expensive precisely because they only appear at 3 a.m. under partial failure, so simulate that in tests.

def test_retry_is_idempotent(tmp_path, moto_s3):
    """Running the same tile twice must leave one object with identical bytes."""
    key = "predictions/_staging/t0001.tif"
    write_tile_atomically(arr, profile, "test-bucket", key)
    first = s3.get_object(Bucket="test-bucket", Key=key)["Body"].read()
    write_tile_atomically(arr, profile, "test-bucket", key)
    second = s3.get_object(Bucket="test-bucket", Key=key)["Body"].read()
    assert first == second
    assert not s3.list_objects_v2(Bucket="test-bucket", Prefix=key + ".").get("Contents")


def test_publish_blocked_by_missing_tiles():
    """A run that lost a tile must not publish."""
    report = validate_staged_tiles(spec, keys=produced_keys[:-1])
    assert report["ok"] is False
    assert "missing" in report["failures"][0]


def test_model_version_is_constant_across_tasks():
    """Every staged tile must record the same model version."""
    versions = {read_tag(k, "model_version") for k in produced_keys}
    assert len(versions) == 1, f"run mixed model versions: {versions}"

Also verify the DAG parses in CI (airflow dags list or prefect deployment build --skip-upload). A DAG that fails to import does not raise an alert — it simply stops being scheduled, which is the quietest possible outage.

Why a staged write survives a killed worker Two write strategies are compared under a worker kill. Writing directly to the final key leaves a partially written object that downstream readers will open without error. Writing to a temporary key and copying into place afterwards leaves the final key either absent or complete, never partial. Direct write to the final key Stage, then move worker tile_0001.tif 60% written killed Next stage opens it without error and reads garbage past the cut Corruption propagates into the mosaic and is discovered days later. worker tile_0001.tmp 60% written killed Final key still absent — nothing reads it retry re-runs and lands identical bytes The temp key is garbage-collected by a lifecycle rule on the staging prefix.

Troubleshooting & Common Errors

Tasks succeed but no output appears — the worker wrote to a container-local path rather than object storage. Assert on the URI scheme inside the task and fail fast when it is not s3:// or gs://.

Broken DAG: No module named 'rasterio' — the scheduler is importing a DAG file that imports the geospatial runtime at module level. Move heavy imports inside the task functions so the scheduler only needs the orchestrator’s own dependencies.

A backfill saturates the source API and everything times out — set max_active_runs on the DAG and a pool on the download task. Rate limits are a property of the source, not of your cluster size.

Retries make the run slower than a full restart — retrying without the tile_already_done skip check means every retry redoes finished work. Check-then-skip is what makes tile granularity pay off; the details are in retrying failed tile tasks without reprocessing everything.

Two runs of the same date produce different mosaics — the model reference was resolved per task instead of per run, so a deployment mid-run split the output. Resolve once in build_spec and pass the resolved URI down.

Sensors occupy every worker slot — a poking sensor in mode="poke" holds its slot for the whole wait. Switch to mode="reschedule", or use a deferrable operator, and the same cluster runs an order of magnitude more concurrent waits.

Performance Optimisation

Right-size the task, do not minimise it. One tile per task is the correct granularity for a 10 km tile taking 30–120 seconds. At one second per task the scheduler becomes the bottleneck — batch 20 tiles into one task instead. Measure scheduler overhead per task in your deployment and set the batch size from that number rather than from intuition.

Cap concurrency where the constraint actually lives. The limit is rarely CPU: it is source-API rate limits, object-store request quotas, or model-server throughput. Airflow pools and Prefect concurrency limits let you bound each independently, which is far better than one global worker count.

Do not push arrays through the orchestrator. Airflow XComs and Prefect results are metadata channels. Pass object keys between tasks and let the workers read the data directly; passing a NumPy array through the metadata database will work in testing and fall over in production.

Prefer a long-lived compute session inside a task. Starting a Dask or Ray cluster per tile costs more than the tile. Start it once per run and submit tile work to it — the pattern in scaling batch inference with Dask and Ray — with the orchestrator supervising the run rather than the individual chunks.

Record run metadata as part of the output. Write model_version, run_id and the source scene id into the GeoTIFF tags of every tile. When a consumer reports a strange prediction six weeks later, that metadata turns a forensic exercise into a one-line query, and it feeds directly into the versioning workflow described in spatial dataset versioning with DVC and lakeFS.

FAQ

Should one task process one tile or one whole scene?

One tile, in almost every case. Failure then costs minutes rather than hours, retries are cheap, progress is observable, and the graph parallelises without code changes. Batch tiles together only when scheduler overhead becomes measurable against per-tile runtime.

Why do my retried tasks produce corrupt output?

Because they write straight to the final path. A worker killed mid-write leaves a truncated GeoTIFF that opens without error. Write to a unique temporary key and move it into place after close, so the final key is either absent or complete.

Airflow or Prefect for geospatial work?

Both are fine. Airflow suits fixed schedules, long backfills and teams already running it. Prefect suits graphs whose shape is decided at run time and teams that prefer plain Python. The design principles here — tile granularity, staged writes, a validation gate — matter far more than the choice of tool.

How should model artifacts be pinned in a scheduled run?

Resolve the version once at the start and pass the resolved URI to every task. Letting each task look up “latest” means a deployment mid-run stitches a mosaic from two models, which is very hard to detect after the fact.

Part of: Geospatial MLOps and Model Deployment