Airflow DAG Templates for Batch Watershed Delineation

Apache Airflow’s TaskFlow API expresses the DEM-to-watershed chain as a scheduled DAG that fans across a list of hydrologic unit codes at runtime, making it the natural fit when watershed delineation runs on a cadence across many basins rather than on demand for one. This template guide sits under pipeline orchestration, part of the Watershed Delineation & Catchment Synchronization section, and focuses on the Airflow-specific machinery: a @dag schedule, dynamic task mapping with .expand, XCom that carries paths rather than payloads, and a @task.sensor that blocks the batch until the input DEM has landed.

Airflow’s strength for this workload is operational. Every mapped basin appears as its own task instance in the grid view, retryable and inspectable in isolation, with batch-wide concurrency governed by max_active_tasks and pool slots. That is the scheduling model an operations team already running Airflow will reach for; the contrast with Prefect’s in-flow model is spelled out in the sibling guide.

Prerequisites

• Apache Airflow 2.8 or newer for the mature TaskFlow API and dynamic task mapping, with the geospatial stack (richdem, rasterio, pysheds, geopandas) installed in the same environment the workers run under.
• A metadata database (Postgres in production; SQLite only for local trials) — every XCom value lands here, which is exactly why you pass paths, not arrays.
• Shared storage reachable by every worker — a mounted volume or object store — so a path written by one task is readable by the next regardless of which worker picks it up.

pip install "apache-airflow>=2.8"
airflow db migrate       # initialise the metadata database before first run

DAG Structure and Scheduling

The @dag decorator turns a function into a DAG; @task functions inside it become nodes, and returning a value from one into another wires the dependency and routes the value through XCom. Scheduling parameters live on the decorator.

The sensor gates the whole batch: nothing downstream runs until the input DEM is present. Once it succeeds, list_basins returns the HUC list and the per-basin pipeline expands into one task instance per basin, all subject to the concurrency cap.

Waiting for Input with @task.sensor

A scheduled batch should not begin the moment the clock fires — it should begin when the day’s DEM has actually arrived on shared storage. The @task.sensor decorator turns a function into a deferrable sensor that polls until its PokeReturnValue reports done.

import logging
from pathlib import Path
from airflow.decorators import task
from airflow.sensors.base import PokeReturnValue

logging.getLogger("airflow.task").setLevel(logging.INFO)


@task.sensor(poke_interval=300, timeout=60 * 60 * 6, mode="reschedule")
def wait_for_dem(dem_dir: str) -> PokeReturnValue:
    """Poll shared storage every 5 minutes (up to 6 hours) for the day's DEM drop.
    mode='reschedule' frees the worker slot between pokes instead of holding it."""
    log = logging.getLogger("airflow.task")
    ready_flag = Path(dem_dir) / "_SUCCESS"
    is_ready = ready_flag.exists()
    log.info("Poking for DEM readiness flag %s: %s", ready_flag, is_ready)
    return PokeReturnValue(is_done=is_ready, xcom_value=str(dem_dir))

mode="reschedule" is the important choice: it releases the worker slot between pokes so a six-hour wait does not occupy a slot the whole time, which matters when many DAGs share a worker pool. The sensor returns the DEM directory as its XCom value so downstream tasks receive the confirmed path.

Dynamic Task Mapping with .expand

.expand generates one mapped task instance per element of a list produced at runtime — the batch size is not fixed in the DAG definition but discovered when list_basins runs. Each mapped instance is independently scheduled, retried, and shown in the grid view.

import logging
import pendulum
from airflow.decorators import dag, task

logging.getLogger("airflow.task").setLevel(logging.INFO)


