Cancelling In-Flight Tile Jobs Cleanly

Cancelling a tile job means four things happening in order, and skipping any one of them leaves debris: signal the whole process group so GDAL’s helpers die with it, wait briefly and then escalate to SIGKILL, discard the partial output rather than promoting it, and release the ledger claim so the work is re-drivable. A cancellation that only raises a Python exception has stopped nothing — the warp is still running, its 12 GB is still allocated, and the next task scheduled on that worker will fail for reasons that have nothing to do with it.

When to Use This Pattern

  • A deadline or timeout fires and the task must stop — the normal case, driven by the timeout budget.
  • An operator cancels a run from the orchestrator UI, which sends a signal to the flow process and expects it to propagate.
  • A spot instance is reclaimed, giving a short warning during which an orderly stop is still possible.
  • A circuit breaker opens mid-batch and the remaining tiles should be abandoned rather than attempted.

Complete Working Example

The runner below owns the whole cancellation path. It registers cleanup before starting work, so an interrupt at any point after that finds a handler in place.

from __future__ import annotations

import os
import signal
import subprocess
import tempfile
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator, Optional, Sequence

GRACE_SECONDS = 10.0


class Cancelled(RuntimeError):
    """The job was stopped deliberately — not a failure of the data."""


@contextmanager
def atomic_output(final_path: Path) -> Iterator[Path]:
    """Yield a temporary path on the SAME filesystem; promote it only on success.

    Same filesystem matters: os.replace is atomic within a mount and a copy
    across mounts, and a copy can be interrupted halfway.
    """
    final_path.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp_name = tempfile.mkstemp(dir=final_path.parent, suffix=".part")
    os.close(fd)
    tmp = Path(tmp_name)
    try:
        yield tmp
        os.replace(tmp, final_path)          # the commit point
    finally:
        # Whether we were cancelled, failed, or promoted the file, no .part survives.
        tmp.unlink(missing_ok=True)


def run_cancellable(argv: Sequence[str], timeout: float) -> None:
    """Run a child so that cancelling us actually cancels it."""
    proc = subprocess.Popen(
        argv,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        start_new_session=True,      # its own process group: signals reach helpers
    )
    try:
        _, err = proc.communicate(timeout=timeout)
        if proc.returncode != 0:
            raise RuntimeError(err.decode("utf-8", "replace")[-400:])
    except (subprocess.TimeoutExpired, KeyboardInterrupt, Cancelled):
        _terminate_group(proc)
        raise
    except BaseException:
        # SystemExit and orchestrator cancellations are BaseException, not Exception.
        # Catching only Exception is the most common way a child survives its parent.
        _terminate_group(proc)
        raise


def _terminate_group(proc: subprocess.Popen) -> None:
    if proc.poll() is not None:
        return
    pgid = os.getpgid(proc.pid)
    os.killpg(pgid, signal.SIGTERM)
    try:
        proc.communicate(timeout=GRACE_SECONDS)
    except subprocess.TimeoutExpired:
        os.killpg(pgid, signal.SIGKILL)      # unconditional; no third chance
        proc.communicate()

With those two pieces, a cancellable tile task is short, and every failure path leaves the same clean state:

import psycopg
from prefect import task


@task
def build_tile(tile, dsn: str, timeout: float) -> None:
    final = Path(f"/data/tiles/{tile.z}/{tile.x}/{tile.y}.tif")

    with psycopg.connect(dsn) as conn:
        try:
            with atomic_output(final) as tmp:
                # A lock_timeout as well as a statement_timeout: waiting for a lock
                # is unbounded by default and survives the client going away.
                with conn.cursor() as cur:
                    cur.execute("SET LOCAL lock_timeout = '30s'")
                    cur.execute("SET LOCAL statement_timeout = '120s'")

                run_cancellable(
                    ["gdalwarp", "-t_srs", "EPSG:3857", "-r", "bilinear",
                     "-dstnodata", "-9999", "-co", "COMPRESS=DEFLATE",
                     tile.source, str(tmp)],
                    timeout=timeout,
                )
        finally:
            # Release the claim whatever happened, so the reaper does not have to
            # wait out the stale window for work we know is not being done.
            release_claim(conn, tile.work_key)
Four things a cancellation must clean upThe process group, the partial output file, the database session and the ledger claim each leave a specific kind of debris if they are not handled.obligationskipped, this surviveskill the process groupa gdalwarp holding 12 GB, invisible to the orchestratordiscard the .part filea truncated GeoTIFF that gdalinfo opens happilybound the DB sessiona query holding locks that block tomorrow’s writerelease the claima tile nobody rebuilds until the stale window expires
None of the four failures announce themselves. Each shows up later as something that looks like an unrelated problem — an OOM, a corrupt tile, a lock storm, a missing tile.

The ordering of those four obligations is not arbitrary. The process group goes first because everything else is unsafe while the child is still writing: deleting a .part file that an active gdalwarp still has open frees no space on Linux until the process exits, and releasing the ledger claim while the work is still running invites a second worker to start the same tile. Once the group is confirmed dead, the remaining three are independent and can happen in any order — but they must all happen, which is why they belong in a finally block rather than in the success path.

Order matters only for the first stepThe process group must be terminated and confirmed dead before the partial file is discarded, the database session closed and the ledger claim released. Those last three are independent of each other.1 · group deadSIGTERM → grace → SIGKILL2a · discard .part2b · close DB session2c · release claimre-drivable
Only step one is ordered. The habit of writing the other three as a checklist in finally is what stops the second-most-common cancellation bug: a cleanup that returns early after the first thing it handles.

Parameter & Option Reference

