Checkpointing and Restarting Long Watershed Pipelines

import hashlib
import json
import logging
import os
import shutil
import tempfile
from contextlib import contextmanager
from pathlib import Path

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


class Checkpoint:
    """
    A content-addressed checkpoint store for expensive pipeline stages.

    Keys are derived from the stage name, the checksums of its input files and
    the parameters it was called with, so a changed input misses automatically.
    """

    def __init__(self, root: str, enabled: bool = True):
        self.root = Path(root)
        self.root.mkdir(parents=True, exist_ok=True)
        self.enabled = enabled
        self.hits = self.misses = 0

    def key(self, stage: str, inputs: list[str], params: dict) -> str:
        h = hashlib.sha256()
        h.update(stage.encode())
        for path in sorted(inputs):
            # Hash the file CONTENT, not its path or mtime. A pipeline that
            # keys on mtime re-runs everything after a copy, and one that keys
            # on path reuses results from a file that has since changed.
            h.update(_file_sha256(path).encode())
        h.update(json.dumps(params, sort_keys=True, default=str).encode())
        return h.hexdigest()[:16]

    def path_for(self, stage: str, key: str, suffix: str = ".tif") -> Path:
        return self.root / f"{stage}__{key}{suffix}"

    def get(self, stage: str, inputs: list[str], params: dict,
            suffix: str = ".tif") -> Path | None:
        if not self.enabled:
            return None
        key = self.key(stage, inputs, params)
        candidate = self.path_for(stage, key, suffix)
        marker = candidate.with_suffix(candidate.suffix + ".done")
        # The marker is written AFTER the payload lands, so a half-written
        # output is never treated as a hit even if the process died mid-write.
        if candidate.exists() and marker.exists():
            self.hits += 1
            log.info("checkpoint HIT  %s (%s) — skipping", stage, key)
            return candidate
        self.misses += 1
        log.info("checkpoint MISS %s (%s) — computing", stage, key)
        return None

    @contextmanager
    def stage(self, stage: str, inputs: list[str], params: dict,
              suffix: str = ".tif"):
        """
        Run a stage under the checkpoint, writing atomically.

        Yields (cached_path_or_None, write_path). When the first is not None
        the body should skip its work entirely.
        """
        cached = self.get(stage, inputs, params, suffix)
        if cached is not None:
            yield cached, None
            return

        key = self.key(stage, inputs, params)
        final = self.path_for(stage, key, suffix)
        # Temporary file in the SAME directory, so the rename is a rename and
        # not a cross-device copy — only a same-filesystem rename is atomic.
        fd, tmp = tempfile.mkstemp(dir=str(self.root), suffix=suffix + ".part")
        os.close(fd)
        try:
            yield None, Path(tmp)
            if not os.path.getsize(tmp):
                raise RuntimeError(f"stage {stage} produced an empty output")
            os.replace(tmp, final)
            final.with_suffix(final.suffix + ".done").write_text(
                json.dumps({"stage": stage, "key": key, "params": params,
                            "inputs": inputs}, default=str))
            log.info("checkpoint SAVE %s (%s) → %s",
                     stage, key, final.name)
        except BaseException:
            # Leave nothing behind on failure: a partial file with no marker is
            # harmless, but cleaning up keeps the directory readable.
            if os.path.exists(tmp):
                os.unlink(tmp)
            raise

    def report(self) -> dict:
        total = self.hits + self.misses
        log.info("checkpoint summary: %d hit(s), %d miss(es) of %d stage(s)",
                 self.hits, self.misses, total)
        if total and self.hits == total:
            log.warning("Every stage was a cache hit — the pipeline produced no "
                        "new work. Is this the run you meant?")
        return {"hits": self.hits, "misses": self.misses}


def _file_sha256(path: str, chunk: int = 1 << 20) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for block in iter(lambda: fh.read(chunk), b""):
            h.update(block)
    return h.hexdigest()


# --- Example usage inside a pipeline ---
# ck = Checkpoint("./checkpoints")
#
# params = {"epsilon": True, "breach_dist": 40}
# with ck.stage("condition", ["raw_dem.tif"], params) as (cached, out):
#     if cached is None:
#         condition_dem("raw_dem.tif", str(out), **params)
#         conditioned = out
#     else:
#         conditioned = cached
#
# with ck.stage("flowdir", [str(conditioned)], {"method": "D8"}) as (cached, out):
#     if cached is None:
#         compute_flowdir(str(conditioned), str(out))
#         fdir = out
#     else:
#         fdir = cached
# ck.report()

INFO checkpoint HIT  condition (a41f9c2e7b3d8105) — skipping
INFO checkpoint HIT  flowdir   (7c02b8ea19f4d663) — skipping
INFO checkpoint MISS accum     (5e8d1a04c9b27f3a) — computing
INFO checkpoint SAVE accum     (5e8d1a04c9b27f3a) → accum__5e8d1a04c9b27f3a.tif
INFO checkpoint MISS delineate (b93c705fa1284de6) — computing
INFO checkpoint summary: 2 hit(s), 2 miss(es) of 4 stage(s)