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.
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
):
|
||||
|
||||
Reference in New Issue
Block a user