Parameter Type Default Spatial notes
start_new_session bool True Creates a process group. Without it killpg cannot reach GDAL’s helper processes and they outlive the parent.
GRACE_SECONDS float 10.0 Between SIGTERM and SIGKILL. Long enough for GDAL to close a file; measure it if your warps read from object storage.
temp file location same mount required os.replace is atomic within a filesystem only. A .part in /tmp renamed onto a network mount is a copy.
lock_timeout 30 s Lock waits are unbounded by default and are not covered by statement_timeout. Set both.
exception scope BaseException Orchestrator cancellation and SystemExit do not inherit from Exception. Catching Exception lets the child survive.
claim release finally Runs on success, failure and cancellation alike, so the reaper is a backstop rather than the primary path.

Verification & Testing

The test worth writing checks the state of the world after a cancellation, not the exception that was raised.

import time
import psutil


def test_cancellation_leaves_nothing_behind(tmp_path, dsn) -> None:
    final = tmp_path / "12" / "2100" / "1300.tif"
    before = {p.pid for p in psutil.process_iter()}

    with pytest.raises(subprocess.TimeoutExpired):
        with atomic_output(final) as tmp:
            run_cancellable(["sleep", "600"], timeout=0.5)

    time.sleep(GRACE_SECONDS + 1)

    # 1. no orphaned child
    leaked = {p.pid for p in psutil.process_iter()} - before
    assert not any("sleep" in " ".join(psutil.Process(p).cmdline()) for p in leaked)

    # 2. no .part file and no half-written output
    assert not final.exists()
    assert list(final.parent.glob("*.part")) == []


def test_sigkill_escalation_for_a_child_that_ignores_sigterm(tmp_path) -> None:
    # A child that traps SIGTERM must still die.
    script = tmp_path / "stubborn.sh"
    script.write_text("#!/bin/sh\ntrap '' TERM\nsleep 600\n")
    script.chmod(0o755)

    started = time.monotonic()
    with pytest.raises(subprocess.TimeoutExpired):
        run_cancellable([str(script)], timeout=0.5)
    elapsed = time.monotonic() - started
    # It should die at roughly the grace period, not hang forever.
    assert GRACE_SECONDS <= elapsed < GRACE_SECONDS + 5

On a running worker the equivalent check is a one-liner, and it is worth putting in the worker’s start-up script: any GDAL process older than the longest possible timeout is, by definition, an orphan from a cancellation that did not complete.

# Orphaned GDAL processes from previous runs — should always print nothing.
ps -eo pid,etimes,comm | awk '$2 > 1900 && $3 ~ /gdal/ { print }'
Why the signal goes to the groupWithout a new session, killing the gdalwarp process leaves its helper processes running under the original session. With start_new_session, one killpg reaches every descendant.proc.kill() — direct child onlygdalwarp ✗warp thrwarp thrwarp thrthree processes still holding memorykillpg on a new sessiongdalwarp ✗thr ✗thr ✗thr ✗one signal, the whole tree stopsThe left case is the usual explanation for “the worker keeps running out of memory even though every task finished”.
One flag on Popen and one call to killpg is the entire difference. It is invisible in every test that only checks the exception.

There is one more check worth automating rather than remembering. Orphaned processes accumulate slowly and only become visible when a worker starts failing for reasons unrelated to the task that fails — an OOM on a small tile, or a disk that filled with .part files nobody promoted. Running the two checks above at worker start-up, and refusing to accept work if either finds something, converts a slow mystery into an immediate, legible failure. It is three lines of shell and it pays for itself the first time a cancellation path regresses.

Common Pitfalls

  • Catching Exception instead of BaseException. Orchestrator cancellations and SystemExit do not inherit from Exception, so the cleanup handler never runs and the child survives the parent it was supposed to die with.
  • Trusting subprocess.run(timeout=…) to kill anything. It raises TimeoutExpired and leaves the child running. The documentation says so; the behaviour still surprises people every time.
  • Promoting the partial file “so the work is not wasted”. A truncated GeoTIFF is a valid file with an invalid final block. gdalinfo opens it, the mosaic includes it, and the corruption is found weeks later in a derived product.
  • Relying only on the ledger’s stale-claim reaper. The reaper is a backstop for workers that die without warning. When you cancelled deliberately, you know the work stopped — release the claim immediately and save the stale window.
  • Forgetting lock_timeout. statement_timeout does not bound the time spent waiting for a lock. A cancelled task can leave a session queued behind another transaction indefinitely, which looks like a hang in the next run rather than in this one.

Frequently Asked Questions

Does this work on Windows workers?

Not as written — os.killpg and process groups are POSIX. On Windows the equivalent is a job object, which subprocess does not expose directly; the practical approach is psutil.Process(pid).children(recursive=True) and terminating each. Most spatial workers run Linux containers, which is why the POSIX path is the one shown.

What about GDAL used as a Python library rather than a subprocess?

In-process GDAL is much harder to cancel, because the C call does not return until it finishes. rasterio supports a progress callback in some operations that can raise, but coverage is patchy. This is the strongest practical argument for running heavy GDAL work out of process: a subprocess can always be killed, and an in-process call sometimes cannot.

Should a cancelled task be retried?

Not automatically, and not by the same retry policy that handles data failures. A cancellation usually means the deadline was reached, so retrying immediately will hit the same wall. Let the work return to the queue and be picked up by the next run, or by a dead-letter re-drive with a fresh budget.

How much warning does a spot reclamation give?

Typically two minutes on the major clouds, delivered through the instance metadata service. That is enough for an orderly stop if — and only if — the cleanup path is already this short. Poll for the notice in a background thread and set a flag the task loop checks between tiles; the same break that respects a budget respects a reclamation.

Timeout Budgets and Cancellation for Geotasks