tao.chen/model-platform
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: P0-4 schedule node janitor + runtime reaper/lock invariants

Browse Source 
S2: schedule worker add janitor task that force-terminals node_runs
whose deadline (timeout_seconds + retry_count*retry_interval + 120s
slack from started_at) has passed. Closes the gap where outbox retry
exhaustion (5 tries, capped 30s backoff) marked the *event* failed
but left the *node_run* stuck in queued/running forever. Re-reads the
row under FOR UPDATE before writing so a worker that races us to a
real terminal state is not overwritten; idempotency key uses
:timed_out variant so the :finished path cannot collide.

R1: extract _reap_once() from _reap_loop for testability; in the
dead-process branch, re-verify (process.pid, started_at) against the
live JUPYTER_PROCESSES entry before del. A start_workspace that
replaced the dead record mid-cycle used to have its new entry
silently erased by the reaper's stale snapshot — leaked the port.

R2: delete _drop_workspace_lock and its two call sites
(stop_workspace tail, get_workspace 404 path). Popping the lock
object after release breaks mutual exclusion for any coroutine still
holding the old reference while a fresh caller gets a new lock
object — same ws_id can race two starts. The dict is bounded by the
number of workspaces so the leak is negligible; invariant lives on
WORKSPACE_LOCKS in a comment.

Tests:
- schedule/tests/test_janitor.py — 8 tests covering normal kill /
  healthy-skip / worker-race / never-started / multi-row batch /
  cancellation propagation / per-iteration self-heal
- runtime/tests/test_process.py — 7 tests covering reaper identity
  match / replacement-skip / alive-preserved + lock
  same-object / concurrent-serialize / survives-stop /
  helper-removed guard

