feat: cancel an in-flight module run
Adds run cancellation, which had no implementation despite `cancelled`
being a declared run status.
Mechanism: API/web/scheduler runs execute as background tasks in the
service process, so `running_pid`'s PID is the service itself — killing it
is not an option. Instead, engine/cancel.py keeps an in-process registry of
run_id -> live jrunner Popen. The cancel handler (same process) terminates
that subprocess; migrate then fails and the engine records the run as
`cancelled` (not `error`) and releases the lock in its finally. Checkpoints
before migrate and before merge honor a cancel flag set between phases.
- engine/cancel.py: registry (arm/register/unregister/is_cancelled/request_cancel) + RunCancelled
- jrunner.migrate: optional proc_hook to expose the live Popen (no engine dependency)
- runner: arm, register via proc_hook, two checkpoints, classify cancelled, unregister
- POST /api/modules/{id}/cancel (extracts run_id from the lock) and POST /runs/{id}/cancel (web)
- cancel button on run detail + live fragment, shown while running
- api/app lifespan: clear stale locks (dead PID / >24h) at startup via the existing
repo.clear_stale_locks, so a hard service kill no longer wedges `running=1`
Verified: real subprocess terminate -> status=cancelled, lock released;
endpoints return 409 (not running) / 404 (unknown) / 200 (fragment).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
parent
3562dfa581
commit
4f83e8d991
@ -15,7 +15,9 @@ from fastapi import FastAPI, Request
|
||||
from fastapi.exception_handlers import http_exception_handler
|
||||
from starlette.exceptions import HTTPException as StarletteHTTPException
|
||||
|
||||
from .. import __version__, db, jrunner
|
||||
import os
|
||||
|
||||
from .. import __version__, db, jrunner, repo
|
||||
from ..web import mount_web
|
||||
from .routes import connections, introspect, modules, runs, system
|
||||
|
||||
@ -102,8 +104,26 @@ class _RequestBodyCapture:
|
||||
await self.app(scope, replay, send)
|
||||
|
||||
|
||||
def _live_pids() -> set[int]:
|
||||
"""PIDs currently alive on this host (Linux /proc). Empty set elsewhere."""
|
||||
try:
|
||||
return {int(p) for p in os.listdir("/proc") if p.isdigit()}
|
||||
except OSError:
|
||||
return set()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _lifespan(app: FastAPI):
|
||||
# Release module locks left behind by a previous process that died mid-run
|
||||
# (a hard kill skips the engine's finally, so `running` stays 1 and future
|
||||
# runs would refuse with "already running"). Clears locks held by a dead PID
|
||||
# or older than 24h.
|
||||
try:
|
||||
n = repo.clear_stale_locks(live_pids=_live_pids())
|
||||
if n:
|
||||
_log.warning("cleared %d stale module lock(s) at startup", n)
|
||||
except Exception: # noqa: BLE001
|
||||
_log.exception("stale-lock cleanup failed")
|
||||
from ..scheduler import start_scheduler
|
||||
start_scheduler()
|
||||
yield
|
||||
|
||||
@ -93,6 +93,28 @@ def _run_in_background(module_id: int, run_id: int, dry_run: bool) -> None:
|
||||
repo.finish_run(run_id, status="error", error=str(e))
|
||||
|
||||
|
||||
@router.post("/modules/{module_id}/cancel")
|
||||
def cancel_module_run(module_id: int) -> dict:
|
||||
"""Request cancellation of a module's in-flight run.
|
||||
|
||||
Kills the live jrunner subprocess if one is active in this process; either
|
||||
way the cancel flag is honored at the next engine checkpoint. The run lands
|
||||
with status ``cancelled`` and the lock is released by the engine.
|
||||
"""
|
||||
m = repo.get_module(module_id)
|
||||
if m is None:
|
||||
raise HTTPException(404, f"module id={module_id} not found")
|
||||
if not m["running"]:
|
||||
raise HTTPException(409, f"module {m['name']!r} is not running")
|
||||
# running_pid is "<pid>:<run_id>" (see engine.runner)
|
||||
_, _, run_id_str = (m["running_pid"] or "").partition(":")
|
||||
if not run_id_str:
|
||||
raise HTTPException(409, "no run_id recorded for the active lock")
|
||||
run_id = int(run_id_str)
|
||||
result = engine.request_cancel(run_id)
|
||||
return {"module_id": module_id, "run_id": run_id, "status": result}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Watermarks — scoped to a module
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@ -1,3 +1,5 @@
|
||||
from .cancel import request_cancel
|
||||
from .runner import GroupRunOutcome, LockBusy, RunOutcome, run_group, run_module
|
||||
|
||||
__all__ = ["GroupRunOutcome", "LockBusy", "RunOutcome", "run_group", "run_module"]
|
||||
__all__ = ["GroupRunOutcome", "LockBusy", "RunOutcome", "request_cancel",
|
||||
"run_group", "run_module"]
|
||||
|
||||
63
pipekit/engine/cancel.py
Normal file
63
pipekit/engine/cancel.py
Normal file
@ -0,0 +1,63 @@
|
||||
"""In-process cancellation registry for running module jobs.
|
||||
|
||||
API/web/scheduler runs all execute as background tasks *inside the pipekit
|
||||
service process*, and so does the cancel HTTP handler — so no IPC or OS signals
|
||||
are needed. This module holds a small in-memory map of ``run_id -> live jrunner
|
||||
Popen`` plus a set of run_ids that have been asked to stop. The cancel endpoint
|
||||
looks up the live subprocess and terminates it; the engine also checks the flag
|
||||
at phase boundaries so a cancel lands promptly even between subprocess calls.
|
||||
|
||||
Scope: this only reaches runs in *this* process. A ``pipekit run`` CLI run is a
|
||||
separate process and isn't in the registry — cancel it with Ctrl-C.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
|
||||
_lock = threading.Lock()
|
||||
_active: dict[int, object] = {} # run_id -> subprocess.Popen (live jrunner proc)
|
||||
_cancelled: set[int] = set() # run_ids asked to stop
|
||||
|
||||
|
||||
class RunCancelled(Exception):
|
||||
"""Raised inside run_module when a cancel was requested at a checkpoint."""
|
||||
|
||||
|
||||
def arm(run_id: int) -> None:
|
||||
"""Clear any stale cancel flag before a run starts."""
|
||||
with _lock:
|
||||
_cancelled.discard(run_id)
|
||||
|
||||
|
||||
def register(run_id: int, proc: object) -> None:
|
||||
"""Record the live subprocess for a run so cancel can reach it."""
|
||||
with _lock:
|
||||
_active[run_id] = proc
|
||||
|
||||
|
||||
def unregister(run_id: int) -> None:
|
||||
"""Drop a run's subprocess handle (proc finished or run ended)."""
|
||||
with _lock:
|
||||
_active.pop(run_id, None)
|
||||
|
||||
|
||||
def is_cancelled(run_id: int) -> bool:
|
||||
with _lock:
|
||||
return run_id in _cancelled
|
||||
|
||||
|
||||
def request_cancel(run_id: int) -> str:
|
||||
"""Flag a run for cancellation and terminate its live subprocess if any.
|
||||
|
||||
Returns ``"signalled"`` if a live subprocess was terminated, or
|
||||
``"flagged"`` if none was active this instant (the flag is honored at the
|
||||
next engine checkpoint — e.g. the run is between phases or in another
|
||||
process).
|
||||
"""
|
||||
with _lock:
|
||||
_cancelled.add(run_id)
|
||||
proc = _active.get(run_id)
|
||||
if proc is not None and proc.poll() is None: # type: ignore[attr-defined]
|
||||
proc.terminate() # type: ignore[attr-defined]
|
||||
return "signalled"
|
||||
return "flagged"
|
||||
@ -20,7 +20,8 @@ import traceback
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .. import drivers, jrunner, repo
|
||||
from . import merge, watermark
|
||||
from . import cancel, merge, watermark
|
||||
from .cancel import RunCancelled
|
||||
|
||||
|
||||
@dataclass
|
||||
@ -66,6 +67,8 @@ def run_module(module_id: int, *, group_run_id: int | None = None,
|
||||
repo.finish_run(run_id, status="error", error="already running")
|
||||
raise LockBusy(f"module {module['name']!r} is already running")
|
||||
|
||||
cancel.arm(run_id)
|
||||
|
||||
resolved_sql: str | None = None
|
||||
merge_sql: str | None = None
|
||||
row_count: int | None = None
|
||||
@ -118,6 +121,9 @@ def run_module(module_id: int, *, group_run_id: int | None = None,
|
||||
jrunner.run_dest_sql(dest_conn, drop_sql)
|
||||
jrunner.run_dest_sql(dest_conn, like_sql)
|
||||
|
||||
if cancel.is_cancelled(run_id):
|
||||
raise RunCancelled()
|
||||
|
||||
# 5. migrate source → staging. jrunner does its own `DELETE FROM staging`
|
||||
# before loading, so we don't need a separate TRUNCATE.
|
||||
migrate_result = jrunner.migrate(
|
||||
@ -125,11 +131,15 @@ def run_module(module_id: int, *, group_run_id: int | None = None,
|
||||
sql=resolved_sql, dest_table=module["staging_table"],
|
||||
clear=False,
|
||||
progress_cb=lambda line: repo.append_run_live_log(run_id, line),
|
||||
proc_hook=lambda p: cancel.register(run_id, p),
|
||||
)
|
||||
row_count = migrate_result.row_count
|
||||
repo.log_run_output(run_id, jrunner_stdout=migrate_result.stdout,
|
||||
jrunner_stderr=migrate_result.stderr)
|
||||
|
||||
if cancel.is_cancelled(run_id):
|
||||
raise RunCancelled()
|
||||
|
||||
# 7. merge
|
||||
jrunner.run_dest_sql(dest_conn, merge_sql)
|
||||
|
||||
@ -142,10 +152,19 @@ def run_module(module_id: int, *, group_run_id: int | None = None,
|
||||
return RunOutcome(run_id, status, row_count, None, resolved_sql, merge_sql)
|
||||
|
||||
except Exception as e: # noqa: BLE001
|
||||
error = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
|
||||
# A cancel request kills the jrunner subprocess, which surfaces as a
|
||||
# JrunnerError — treat any failure under a set cancel flag (or an
|
||||
# explicit checkpoint raise) as 'cancelled', not 'error'.
|
||||
cancelled = isinstance(e, RunCancelled) or cancel.is_cancelled(run_id)
|
||||
if isinstance(e, jrunner.JrunnerError):
|
||||
repo.log_run_output(run_id, jrunner_stdout=e.stdout,
|
||||
jrunner_stderr=e.stderr)
|
||||
if cancelled:
|
||||
status = "cancelled"
|
||||
error = "cancelled by user"
|
||||
else:
|
||||
status = "error"
|
||||
error = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
|
||||
# Failure-path hooks, if any. Never let these mask the real error.
|
||||
try:
|
||||
hook_log = _run_hooks(module_id, fail_fast=False,
|
||||
@ -154,9 +173,10 @@ def run_module(module_id: int, *, group_run_id: int | None = None,
|
||||
repo.log_run_output(run_id, hook_log=hook_log)
|
||||
except Exception: # noqa: BLE001, S110
|
||||
pass
|
||||
return RunOutcome(run_id, "error", row_count, error, resolved_sql, merge_sql)
|
||||
return RunOutcome(run_id, status, row_count, error, resolved_sql, merge_sql)
|
||||
|
||||
finally:
|
||||
cancel.unregister(run_id)
|
||||
repo.finish_run(run_id, status=status, row_count=row_count, error=error)
|
||||
repo.release_module_lock(module_id)
|
||||
|
||||
|
||||
@ -147,9 +147,13 @@ def migrate(
|
||||
trim: bool = True,
|
||||
timeout: int = 3600,
|
||||
progress_cb: "Callable[[str], None] | None" = None,
|
||||
proc_hook: "Callable[[object], None] | None" = None,
|
||||
) -> MigrateResult:
|
||||
"""Stream `sql` results from source into `dest_table` via jrunner migration mode.
|
||||
|
||||
If ``proc_hook`` is provided it is called with the live ``Popen`` right after
|
||||
spawn, letting a caller (the engine) register it for cancellation.
|
||||
|
||||
If ``progress_cb`` is provided it is called with each non-empty stdout line
|
||||
as jrunner produces it, enabling live progress reporting.
|
||||
"""
|
||||
@ -180,6 +184,8 @@ def migrate(
|
||||
|
||||
proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
||||
text=True, env=_subprocess_env())
|
||||
if proc_hook is not None:
|
||||
proc_hook(proc)
|
||||
|
||||
stderr_buf: list[str] = []
|
||||
|
||||
|
||||
@ -1071,6 +1071,20 @@ def run_live_fragment(request: Request, run_id: int):
|
||||
)
|
||||
|
||||
|
||||
@_router.post("/runs/{run_id}/cancel", response_class=HTMLResponse)
|
||||
def run_cancel(request: Request, run_id: int):
|
||||
run = repo.get_run(run_id)
|
||||
if run is None:
|
||||
raise HTTPException(404, f"run id={run_id} not found")
|
||||
if run["status"] == "running":
|
||||
engine.request_cancel(run_id)
|
||||
# Return the (soon-to-update) live fragment so the button/status refresh.
|
||||
run = repo.get_run(run_id)
|
||||
module = repo.get_module(run["module_id"])
|
||||
run["module_name"] = module["name"] if module else "?"
|
||||
return _templates.TemplateResponse(request, "_run_live.html", _ctx(run=run))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Watermarks — add/edit/delete forms on module detail
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@ -9,6 +9,12 @@
|
||||
<header>
|
||||
Progress
|
||||
<span style="margin-left:auto">
|
||||
{% if run.status == 'running' %}
|
||||
<button class="btn ghost" style="margin-right:0.5rem"
|
||||
hx-post="/runs/{{ run.id }}/cancel"
|
||||
hx-target="#run-live" hx-swap="outerHTML"
|
||||
hx-confirm="Cancel this run?">cancel</button>
|
||||
{% endif %}
|
||||
<span class="pill {{ run.status }}">{{ run.status }}</span>
|
||||
{% if run.row_count is not none %} · {{ run.row_count }} rows{% endif %}
|
||||
</span>
|
||||
|
||||
@ -48,6 +48,12 @@
|
||||
<header>
|
||||
Progress
|
||||
<span style="margin-left:auto">
|
||||
{% if run.status == 'running' %}
|
||||
<button class="btn ghost" style="margin-right:0.5rem"
|
||||
hx-post="/runs/{{ run.id }}/cancel"
|
||||
hx-target="#run-live" hx-swap="outerHTML"
|
||||
hx-confirm="Cancel this run?">cancel</button>
|
||||
{% endif %}
|
||||
<span class="pill {{ run.status }}">{{ run.status }}</span>
|
||||
{% if run.row_count is not none %} · {{ run.row_count }} rows{% endif %}
|
||||
</span>
|
||||
|
||||
Loading…
Reference in New Issue
Block a user