From bca28dbf75e217baebffcb379a445b343fe29e9f Mon Sep 17 00:00:00 2001 From: stroblme Date: Mon, 5 Oct 2026 11:23:54 +0200 Subject: [PATCH] Keep a run's lease while its pipeline builds; allow cancelling it then A claimed run only joined the leased set once its pipeline was built, so a build waiting for a busy worker pool got the run swept as abandoned. Cancel now also stops such a run, and a terminal status being overwritten is logged. --- backend/fluksio/flow/runs.py | 44 ++++++++++++---- backend/tests/api/routes/test_runs.py | 73 +++++++++++++++++++++++++++ 2 files changed, 108 insertions(+), 9 deletions(-) diff --git a/backend/fluksio/flow/runs.py b/backend/fluksio/flow/runs.py index f290d5c..2525d2a 100644 --- a/backend/fluksio/flow/runs.py +++ b/backend/fluksio/flow/runs.py @@ -809,8 +809,9 @@ class RunService: self._consumer: threading.Thread | None = None self._keeper: threading.Thread | None = None # Runs this process is driving, and the pipeline each is running, so a - # cancel has something to hold on to. - self._active: dict[str, Pipeline] = {} + # cancel has something to hold on to. None while it is still being + # built — the run is claimed, so its lease is ours to keep all the same. + self._active: dict[str, Pipeline | None] = {} self._cancelled: set[str] = set() self._lock = threading.Lock() @@ -1105,15 +1106,21 @@ class RunService: def cancel(self, run_id: str) -> bool: """Stop a run: kill what it is executing, schedule nothing further.""" with self._lock: - pipeline = self._active.get(run_id) - if pipeline is None: + if run_id not in self._active: # Not running here — if it is still queued, refusing to start # is all the cancelling it needs. - cancelled = self._finish_queued(run_id) + cancelled = self._finish_unstarted(run_id, "queued") if cancelled: self._cancelled.add(run_id) return cancelled self._cancelled.add(run_id) + pipeline = self._active[run_id] + if pipeline is None: + # Claimed, and its build still waiting for a worker: the driver + # checks the mark before it starts anything, and the row says + # so now rather than once a worker frees up. + self._finish_unstarted(run_id, "running") + return True # The gate first, so nothing new is submitted while the running nodes # are being killed; a gated node is never handed to the executor, so # the graph drains instead of going further. @@ -1165,11 +1172,12 @@ class RunService: run = session.get(Run, run_id) return run.flow if run else None - def _finish_queued(self, run_id: str) -> bool: + def _finish_unstarted(self, run_id: str, status: str) -> bool: + """Cancel a run nothing of has executed, if it is still in ``status``.""" with Session(db_engine) as session: result = session.exec( update(Run) - .where(col(Run.id) == run_id, col(Run.status) == "queued") + .where(col(Run.id) == run_id, col(Run.status) == status) .values( status="cancelled", finished_at=datetime.now(UTC), @@ -1382,6 +1390,11 @@ class RunService: ), ) try: + with self._lock: + # Leased from the claim, not from when its pipeline exists: + # building one loads every python node in a worker, and waits + # for one while all of them are busy — for longer than a lease. + self._active[run_id] = None flow = self.controller.store.read_flow(run.flow, draft=run.draft) # Read again, now that it is this run's turn: a sweep queues every # run at once, and the code on disk is free to move in the hours @@ -1406,9 +1419,12 @@ class RunService: run_cache=None if run.no_cache else self._cache, ) with self._lock: - self._active[run_id] = pipeline if run_id in self._cancelled: - pipeline.pause(flow.name) + # Cancelled while it was being built, so nothing of it + # ran; the finally below records it. + status, reason = "cancelled", "Cancelled before it started" + return + self._active[run_id] = pipeline self._publish(run, "run_started") pipeline.run(seed_values(flow, run.params, run.seed)) result = collect_result(flow, state) @@ -1529,6 +1545,16 @@ class RunService: ) -> None: try: with Session(db_engine) as session: + was = session.exec( + select(Run.status).where(col(Run.id) == run_id) + ).first() + if was not in (None, "running", status): + # Swept as abandoned, say, while it was in fact still going. + # What it ended as is the truth, but a client may already + # have taken the first answer as final. + logger.warning( + "Run %s was %s and has now ended %s", run_id, was, status + ) session.exec( update(Run) .where(col(Run.id) == run_id) diff --git a/backend/tests/api/routes/test_runs.py b/backend/tests/api/routes/test_runs.py index 8d0a246..dcf4fc8 100644 --- a/backend/tests/api/routes/test_runs.py +++ b/backend/tests/api/routes/test_runs.py @@ -7,6 +7,7 @@ The pipeline half — what a hit restores and what a key is made of — is in import base64 import csv import json +import threading from datetime import UTC, datetime, timedelta import pytest @@ -19,6 +20,7 @@ from fluksio.flow.controller import FlowController from fluksio.flow.messages import DType, MessageSpec from fluksio.flow.pipeline import NodeOutcome from fluksio.flow.runs import ( + LEASE_STALE_S, OUTPUT_CAP, RunCache, RunRejected, @@ -778,6 +780,77 @@ def test_runs_left_queued_are_woken_by_the_next_engine(): session.commit() +class _BusyWorkers(_OneFlow): + """Builds a run's pipeline the way a full worker pool does: eventually.""" + + events = None + + def __init__(self, flow): + super().__init__(flow) + self.building = threading.Event() + self.free = threading.Event() + + def build_run_pipeline(self, flow, **_kwargs): + self.building.set() + self.free.wait(10) + return object() # never run: the test cancels it first + + +def test_a_run_waiting_for_a_worker_keeps_its_lease_and_can_be_cancelled( + monkeypatch, +): + """Only a run whose pipeline was built had its lease kept. + + Building one loads every python node in a worker, so a run claimed while + all of them were busy was swept as `abandoned` and finished `ok` minutes + later, and a cancel in that window was refused. + """ + controller = _BusyWorkers(FlowDef(name="study", mode="batch")) + service = RunService(controller=controller, queue=_Collect()) + run = service.submit("study", {}) + driver = threading.Thread(target=service._drive, args=(run.id,)) + driver.start() + try: + assert controller.building.wait(5) + # As if it had waited past the cutoff, then one pass of the keeper: + # its lease is renewed before the sweep looks at it. + with Session(db_engine) as session: + session.get(Run, run.id).lease_at = datetime.now(UTC) - timedelta( + seconds=2 * LEASE_STALE_S + ) + session.commit() + monkeypatch.setattr("fluksio.flow.runs.LEASE_INTERVAL_S", 0) + monkeypatch.setattr("fluksio.flow.runs.SWEEP_INTERVAL_S", 0) + sweep = service._sweep_abandoned + + def sweep_once(now): + sweep(now) + service._stop.set() + + monkeypatch.setattr(service, "_sweep_abandoned", sweep_once) + service._keep_leases() + with Session(db_engine) as session: + assert session.get(Run, run.id).status == "running" + + assert service.cancel(run.id) + with Session(db_engine) as session: + assert session.get(Run, run.id).status == "cancelled" + + # A worker frees up: the build returns and the run must not start. + controller.free.set() + driver.join(5) + with Session(db_engine) as session: + stored = session.get(Run, run.id) + assert stored.status == "cancelled" + assert stored.status_reason == "Cancelled before it started" + finally: + controller.free.set() + driver.join(5) + with Session(db_engine) as session: + session.delete(session.get(Run, run.id)) + session.commit() + + def test_a_key_nobody_used_submits_normally( client, superuser_token_headers, monkeypatch ):