fix: P0-4 schedule node janitor + runtime reaper/lock invariants
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
co-authored by
Claude Fable 5
Claude Fable 5
parent
e5633cc95a
commit
4acfbb162f
@@ -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
|
||||
Reference in New Issue
Block a user