Watershed Modeling Python for hydrology automation

Building a Prefect Flow for DEM-to-Watershed Automation

Prefect 2.x turns the DEM-to-watershed chain into a handful of @task functions composed inside one @flow, with retries, caching, and concurrency configured as decorator arguments rather than hand-rolled control flow. This guide builds that flow concretely, and sits under the pipeline orchestration topic within the broader Watershed Delineation & Catchment Synchronization section. Where the parent guide compares orchestrators and lays out the nine-stage graph, this page is Prefect-specific: exactly which decorator arguments to set, why a cache key must hash file contents, and how .map() fans one flow across many pour points.

The payoff of doing this in Prefect rather than a bash loop is that a re-run costs almost nothing. Deterministic stages return their cached result instantly, transient network failures retry themselves, and a batch of basins runs concurrently under a single task runner with per-basin failure isolation baked in.

Prerequisites

This page assumes you have a working geospatial environment and a conditioned understanding of the individual stages. If depression filling or routing is unfamiliar, build those first — the flow here calls into the same operations documented in D8 flow direction implementation.

  • Python 3.10 or newer, with prefect>=2.14 installed alongside a conda-forge geospatial base (richdem, rasterio, pysheds, geopandas).
  • A staged DEM per basin on local or object storage. The acquisition task validates its presence; swap in a real 3DEP or Copernicus request for production.
  • A manifest of pour points as (x, y) coordinate pairs in the DEM’s projected CRS.
bash 
pip install "prefect>=2.14"
prefect version    # confirm the 2.x line before running the flow below

Flow Structure: @flow and @task

A Prefect flow is a normal Python function decorated with @flow; the tasks it calls are functions decorated with @task. Dependencies are never declared explicitly — Prefect infers the graph from how return values thread from one task into the next. That is why keeping each stage a separate task matters: it is the unit Prefect retries, caches, and schedules.

Prefect Flow with Mapped Delineation Fan-Out A single conditioning chain runs left to right: acquire DEM, then reproject with a cache key, then fill and route with a cache key. The routed output feeds a delineate task that is mapped with dot map over three pour points, producing three concurrent catchment tasks that each feed a shared validate task. Annotations mark retries on acquire and reproject, caching on reproject and fill-and-route, and the ConcurrentTaskRunner covering the mapped branch. acquire_dem @task retries=3 reproject + fill_route cache_key_fn delineate .map(points) fan-out catchment[0] catchment[1] catchment[2] ConcurrentTaskRunner

The linear head of the flow conditions the DEM once per basin; the .map() fan-out at the tail delineates many pour points against that single conditioned grid concurrently. Sharing the conditioning stage across pour points is the efficiency win — you fill and route the DEM once, then delineate ten outlets from it without re-routing.

Task Caching: cache_key_fn and cache_expiration

Prefect caches a task’s result when you supply a cache_key_fn. On each call it computes a key; if a completed run with that key exists and has not expired, Prefect returns the stored result instead of re-executing. The default helper, task_input_hash, keys on the argument values — which for a file path means the string, not the file. That is a trap in hydrology: a DEM tile re-downloaded or corrected keeps its path but changes its bytes, and a path-based key would serve the stale prior result.

The fix is a cache key that hashes the input file’s contents:

python 
import hashlib
import logging
from pathlib import Path
from datetime import timedelta

from prefect import task, get_run_logger
from prefect.tasks import task_input_hash

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")


def content_cache_key(context, arguments) -> str:
    """Cache key from the first existing file argument's SHA-256, so the key changes
    whenever the DEM bytes change. Falls back to the default argument hash otherwise."""
    for value in arguments.values():
        candidate = Path(str(value))
        if candidate.exists() and candidate.is_file():
            digest = hashlib.sha256(candidate.read_bytes()).hexdigest()[:16]
            return f"{context.task.name}-{digest}"
    return task_input_hash(context, arguments)


@task(cache_key_fn=content_cache_key, cache_expiration=timedelta(days=14))
def fill_and_route(proj_dem_path: str) -> dict:
    """Deterministic conditioning + routing. Cached on DEM content for 14 days."""
    logger = get_run_logger()
    logger.info("fill_and_route cache miss for %s; computing", proj_dem_path)
    # RichDEM fill + ResolveFlats + FlowDirection + FlowAccumulation go here.
    return {"flow_dir": proj_dem_path + ".dir", "flow_acc": proj_dem_path + ".acc"}

cache_expiration bounds how long a cached result stays valid — 14 days here means the conditioned grid is trusted for two weeks before recomputation, a sensible ceiling for DEMs that update on quarterly release cycles. For a large content hash, reading the whole file on every call is itself a cost; for multi-gigabyte tiles, hash a cheap proxy instead (file size plus modification time, or a stored checksum sidecar) rather than the full bytes.

Retries and Backoff: retries and retry_delay_seconds

Attach retries to the tasks that fail transiently — acquisition and reprojection reach across a network — and leave the deterministic compute tasks without them so a genuine data error surfaces immediately.

python 
from prefect import task, get_run_logger


@task(retries=3, retry_delay_seconds=[5, 30, 120])
def acquire_dem(huc_id: str, work_dir: str) -> str:
    """Network-bound fetch. Three retries with escalating backoff of 5, 30, 120 s."""
    logger = get_run_logger()
    logger.info("Fetching DEM for %s", huc_id)
    # A real 3DEP/Copernicus request goes here; a transient 503 triggers a retry.
    return f"{work_dir}/{huc_id}/dem_raw.tif"