@dag(
    schedule="@daily",
    start_date=pendulum.datetime(2026, 1, 1, tz="UTC"),
    catchup=False,
    max_active_tasks=4,          # cap concurrent RichDEM stages to fit worker RAM
    default_args={"retries": 2, "retry_delay": pendulum.duration(minutes=5)},
    tags=["hydrology", "delineation"],
)
def batch_watershed_delineation():
    """Scheduled batch: wait for the DEM drop, list basins, delineate each, collect."""

    @task
    def list_basins(dem_dir: str) -> list[dict]:
        """Read the day's basin manifest and return one dict per HUC to map over."""
        log = logging.getLogger("airflow.task")
        basins = [
            {"huc_id": "03010101", "dem_dir": dem_dir, "outlet_xy": [1_450_000.0, 1_920_000.0]},
            {"huc_id": "03010102", "dem_dir": dem_dir, "outlet_xy": [1_462_500.0, 1_905_000.0]},
            {"huc_id": "03010103", "dem_dir": dem_dir, "outlet_xy": [1_478_000.0, 1_888_500.0]},
        ]
        log.info("Discovered %d basins for this run", len(basins))
        return basins

    @task
    def delineate_basin(params: dict) -> str:
        """Full per-basin pipeline. Reads the DEM path from params, writes a catchment
        file to shared storage, and returns ONLY its path via XCom (never the raster)."""
        log = logging.getLogger("airflow.task")
        huc_id = params["huc_id"]
        dem_path = f"{params['dem_dir']}/{huc_id}/dem_raw.tif"
        # fill + route (RichDEM) and snap + delineate (pysheds) run here on dem_path.
        out_path = f"{params['dem_dir']}/{huc_id}/catchment.gpkg"
        log.info("Delineated %s -> %s", huc_id, out_path)
        return out_path      # a path string: small, XCom-safe

    @task
    def collect_results(catchment_paths: list[str]) -> None:
        """Gather the mapped outputs (paths) for downstream cataloguing."""
        log = logging.getLogger("airflow.task")
        log.info("Batch produced %d catchments", len(catchment_paths))

    dem_dir = wait_for_dem(dem_dir="/data/dem_drop")
    basins = list_basins(dem_dir)
    catchments = delineate_basin.expand(params=basins)   # one instance per basin
    collect_results(catchments)


batch_watershed_delineation()

delineate_basin.expand(params=basins) is the whole mechanism: Airflow reads the list list_basins produced and materializes a mapped instance per element, each running delineate_basin with one basin’s parameters. collect_results receives the list of returned paths automatically once every mapped instance finishes. The per-basin body condenses the same fill, route, snap, and delineate operations detailed in the parent orchestration guide.

Parameter Reference

Gotchas

• XCom carries paths, not arrays. Every value returned from a task is serialized into the metadata database, which is sized for small control-plane payloads. Returning a NumPy raster bloats the database and can exceed the backend’s XCom size limit outright. Write the raster to shared storage and return only its path string, as every task above does; the next task opens the file itself.

• Tasks must be idempotent for retries and backfill. Airflow retries a failed mapped instance and may re-run a scheduling interval. A task that appends to a shared file or writes to a timestamp-only path duplicates records on rerun. Key every output on the HUC identifier and write to a deterministic path so a retry overwrites cleanly rather than accumulating.

• Executor choice governs memory contention. The LocalExecutor runs tasks as subprocesses on one node — fine for a powerful, high-RAM machine and a modest basin count. The CeleryExecutor or KubernetesExecutor distribute mapped instances across workers for large batches. Whichever you pick, RichDEM’s resident arrays are the constraint: cap max_active_tasks and use pools so concurrent DEMs never exhaust a worker.

• catchup=False on first deploy. A DAG with a past start_date and catchup=True immediately backfills every missed interval, launching a flood of batch runs. Set catchup=False unless you genuinely want the history reprocessed.

Related Topics

• Pipeline Orchestration: DEM to Watershed Automation — the parent guide laying out the full nine-stage task graph and orchestration guarantees
• Building a Prefect flow for DEM-to-watershed automation — the same pipeline in Prefect 2.x, for a direct contrast in scheduling model, caching, and fan-out mechanics