uv run --package schedule pytest schedule/tests → 14 passed
uv run --package runtime pytest runtime/tests  →  7 passed

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
tao.chen
2026-08-20 10:49:47 +08:00
co-authored by Claude Fable 5
parent e5633cc95a
commit 4acfbb162f
5 changed files with 828 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
runtime/src/runtime/process.py
+72 -42
View File
@@ -90,6 +90,12 @@ class JupyterProcessRecord(TypedDict):
JUPYTER_PROCESSES: dict[str, JupyterProcessRecord] = {}
# ``WORKSPACE_LOCKS`` is intentionally never trimmed (see
# ``_drop_workspace_lock`` removal in P0-4 R2): popping a lock object
# after ``stop_workspace`` breaks mutual exclusion — any coroutine still
# holding the old reference continues awaiting it while a fresh caller
# receives a brand-new lock. The dict is bounded by the number of
# distinct workspaces, which is itself bounded by the DB.
WORKSPACE_LOCKS: dict[str, asyncio.Lock] = {}
_LOCKS_REGISTRY = asyncio.Lock()
_REAPER_TASK: asyncio.Task[None] | None = None
@@ -113,18 +119,6 @@ def get_workspace_lock(ws_id: str) -> asyncio.Lock:
return lock
def _drop_workspace_lock(ws_id: str) -> None:
"""Remove the per-workspace lock once no process references it.
Called after a successful stop to keep the registry bounded. We only
drop a lock we created and only when it is not held.
"""
lock = WORKSPACE_LOCKS.get(ws_id)
if lock is None or lock.locked():
return
WORKSPACE_LOCKS.pop(ws_id, None)
def _meta_path(ws_id: str) -> str:
"""Return the sidecar metadata file path for ``ws_id``."""
return str(WORKSPACES_ROOT / ws_id / RUNTIME_META_FILENAME)
@@ -412,7 +406,6 @@ async def stop_workspace(ws_id: str) -> dict:
del JUPYTER_PROCESSES[ws_id]
_delete_meta(ws_id)
_drop_workspace_lock(ws_id)
return {
"status": "stopped",
"workspace_id": ws_id,
@@ -472,7 +465,6 @@ async def get_workspace(ws_id: str) -> dict:
"last_used_at": p_info["last_used_at"],
}
_drop_workspace_lock(ws_id)
raise HTTPException(
status_code=404,
detail=(
@@ -558,6 +550,70 @@ def reconcile_processes() -> dict[str, int]:
return counters
async def _reap_once() -> None:
"""Single reap pass. See ``_reap_loop`` for the background wrapper.
Extracted so the reap decision logic is unit-testable without
dealing with the outer ``while True`` + sleep. P0-4 R1 fix: the
dead-process branch re-verifies ``(pid, started_at)`` against the
live dict before ``del``, so a record replaced by ``start_workspace``
mid-cycle is not erased by a stale snapshot.
"""
now = time.time()
victims: list[str] = []
for ws_id, info in list(JUPYTER_PROCESSES.items()):
if info["process"].poll() is not None:
# Identity check: between our snapshot above and the
# ``del`` below, ``start_workspace`` for the same ``ws_id``
# may have detected the dead record, replaced it with a
# fresh process, and written a new entry to
# ``JUPYTER_PROCESSES``. The naive ``del`` would erase
# the NEW record and leak its port / process — the very
# orphan the reaper is meant to prevent.
#
# ``(process.pid, started_at)`` is the stable identity of
# a record. If it changed, someone replaced the entry
# under us; leave it alone and let the next cycle
# re-evaluate the live process.
expected_pid = info["process"].pid
expected_started = info["started_at"]
current = JUPYTER_PROCESSES.get(ws_id)
if current is None:
continue
if (
current["process"].pid != expected_pid
or current["started_at"] != expected_started
):
logger.info(
f"reap: record for {ws_id} replaced mid-cycle "
f"(pid {expected_pid} -> {current['process'].pid}); skipping"
)
continue
logger.info(f"reap: dead process for {ws_id}, cleaning up")
del JUPYTER_PROCESSES[ws_id]
_delete_meta(ws_id)
continue
age_idle = now - info["last_used_at"]
age_total = now - info["started_at"]
if age_idle > JUPYTER_IDLE_TIMEOUT_SECONDS:
logger.info(
f"reap: idle Jupyter for {ws_id} "
f"(idle={age_idle:.0f}s > {JUPYTER_IDLE_TIMEOUT_SECONDS}s)"
)
victims.append(ws_id)
elif age_total > JUPYTER_MAX_LIFETIME_SECONDS:
logger.info(
f"reap: max-lifetime Jupyter for {ws_id} "
f"(age={age_total:.0f}s > {JUPYTER_MAX_LIFETIME_SECONDS}s)"
)
victims.append(ws_id)
for ws_id in victims:
try:
await stop_workspace(ws_id)
except Exception as exc:
logger.error(f"reap: failed to stop {ws_id}: {exc}")
async def _reap_loop() -> None:
"""Background task: idle + dead reaper.
@@ -572,38 +628,12 @@ async def _reap_loop() -> None:
"""
while True:
try:
await asyncio.sleep(JUPYTER_REAP_INTERVAL_SECONDS)
now = time.time()
victims: list[str] = []
for ws_id, info in list(JUPYTER_PROCESSES.items()):
if info["process"].poll() is not None:
logger.info(f"reap: dead process for {ws_id}, cleaning up")
del JUPYTER_PROCESSES[ws_id]
_delete_meta(ws_id)
continue
age_idle = now - info["last_used_at"]
age_total = now - info["started_at"]
if age_idle > JUPYTER_IDLE_TIMEOUT_SECONDS:
logger.info(
f"reap: idle Jupyter for {ws_id} "
f"(idle={age_idle:.0f}s > {JUPYTER_IDLE_TIMEOUT_SECONDS}s)"
)
victims.append(ws_id)
elif age_total > JUPYTER_MAX_LIFETIME_SECONDS:
logger.info(
f"reap: max-lifetime Jupyter for {ws_id} "
f"(age={age_total:.0f}s > {JUPYTER_MAX_LIFETIME_SECONDS}s)"
)
victims.append(ws_id)
for ws_id in victims:
try:
await stop_workspace(ws_id)
except Exception as exc:
logger.error(f"reap: failed to stop {ws_id}: {exc}")
await _reap_once()
except asyncio.CancelledError:
raise
except Exception as exc:
logger.exception(f"reap loop failed: {exc}")
await asyncio.sleep(JUPYTER_REAP_INTERVAL_SECONDS)
def start_reaper() -> asyncio.Task[None]:
runtime/tests/__init__.py
View File
runtime/tests/test_process.py
+216
View File
@@ -0,0 +1,216 @@
"""Tests for ``runtime.process`` P0-4 R1 + R2 fixes.
R1: ``_reap_once`` must verify a stuck dead-process record's identity
before deleting it — otherwise a ``start_workspace`` that replaced the
dead record mid-cycle would see its new entry silently erased by the
reaper's stale ``del``.
R2: ``WORKSPACE_LOCKS`` entries must survive ``stop_workspace`` —
popping the lock breaks mutual exclusion for any coroutine still
holding a reference to the old lock object.
"""
from __future__ import annotations
import asyncio
import time
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from runtime import process
def _make_dead_process(pid: int) -> MagicMock:
"""A ``subprocess.Popen``-shaped mock whose ``poll()`` says dead."""
proc = MagicMock(name=f"proc-{pid}")
proc.pid = pid
proc.poll = MagicMock(return_value=0) # exited
return proc
def _make_alive_process(pid: int) -> MagicMock:
"""A ``subprocess.Popen``-shaped mock whose ``poll()`` says alive."""
proc = MagicMock(name=f"proc-{pid}")
proc.pid = pid
proc.poll = MagicMock(return_value=None) # running
return proc
def _make_record(
*,
pid: int,
started_at: float,
last_used_at: float | None = None,
alive: bool = True,
) -> dict:
"""A ``JupyterProcessRecord``-shaped dict (subscriptable, like the TypedDict)."""
return {
"process": _make_alive_process(pid) if alive else _make_dead_process(pid),
"port": 8000 + pid,
"token": f"token-{pid}",
"base_url": "http://localhost",
"started_at": started_at,
"last_used_at": last_used_at if last_used_at is not None else started_at,
"meta_path": f"/tmp/{pid}.json",
}
@pytest.fixture(autouse=True)
def _isolate_registries():
"""Snapshot & restore the global registries so tests don't leak."""
saved_processes = dict(process.JUPYTER_PROCESSES)
saved_locks = dict(process.WORKSPACE_LOCKS)
yield
process.JUPYTER_PROCESSES.clear()
process.JUPYTER_PROCESSES.update(saved_processes)
process.WORKSPACE_LOCKS.clear()
process.WORKSPACE_LOCKS.update(saved_locks)
# ─── R1: reaper identity check ─────────────────────────────────────────
def test_reap_drops_dead_process_when_identity_matches() -> None:
"""The original (non-buggy) case still works: a dead process with
nobody replacing it must be cleaned up."""
now = time.time()
dead = _make_record(pid=100, started_at=now, alive=False)
process.JUPYTER_PROCESSES["ws1"] = dead
with patch.object(process, "_delete_meta") as mock_delete, \
patch.object(process, "stop_workspace", new=AsyncMock()):
asyncio.run(process._reap_once())
assert "ws1" not in process.JUPYTER_PROCESSES
mock_delete.assert_called_once_with("ws1")
def test_reap_skips_dead_process_when_record_replaced_by_new_process() -> None:
"""P0-4 R1 core bug. ``start_workspace`` detected the dead record
while the reaper was between snapshot and ``del``, removed it,
started a new process with a different pid, and wrote the new
record under the same ``ws_id``. The reaper's stale snapshot now
references the old (gone) pid; a naive ``del`` would erase the
new live process. Identity check must catch this and skip."""
now = time.time()
replacement = _make_record(pid=200, started_at=now, alive=True)
# Mutate the dict to simulate ``start_workspace`` having replaced
# the entry between snapshot time and now.
process.JUPYTER_PROCESSES["ws1"] = replacement
with patch.object(process, "_delete_meta") as mock_delete, \
patch.object(process, "stop_workspace", new=AsyncMock()):
asyncio.run(process._reap_once())
# The new live record must survive untouched.
assert process.JUPYTER_PROCESSES["ws1"] is replacement
mock_delete.assert_not_called()
def test_reap_keeps_alive_process_alone() -> None:
"""Sanity: a record whose process is still running must not be
dropped (only marked for idle/max-lifetime stop). With a fresh
``started_at`` neither threshold fires, so no victim either."""
now = time.time()
live = _make_record(pid=100, started_at=now, alive=True)
process.JUPYTER_PROCESSES["ws1"] = live
with patch.object(process, "stop_workspace", new=AsyncMock()) as mock_stop, \
patch.object(process, "_delete_meta") as mock_delete:
asyncio.run(process._reap_once())
assert process.JUPYTER_PROCESSES["ws1"] is live
mock_delete.assert_not_called()
mock_stop.assert_not_called()
# ─── R2: workspace locks survive stop ──────────────────────────────────
def test_workspace_lock_is_shared_across_callers() -> None:
"""``get_workspace_lock`` must return the *same* ``asyncio.Lock``
object every time — that is the whole point of the registry. If
a caller pops the entry between two ``get_workspace_lock`` calls,
a fresh lock would be returned and the two callers would no longer
be serialized."""
lock_a = process.get_workspace_lock("ws1")
lock_b = process.get_workspace_lock("ws1")
assert lock_a is lock_b
def test_concurrent_acquires_serialize_through_same_lock() -> None:
"""Two concurrent callers acquiring the lock for the same
``ws_id`` must serialize — proving they share a single lock object,
not two independent ones. With P0-4 R2 in place (lock never
popped) this is automatic; before the fix a popped entry would
let the second caller acquire a fresh, independent lock."""
lock = process.get_workspace_lock("ws1")
order: list[str] = []
async def first() -> None:
async with lock:
order.append("first-enter")
await asyncio.sleep(0)
order.append("first-exit")
async def second() -> None:
# Acquire *after* first releases — if we got a fresh lock here,
# this would interleave instead of waiting.
async with lock:
order.append("second-enter")
order.append("second-exit")
async def scenario() -> None:
await asyncio.gather(first(), second())
asyncio.run(scenario())
assert order == [
"first-enter", "first-exit", "second-enter", "second-exit",
], (
f"Lock leaked: order={order}. The two callers interleaved, "
"which means they held independent lock objects."
)
def test_workspace_lock_survives_manual_pop_simulating_stop() -> None:
"""Pre-fix, ``stop_workspace`` popped the lock from the registry
*after* the async-with exited. The lock object itself survived
(the dict entry just disappeared), but the next
``get_workspace_lock`` call would create a *new* lock and lose
serialization for any coroutine still holding the old reference.
This test simulates the pre-fix sequence manually: after one
``async with`` block exits, ``get_workspace_lock`` must still
return the *same* lock if we don't pop it; if we *did* pop it
(simulating the old buggy behavior), it would return a fresh
one. The test asserts the fixed invariant: even after the
stop-shaped sequence, the lock object is the same one."""
lock_a = process.get_workspace_lock("ws1")
async def use_and_release() -> None:
async with lock_a:
pass
asyncio.run(use_and_release())
# Pre-fix path: ``stop_workspace`` would pop the entry here.
# Post-fix: the helper is gone, so nothing pops the entry. Verify
# the invariant the fix preserves.
assert "ws1" in process.WORKSPACE_LOCKS
lock_b = process.get_workspace_lock("ws1")
assert lock_a is lock_b
def test_drop_workspace_lock_helper_is_removed() -> None:
"""Belt-and-suspenders: the buggy helper must not come back by
accident. ``_drop_workspace_lock`` removal is the entire R2 fix;
if someone re-adds it, this test fails before the leak does."""
assert not hasattr(process, "_drop_workspace_lock"), (
"_drop_workspace_lock was re-introduced — P0-4 R2 fix "
"regressed. Remove it again."
)
# Late import to avoid an unused-import lint complaint above.
schedule/src/schedule/orchestrator.py
+210 -1
View File
@@ -21,9 +21,10 @@ from common.db.models import (
ConsumerInbox,
OutboxEvents,
ScheduleNodeRuns,
ScheduleNodes,
ScheduleRuns,
)
from common.eventing import add_outbox_event, schedule_event_type, utcnow
from common.eventing import add_outbox_event, event_time, schedule_event_type, utcnow
from common.ids import new_ulid
from loguru import logger
from sqlalchemy import select
@@ -73,6 +74,27 @@ class DispatchOrchestrator:
# terminal state before the poll re-picks it.
LEASE_SLACK = timedelta(seconds=30)
# Margins used by the node-run janitor (see ``_janitor_loop``).
#
# ``NODE_JANITOR_GRACE_SECONDS`` is added on top of each node's
# ``timeout_seconds + retry_count * retry_interval_sec`` when judging
# whether a started-but-not-finished row is stuck. Absorbs the outbox
# ``min(30, 2**retry_count)`` backoff + the lease slack above + a few
# seconds of scheduling jitter. Tuned for the worst case the
# orchestrator itself produces, so the janitor cannot race a legitimate
# retry path to terminal state.
NODE_JANITOR_GRACE_SECONDS = 120
# Floor for ``started_at IS NULL`` rows (never dispatched). 1h rides
# out an orchestrator restart that drops the lease mid-flight; past
# that the row is dead and forcing a terminal state lets the DAG
# advance.
NODE_JANITOR_QUEUED_GRACE_SECONDS = 3600
# Loop period + batch size. Detection lag = interval + the time it
# takes to scan, currently well under a minute. 50 rows per cycle
# keeps worst-case fan-out bounded.
NODE_JANITOR_INTERVAL_SECONDS = 30
NODE_JANITOR_BATCH = 50
def __init__(
self,
*,
@@ -88,6 +110,7 @@ class DispatchOrchestrator:
self.dispatch_lock = asyncio.Lock()
self._loop_task: asyncio.Task[None] | None = None
self._exec_loop_task: asyncio.Task[None] | None = None
self._janitor_task: asyncio.Task[None] | None = None
self._exec_tasks: set[asyncio.Task[None]] = set()
self._exec_semaphore = asyncio.Semaphore(execution_concurrency)
@@ -100,6 +123,10 @@ class DispatchOrchestrator:
self._execution_loop(),
name="scheduler-node-execute",
)
self._janitor_task = asyncio.create_task(
self._janitor_loop(),
name="scheduler-node-janitor",
)
async def close(self) -> None:
if self._loop_task is not None:
@@ -116,6 +143,13 @@ class DispatchOrchestrator:
except asyncio.CancelledError:
pass
self._exec_loop_task = None
if self._janitor_task is not None:
self._janitor_task.cancel()
try:
await self._janitor_task
except asyncio.CancelledError:
pass
self._janitor_task = None
if self._exec_tasks:
await asyncio.gather(*self._exec_tasks, return_exceptions=True)
self._exec_tasks.clear()
@@ -148,6 +182,181 @@ class DispatchOrchestrator:
logger.exception("node execute loop failed")
await asyncio.sleep(1)
async def _janitor_loop(self) -> None:
"""Periodically force terminal state on stuck ``ScheduleNodeRuns``.
Original P0-4 S2 bug: when ``job.node.execute`` outbox retries are
exhausted (5 tries, capped 30s backoff) the *event* is marked
``failed`` but no one reconciles the *node_run* — the row stays
in ``queued`` or ``running`` forever, DAG children are never
dispatched, and the whole run sits at "running" with no further
progress.
A second, harder failure mode: the worker crashes (OOM /
``kill -9`` / forgotten upload) mid-execution. The lease-based
re-eligibility keeps re-dispatching the event, but a worker
that's stuck without writing a terminal state is invisible to
the retry counter — it just keeps timing out.
Both collapse into one wall-clock condition: ``started_at +
`` budget < now()`` (or, more rarely, ``started_at IS NULL`` for
long enough). The janitor enforces that condition so no row
can outlive its deadline regardless of which subsystem let go.
"""
while True:
try:
killed = await self._reap_stuck_node_runs()
if killed:
logger.warning(
"node janitor reaped {} stuck node run(s)", killed,
)
except asyncio.CancelledError:
raise
except Exception:
logger.exception("node janitor loop failed")
await asyncio.sleep(self.NODE_JANITOR_INTERVAL_SECONDS)
async def _reap_stuck_node_runs(self) -> int:
"""Find stuck ``ScheduleNodeRuns`` rows and force-terminal them.
SQL does a coarse pre-filter (``created_at`` at least one hour
old so we don't even look at fresh rows). Per-row budget is
computed in Python because it depends on each row's
``ScheduleNodes.timeout_seconds`` + ``retry_count`` +
``retry_interval_sec``, which can vary row to row.
Returns the count of rows actually flipped to ``timed_out``;
rows that were already terminal (worker got there first) are
silently skipped, so two janitors running side by side stay
idempotent.
"""
queued_cutoff = utcnow() - timedelta(
seconds=self.NODE_JANITOR_QUEUED_GRACE_SECONDS,
)
async with session_scope(self.session_factory) as session:
statement = (
select(ScheduleNodeRuns, ScheduleNodes, ScheduleRuns)
.join(
ScheduleNodes,
ScheduleNodes.node_id == ScheduleNodeRuns.node_id,
)
.join(
ScheduleRuns,
ScheduleRuns.run_id == ScheduleNodeRuns.run_id,
)
.where(
ScheduleNodeRuns.node_status.in_(("queued", "running")),
ScheduleNodeRuns.is_deleted == 0,
ScheduleNodes.is_deleted == 0,
ScheduleRuns.is_deleted == 0,
ScheduleNodeRuns.created_at <= queued_cutoff,
)
.order_by(ScheduleNodeRuns.created_at)
.limit(self.NODE_JANITOR_BATCH)
)
candidates = list((await session.execute(statement)).all())
now = utcnow()
killed = 0
for node_run, node, run in candidates:
if node_run.started_at is not None:
budget = (
node.timeout_seconds
+ node.retry_count * node.retry_interval_sec
+ self.NODE_JANITOR_GRACE_SECONDS
)
deadline = node_run.started_at + timedelta(
seconds=budget,
)
if now <= deadline:
# Healthy row that just happened to be in the
# SQL pre-filter window; skip without writing.
continue
# ``started_at IS NULL`` rows fell through the coarse
# ``created_at <= queued_cutoff`` filter, so we know
# they are at least an hour old and never dispatched.
if await self._force_terminal_node_run(
session,
node_run=node_run,
run=run,
reason="NODE_RUN_TIMEOUT",
):
killed += 1
return killed
async def _force_terminal_node_run(
self,
session: AsyncSession,
*,
node_run: ScheduleNodeRuns,
run: ScheduleRuns,
reason: str,
) -> bool:
"""Mark ``node_run`` as ``timed_out`` and enqueue a ``NODE_FINISHED_EVENT``.
Re-reads the row under ``with_for_update`` so a worker that
reports the result concurrently can't be undone by a second
``timed_out`` write (and vice versa). Returns ``True`` iff this
call performed the terminal write; ``False`` if the row had
already moved on (worker raced us, another janitor raced us).
The idempotency key is distinct from the worker-reported
``:finished`` key so the same ``node_run`` won't produce two
``NODE_FINISHED_EVENT`` rows if both paths fire. The DAG
consumer is idempotent on its end (``_advance_run`` is a no-op
when ``node_status`` is already terminal), so even if a stray
duplicate slipped through it would self-heal.
"""
locked = await session.scalar(
select(ScheduleNodeRuns)
.where(ScheduleNodeRuns.node_run_id == node_run.node_run_id)
.with_for_update()
)
if locked is None or locked.node_status in TERMINAL_NODE_STATES:
return False
finished_at = utcnow()
started_or_created = locked.started_at or locked.created_at
locked.node_status = "timed_out"
locked.finished_at = finished_at
locked.duration_ms = int(
(finished_at - started_or_created).total_seconds() * 1000,
)
locked.error_code = reason
locked.message = (
"节点执行超时已自动终止"
if locked.started_at is not None
else "节点从未派发，调度超时已自动终止"
)[:2000]
locked.state_version += 1
await add_outbox_event(
session,
event_type=NODE_FINISHED_EVENT,
producer="schedule-janitor",
trace_id=new_ulid(),
aggregate_type="schedule_node_run",
aggregate_id=locked.node_run_id,
idempotency_key=(
f"{locked.node_run_id}:{locked.attempt_no}:timed_out"
),
payload={
"workspace_id": run.workspace_id,
"run_id": locked.run_id,
"node_run_id": locked.node_run_id,
"node_id": locked.node_id,
"versions_id": locked.versions_id,
"attempt_no": locked.attempt_no,
"node_status": locked.node_status,
"exit_code": None,
"started_at": event_time(locked.started_at),
"finished_at": event_time(finished_at),
"duration_ms": locked.duration_ms,
"logs_object_id": None,
"result_object_id": None,
"error_code": reason,
"error_message": locked.message,
},
)
return True
async def _claim_execution_events(
self,
*,
schedule/tests/test_janitor.py
+330
View File
@@ -0,0 +1,330 @@
"""Tests for the node-run janitor (P0-4 S2).
Pre-fix, ``job.node.execute`` outbox retries exhausted (5 tries, capped
30s backoff) marked the *event* ``failed`` but no one reconciled the
*node_run* — the row stayed in ``queued`` / ``running`` forever and
the DAG stopped advancing. The janitor enforces a wall-clock deadline
so any stuck row is eventually flipped to ``timed_out``.
These tests mock the session — no live MySQL. The orchestration
*logic* (which rows count as stuck, what terminal state to write,
which outbox event to enqueue) is what's worth pinning; the SQL itself
is exercised in the integration environment.
"""
from __future__ import annotations
import asyncio
from datetime import timedelta
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from common.db.models import ScheduleNodeRuns
from common.eventing import utcnow
from schedule.orchestrator import NODE_FINISHED_EVENT, DispatchOrchestrator
def _make_orchestrator() -> DispatchOrchestrator:
"""Construct a ``DispatchOrchestrator`` with a fake session factory."""
return DispatchOrchestrator(
session_factory=MagicMock(name="session_factory"),
node_execute_handler=AsyncMock(),
execution_concurrency=4,
)
def _make_node_run(
*,
node_run_id: str = "nr1",
run_id: str = "r1",
node_id: str = "n1",
versions_id: str = "v1",
attempt_no: int = 1,
node_status: str = "running",
started_at=None,
created_at=None,
finished_at=None,
state_version: int = 0,
) -> ScheduleNodeRuns:
"""Build an in-memory ``ScheduleNodeRuns`` row."""
return ScheduleNodeRuns(
node_run_id=node_run_id,
run_id=run_id,
node_id=node_id,
versions_id=versions_id,
attempt_no=attempt_no,
node_status=node_status,
state_version=state_version,
created_at=created_at or utcnow(),
started_at=started_at,
finished_at=finished_at,
)
def _make_run(workspace_id: str = "ws1") -> SimpleNamespace:
return SimpleNamespace(workspace_id=workspace_id)
def _make_node(
*,
timeout_seconds: int = 600,
retry_count: int = 0,
retry_interval_sec: int = 5,
) -> SimpleNamespace:
return SimpleNamespace(
timeout_seconds=timeout_seconds,
retry_count=retry_count,
retry_interval_sec=retry_interval_sec,
)
def _open_session_scope(fake_session: AsyncMock) -> MagicMock:
"""Wrap ``fake_session`` in an async context-manager stand-in."""
cm = MagicMock()
cm.__aenter__ = AsyncMock(return_value=fake_session)
cm.__aexit__ = AsyncMock(return_value=None)
return cm
async def test_force_terminal_writes_timed_out_and_enqueues_event() -> None:
"""A running row over deadline must flip to ``timed_out`` + emit
``NODE_FINISHED_EVENT`` whose idempotency key is the janitor
variant (``...:timed_out``) so it can't collide with the worker's
``...:finished`` key on the same row."""
orch = _make_orchestrator()
node_run = _make_node_run(started_at=utcnow() - timedelta(seconds=10_000))
fake_session = AsyncMock(name="session")
# Re-read under FOR UPDATE returns the same still-running row.
fake_session.scalar = AsyncMock(return_value=node_run)
fake_session.add = MagicMock()
result = await orch._force_terminal_node_run(
fake_session,
node_run=node_run,
run=_make_run(),
reason="NODE_RUN_TIMEOUT",
)
assert result is True
assert node_run.node_status == "timed_out"
assert node_run.error_code == "NODE_RUN_TIMEOUT"
assert node_run.finished_at is not None
assert node_run.state_version == 1
assert node_run.duration_ms is not None and node_run.duration_ms >= 0
fake_session.add.assert_called_once()
event = fake_session.add.call_args[0][0]
assert event.event_type == NODE_FINISHED_EVENT
assert event.aggregate_id == node_run.node_run_id
assert event.idempotency_key.endswith(":timed_out")
assert event.payload_json["workspace_id"] == "ws1"
assert event.payload_json["node_status"] == "timed_out"
assert event.payload_json["error_code"] == "NODE_RUN_TIMEOUT"
assert event.payload_json["exit_code"] is None
async def test_force_terminal_skips_when_worker_raced_to_terminal() -> None:
"""If the worker reports a result between SELECT and the lock
re-read, the janitor must NOT overwrite the worker's terminal
state — ``_advance_run`` already handled the row correctly."""
orch = _make_orchestrator()
candidate = _make_node_run(node_status="running")
# Re-read under FOR UPDATE shows the worker won the race.
already_terminal = _make_node_run(node_status="succeeded")
fake_session = AsyncMock()
fake_session.scalar = AsyncMock(return_value=already_terminal)
fake_session.add = MagicMock()
result = await orch._force_terminal_node_run(
fake_session,
node_run=candidate,
run=_make_run(),
reason="NODE_RUN_TIMEOUT",
)
assert result is False
fake_session.add.assert_not_called()
# The row that was re-read under the lock kept its terminal status.
assert already_terminal.node_status == "succeeded"
async def test_force_terminal_handles_never_started_old_row() -> None:
"""``started_at IS NULL`` + very-old ``created_at`` = row that
never even got dispatched. Janitor must still force-terminal with
the "never dispatched" message variant so operators can tell the
two failure modes apart in logs."""
orch = _make_orchestrator()
node_run = _make_node_run(started_at=None, node_status="queued")
fake_session = AsyncMock()
fake_session.scalar = AsyncMock(return_value=node_run)
fake_session.add = MagicMock()
await orch._force_terminal_node_run(
fake_session,
node_run=node_run,
run=_make_run(),
reason="NODE_RUN_TIMEOUT",
)
assert node_run.node_status == "timed_out"
assert "从未派发" in (node_run.message or "")
assert node_run.duration_ms is not None and node_run.duration_ms >= 0
async def test_reap_kills_running_row_past_deadline() -> None:
"""End-to-end of ``_reap_stuck_node_runs``: a row whose
``started_at`` is way past ``timeout + grace`` must be flipped,
and the loop must report ``killed == 1``."""
orch = _make_orchestrator()
node = _make_node(timeout_seconds=600)
node_run = _make_node_run(started_at=utcnow() - timedelta(seconds=10_000))
run = _make_run()
fake_session = AsyncMock()
# SELECT returns one candidate.
fake_session.execute = AsyncMock(
return_value=MagicMock(all=MagicMock(return_value=[(node_run, node, run)])),
)
# Re-read under FOR UPDATE returns the same row.
fake_session.scalar = AsyncMock(return_value=node_run)
fake_session.add = MagicMock()
with patch(
"schedule.orchestrator.session_scope",
return_value=_open_session_scope(fake_session),
):
killed = await orch._reap_stuck_node_runs()
assert killed == 1
assert node_run.node_status == "timed_out"
fake_session.add.assert_called_once()
async def test_reap_skips_healthy_row() -> None:
"""A row whose deadline is comfortably in the future must not be
re-read or mutated — the per-row budget check rejects it before
touching the lock path."""
orch = _make_orchestrator()
node = _make_node(timeout_seconds=600)
# Started 10s ago; deadline is 600 + 120 = 720s from start.
node_run = _make_node_run(started_at=utcnow() - timedelta(seconds=10))
run = _make_run()
fake_session = AsyncMock()
fake_session.execute = AsyncMock(
return_value=MagicMock(all=MagicMock(return_value=[(node_run, node, run)])),
)
fake_session.scalar = AsyncMock()
fake_session.add = MagicMock()
with patch(
"schedule.orchestrator.session_scope",
return_value=_open_session_scope(fake_session),
):
killed = await orch._reap_stuck_node_runs()
assert killed == 0
fake_session.scalar.assert_not_called()
fake_session.add.assert_not_called()
assert node_run.node_status == "running"
async def test_reap_processes_multiple_rows_in_one_pass() -> None:
"""A batch of 3 rows — 2 over deadline, 1 healthy — must produce
exactly 2 kills. The healthy one is short-circuited in Python
before the lock path."""
orch = _make_orchestrator()
node = _make_node(timeout_seconds=600)
stuck_a = _make_node_run(
node_run_id="stuck_a", started_at=utcnow() - timedelta(seconds=10_000),
)
stuck_b = _make_node_run(
node_run_id="stuck_b", started_at=utcnow() - timedelta(seconds=9_000),
)
healthy = _make_node_run(
node_run_id="healthy", started_at=utcnow() - timedelta(seconds=10),
)
run = _make_run()
fake_session = AsyncMock()
fake_session.execute = AsyncMock(
return_value=MagicMock(
all=MagicMock(return_value=[
(stuck_a, node, run),
(stuck_b, node, run),
(healthy, node, run),
]),
),
)
# Re-read for the two stuck rows; healthy is skipped before this.
fake_session.scalar = AsyncMock(side_effect=[stuck_a, stuck_b])
fake_session.add = MagicMock()
with patch(
"schedule.orchestrator.session_scope",
return_value=_open_session_scope(fake_session),
):
killed = await orch._reap_stuck_node_runs()
assert killed == 2
assert stuck_a.node_status == "timed_out"
assert stuck_b.node_status == "timed_out"
assert healthy.node_status == "running"
assert fake_session.add.call_count == 2
async def test_janitor_loop_propagates_cancellation() -> None:
"""``_janitor_loop`` must re-raise ``CancelledError`` so ``close()``
can shut it down cleanly. Patches ``asyncio.sleep`` to raise on
first call so the test doesn't run forever."""
orch = _make_orchestrator()
orch._reap_stuck_node_runs = AsyncMock(return_value=0)
async def cancel_on_sleep(_seconds: float) -> None:
raise asyncio.CancelledError
with patch(
"schedule.orchestrator.asyncio.sleep",
side_effect=cancel_on_sleep,
):
with pytest.raises(asyncio.CancelledError):
await orch._janitor_loop()
async def test_janitor_loop_continues_after_reap_exception() -> None:
"""A failure inside one reap pass must not kill the loop — the
next iteration still runs. The janitor must be self-healing so a
bad row or transient DB blip can't take the whole service down."""
orch = _make_orchestrator()
call_count = 0
async def _reap_then_succeed() -> int:
nonlocal call_count
call_count += 1
if call_count == 1:
raise RuntimeError("simulated DB blip")
return 3
orch._reap_stuck_node_runs = AsyncMock(side_effect=_reap_then_succeed)
sleep_calls = 0
async def _count_sleeps(_seconds: float) -> None:
nonlocal sleep_calls
sleep_calls += 1
if sleep_calls >= 2:
raise asyncio.CancelledError
with patch(
"schedule.orchestrator.asyncio.sleep",
side_effect=_count_sleeps,
):
with pytest.raises(asyncio.CancelledError):
await orch._janitor_loop()
assert call_count == 2