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
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