Passing a list to retry_delay_seconds gives per-attempt backoff — 5 seconds after the first failure, 30 after the second, 120 after the third — which spaces retries out enough to ride through a rate limit or a mirror hiccup without hammering the server. A single integer applies the same delay to every attempt. Reserve retries for I/O; a RichDEM call that raises on a malformed array will raise identically three times and only delay the real signal.

Fanning Out with .map()

.map() calls a task once per element of an iterable, submitting them all to the task runner so they execute concurrently. It is how one flow delineates many pour points against a single conditioned DEM.

python 
import logging
from prefect import flow, task, get_run_logger
from prefect.task_runners import ConcurrentTaskRunner

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")


@task
def delineate_at(routed: dict, outlet_xy: tuple) -> str:
    """Snap one pour point and delineate its catchment from the shared routed grids."""
    logger = get_run_logger()
    logger.info("Delineating catchment at %s", outlet_xy)
    # pysheds snap + catchment + polygonize go here; return the output path.
    return f"catchment_{outlet_xy[0]:.0f}_{outlet_xy[1]:.0f}.gpkg"


@flow(name="prefect-dem-to-watershed", task_runner=ConcurrentTaskRunner())
def multi_outlet_flow(huc_id: str, outlets: list[tuple], work_dir: str = "./work") -> list[str]:
    logger = get_run_logger()
    dem_raw = acquire_dem(huc_id, work_dir)
    routed = fill_and_route(dem_raw)          # conditioned once per basin
    # .map() fans delineation across every outlet concurrently:
    catchments = delineate_at.map(routed=routed, outlet_xy=outlets)
    paths = [c.result() for c in catchments]
    logger.info("Delineated %d catchments for HUC %s", len(paths), huc_id)
    return paths


if __name__ == "__main__":
    multi_outlet_flow(
        huc_id="03010101",
        outlets=[(1_450_000.0, 1_920_000.0),
                 (1_462_500.0, 1_905_000.0),
                 (1_478_000.0, 1_888_500.0)],
    )

When you .map() a scalar argument like routed alongside an iterable like outlets, Prefect broadcasts the scalar to every mapped call — every delineation sees the same conditioned grid, each with its own pour point. The ConcurrentTaskRunner on the flow decorator is what makes the mapped tasks overlap rather than run serially.

Parameter Reference

Parameter Where it goes Effect
retries @task Number of automatic re-executions on failure; set on network tasks only
retry_delay_seconds @task Scalar or per-attempt list of backoff waits between retries
cache_key_fn @task Function returning a cache key; hash file contents, not the path
cache_expiration @task timedelta after which a cached result is recomputed
persist_result @task / @flow Persists the return value so caching survives across processes
task_runner @flow ConcurrentTaskRunner for I/O concurrency; DaskTaskRunner for multi-process

Gotchas

  • Return serializable handles, not arrays. Prefect persists a task’s return value to back its cache. Return file paths or small dictionaries of paths, never a multi-hundred-megabyte NumPy array — the heavy raster lives on disk as a GeoTIFF, and the cached value is just the pointer to it. Returning the array itself bloats the result store and can fail serialization outright.

  • Cache keys on file inputs must hash content. As shown above, task_input_hash on a path is stale-prone. If a corrected DEM reuses its filename, only a content hash invalidates the cache. This is the single most common Prefect caching bug in geospatial pipelines.

  • ConcurrentTaskRunner shares one process. It overlaps I/O well and lets GIL-releasing RichDEM C++ calls run in parallel across threads, but every concurrent DEM holds a full array resident. Memory, not CPU, is the ceiling — cap the number of mapped tasks in flight to available RAM, or move to DaskTaskRunner for true multi-process scaling across machines.

  • .result() inside a flow forces resolution. Calling .result() on mapped futures blocks until they finish, which is fine at the end of a flow but serializes execution if done mid-chain. Pass futures directly between tasks and let Prefect resolve them lazily wherever possible.

Frequently Asked Questions

Why must a Prefect cache_key_fn hash file contents instead of the path?

Prefect’s default task_input_hash keys on the argument values, so two runs with the same file path return the cached result even after the file on disk changes. A DEM tile that is re-downloaded or corrected keeps the same path but different bytes, so a path-based key serves a stale result. Hashing the file’s bytes makes the cache key change whenever the input content changes, which is exactly the behaviour a reproducible pipeline needs.

Does ConcurrentTaskRunner give real parallelism for RichDEM stages?

ConcurrentTaskRunner runs tasks on an asyncio event loop backed by a thread pool, which overlaps I/O-bound work well but shares a single process. CPU-bound RichDEM calls release the GIL inside their C++ core, so they do parallelize across threads, but memory is the real constraint: each concurrent DEM holds a full array resident. Cap concurrency to available RAM, or switch to DaskTaskRunner for multi-process and multi-machine scaling.

What results can Prefect cache and persist between runs?

Prefect persists a task’s return value when persist_result is enabled and the value is serializable. Return small, serializable handles — file paths as strings or dictionaries of paths — rather than large in-memory NumPy arrays. The heavy artifacts live on disk as GeoTIFFs; the cached return value is just the path that points to them, which keeps the result store small and serialization reliable.