How to Structure a DAG for Raster Processing

from __future__ import annotations

import json
from dataclasses import asdict
from pathlib import Path
from typing import Optional

from prefect import flow, task, get_run_logger
from prefect.concurrency.sync import concurrency


@task(retries=2, timeout_seconds=300)
def inspect(scene_uri: str, dst_crs: str) -> SceneFacts:
    """Stage 1 — headers only. Nothing here reads a pixel."""
    facts = read_scene_facts(scene_uri, dst_crs)     # megapixels, footprint, CRS, bands
    get_run_logger().info(
        "scene inspected: %.0f Mpx, %s → %s", facts.megapixels, facts.src_crs, dst_crs
    )
    return facts


@task(timeout_seconds=300)
def plan(facts: SceneFacts, zoom: int, manifest_path: Path) -> list[TileItem]:
    """Stage 2 — the graph's width, decided from the footprint and recorded."""
    items = [
        TileItem(zoom, t.x, t.y, facts.uri, facts.digest)
        for t in tiles_intersecting(facts.footprint, zoom)     # footprint, not bbox
    ]
    manifest_path.write_text(json.dumps([asdict(i) for i in items]))
    if not items:
        raise ValueError("empty manifest — the footprint and the zoom disagree")
    if len(items) > 20_000:
        raise ValueError(f"{len(items)} tiles is implausible — check the zoom")
    return items


@task(retries=2, retry_delay_seconds=[15, 60], timeout_seconds=900, tags=["tile"])
def process(item: TileItem, scratch: Path) -> Optional[str]:
    """Stage 3 — the only stage that fans out. Idempotent and bounded."""
    dest = scratch / f"{item.z}_{item.x}_{item.y}.tif"
    if ledger_has(item.work_key):
        return str(dest)
    with concurrency("scene-reads", occupy=1):
        warp_window(item, dest)
    ledger_mark(item.work_key)
    return str(dest)


@task(retries=1, timeout_seconds=3600)
def reduce(paths: list[Optional[str]], expected: int, dest: Path) -> Path:
    """Stage 4 — many to one, by reference. Never holds the mosaic in memory."""
    present = [p for p in paths if p]
    coverage = len(present) / expected
    if coverage < 0.995:
        raise RuntimeError(f"coverage {coverage:.3f} below threshold — not assembling")
    build_vrt(present, dest.with_suffix(".vrt"))
    translate(dest.with_suffix(".vrt"), dest)      # one sequential pass
    add_overviews(dest)
    return dest


@task(retries=2, timeout_seconds=1800)
def publish(mosaic: Path, coverage: float, key: str) -> str:
    """Stage 5 — the commit point, with provenance attached to the artefact."""
    upload_with_metadata(mosaic, key, metadata={"tile-coverage": f"{coverage:.4f}"})
    invalidate_downstream(key)
    return key


@flow(name="raster-scene", timeout_seconds=7200)
def raster_scene(scene_uri: str, zoom: int, scratch: Path,
                 dst_crs: str = "EPSG:3857") -> str:
    facts = inspect(scene_uri, dst_crs)
    items = plan(facts, zoom, scratch / "manifest.json")
    states = process.map(items, scratch=scratch, return_state=True)
    paths = [s.result() if s.is_completed() else None for s in states]
    mosaic = reduce(paths, expected=len(items), dest=scratch / "mosaic.tif")
    return publish(mosaic, coverage=sum(p is not None for p in paths) / len(items),
                   key=f"mosaics/{facts.digest}.tif")

from prefect.testing.utilities import prefect_test_harness


def test_reduce_waits_for_every_tile(monkeypatch, tmp_path) -> None:
    order: list[str] = []
    monkeypatch.setattr("mymodule.warp_window", lambda *a: order.append("process"))
    monkeypatch.setattr("mymodule.build_vrt", lambda *a: order.append("reduce"))
    with prefect_test_harness():
        raster_scene("fixtures/scene.tif", zoom=12, scratch=tmp_path)
    assert order.index("reduce") > order.rindex("process")


def test_implausible_fanout_is_refused(tmp_path) -> None:
    facts = SceneFacts(uri="x", footprint=world_footprint(), digest="d")
    with pytest.raises(ValueError, match="implausible"):
        plan.fn(facts, zoom=18, manifest_path=tmp_path / "m.json")


def test_low_coverage_does_not_publish(tmp_path) -> None:
    paths = [str(tmp_path / f"{i}.tif") for i in range(300)] + [None] * 112
    with pytest.raises(RuntimeError, match="coverage"):
        reduce.fn(paths, expected=412, dest=tmp_path / "mosaic.tif")

scene inspected: 8420 Mpx, EPSG:32633 → EPSG:3857
manifest: 412 tiles at z12 (footprint 61% of bbox)
reduce: 412/412 present, coverage 1.0000 → mosaics/sha256-ab12….tif