tests/test_cancel_signal_guards.py
Rama main · réplica de lectura del repositorio original.
"""Audit 2026-07-05 regression tests: preempt-signal write guards (A1) and
cancel-honoring claim teardown (A2), plus the smaller correctness fixes from
the same batch (F-3 finished_at on cwd_unreachable, phase-2 emit-after-commit,
bounded /log reads, /events reserved-key 422).
A1 — `signal` writers race state transitions like state writers do: the
sweeper's auto-preempt and /preempt-blockers stamped `signal='preempt'` with
plain ORM writes, and the /next-job claim never cleared `signal`, so a stale
preempt could ride a requeue onto a fresh claim and kill it seconds in.
A2 — every requeue path cleared `signal` unconditionally (M1), which erased a
pending user CANCEL: the user was told "cancelling", the claim was torn down
(admission refusal / dead worker), and the job silently re-ran elsewhere.
Teardown now honors a pending cancel (job -> CANCELLED) and only clears
non-cancel signals.
"""
from __future__ import annotations
from datetime import UTC, datetime, timedelta
import pytest
from fastapi.testclient import TestClient
from sqlalchemy import update
from sqlalchemy.orm import Session, sessionmaker
from jobd.broker import sweeper as sweeper_mod
from jobd.broker.constants import MAX_LOG_CHUNK_BYTES
# The cancel endpoint resolves _cas_state from the routes module's globals
# (Stage-3 split) — patch THERE; a patch on jobd.app would be a silent no-op.
from jobd.broker.routes import jobs as jobs_routes_mod
from jobd.db import Job
from jobd.models import JobState
def _hb(client: TestClient, host: str, in_flight: list[int] | None = None):
body: dict = {
"host": host,
"free_vram_gb": 0,
"unregistered_vram_gb": 0,
"free_ram_gb": 8,
"idle_cpus": 4,
}
if in_flight is not None:
body["in_flight_job_ids"] = in_flight
r = client.post("/heartbeat", json=body)
assert r.status_code == 200, r.text
return r
def _submit(client: TestClient, **overrides) -> int:
payload = {"cmd": ["true"], "cwd": "/tmp", "project": "project-a", **overrides}
r = client.post("/submit", json=payload)
assert r.status_code == 200, r.text
return r.json()["id"]
def _claim(client: TestClient, host: str) -> dict | None:
r = client.post(
"/next-job",
json={
"host": host,
"free_vram_gb": 0,
"unregistered_vram_gb": 0,
"free_ram_gb": 8,
"idle_cpus": 4,
},
)
return r.json()
def _db_update(client: TestClient, where_id: int, **values) -> None:
engine = client.app.state.engine
with engine.begin() as conn:
conn.execute(update(Job).where(Job.id == where_id).values(**values))
# ---------------------------------------------------------------------------
# A1 — stale-signal guards
# ---------------------------------------------------------------------------
def test_claim_clears_stale_signal_on_queued_row(client):
"""A1: a leftover signal on a QUEUED row (stamped by any unguarded writer
mid-requeue) must not survive the queued->assigned claim — the fresh run
would honor it on the first /signal poll and die ~2s in."""
_hb(client, "w1")
job_id = _submit(client)
_db_update(client, job_id, signal="preempt")
got = _claim(client, "w1")
assert got is not None and got["id"] == job_id
assert client.get(f"/jobs/{job_id}/signal").json()["signal"] is None
def test_sweeper_auto_preempt_signal_write_is_cas_guarded(client, monkeypatch):
"""A1: the sweeper's auto-preempt must not stamp 'preempt' onto a candidate
that left RUNNING between the candidate SELECT and the signal write (a
/complete or reconcile-requeue landing in the window). Injected exactly
like the F1 reclaim test: flip the row terminal right as the guard runs,
then delegate to the real CAS — which must match 0 rows and emit nothing."""
_hb(client, "w1")
candidate_id = _submit(client)
assert _claim(client, "w1")["id"] == candidate_id
client.post(f"/jobs/{candidate_id}/started")
# Preemptible, low priority, running long enough to clear the runtime floor.
_db_update(
client,
candidate_id,
preemptible=True,
priority=10,
started_at=datetime.now(UTC) - timedelta(hours=1),
)
# High-priority queued job, old enough to trip the blocked-queue probe.
queued_id = _submit(client)
_db_update(client, queued_id, priority=90, submitted_at=datetime.now(UTC) - timedelta(hours=1))
real_cas = sweeper_mod._cas_state
def racing_cas(session, jid, expected, **values):
if values.get("signal") == "preempt":
monkeypatch.setattr(sweeper_mod, "_cas_state", real_cas) # fire once
session.execute(
update(Job).where(Job.id == jid).values(state=JobState.COMPLETED.value, exit_code=0)
)
return real_cas(session, jid, expected, **values)
monkeypatch.setattr(sweeper_mod, "_cas_state", racing_cas)
client.app.state.sweep_once()
got = client.get(f"/jobs/{candidate_id}").json()
assert got["state"] == "completed"
assert client.get(f"/jobs/{candidate_id}/signal").json()["signal"] is None
assert client.get("/events", params={"event": "auto_preempt"}).json() == []
def test_sweeper_auto_preempt_still_fires_on_live_candidate(client):
"""Companion to the race test: with no interference the auto-preempt CAS
wins, the signal lands, and the event is emitted (after commit)."""
_hb(client, "w1")
candidate_id = _submit(client)
assert _claim(client, "w1")["id"] == candidate_id
client.post(f"/jobs/{candidate_id}/started")
_db_update(
client,
candidate_id,
preemptible=True,
priority=10,
started_at=datetime.now(UTC) - timedelta(hours=1),
)
queued_id = _submit(client)
_db_update(client, queued_id, priority=90, submitted_at=datetime.now(UTC) - timedelta(hours=1))
client.app.state.sweep_once()
assert client.get(f"/jobs/{candidate_id}/signal").json()["signal"] == "preempt"
events = client.get("/events", params={"event": "auto_preempt"}).json()
assert [e["job_id"] for e in events] == [candidate_id]
def test_preempt_blockers_lost_race_returns_unsignaled(client, monkeypatch):
"""A1: /preempt-blockers must not stamp 'preempt' onto a candidate that
completed between its SELECT and the write; the caller gets signaled=None
instead of a false success."""
_hb(client, "w1")
candidate_id = _submit(client)
assert _claim(client, "w1")["id"] == candidate_id
client.post(f"/jobs/{candidate_id}/started")
_db_update(client, candidate_id, preemptible=True, priority=10)
queued_id = _submit(client)
_db_update(client, queued_id, priority=90)
real_cas = jobs_routes_mod._cas_state
def racing_cas(session, jid, expected, **values):
if values.get("signal") == "preempt":
monkeypatch.setattr(jobs_routes_mod, "_cas_state", real_cas) # fire once
session.execute(
update(Job).where(Job.id == jid).values(state=JobState.COMPLETED.value, exit_code=0)
)
return real_cas(session, jid, expected, **values)
monkeypatch.setattr(jobs_routes_mod, "_cas_state", racing_cas)
r = client.post(f"/jobs/{queued_id}/preempt-blockers", params={"force": True})
assert r.status_code == 200
assert r.json()["signaled"] is None
assert client.get(f"/jobs/{candidate_id}").json()["state"] == "completed"
assert client.get(f"/jobs/{candidate_id}/signal").json()["signal"] is None
assert client.get("/events", params={"event": "auto_preempt"}).json() == []
# ---------------------------------------------------------------------------
# A2 — pending cancel honored on claim teardown
# ---------------------------------------------------------------------------
def _assigned_job_with_pending_cancel(client, host: str = "w1") -> int:
_hb(client, host)
job_id = _submit(client)
assert _claim(client, host)["id"] == job_id
r = client.post(f"/jobs/{job_id}/cancel")
assert r.status_code == 200
assert client.get(f"/jobs/{job_id}/signal").json()["signal"] == "cancel"
return job_id
def test_refuse_admission_honors_pending_cancel_gpu_contention(client):
"""A2: user cancels an ASSIGNED job (told "cancelling"), then the worker
refuses admission. The old requeue erased the cancel and the job re-ran
elsewhere; it must land CANCELLED instead."""
job_id = _assigned_job_with_pending_cancel(client)
r = client.post(
f"/jobs/{job_id}/refuse-admission",
json={"reason": "gpu_contention", "required_gb": 8, "free_gb": 1},
)
assert r.status_code == 200, r.text
got = client.get(f"/jobs/{job_id}").json()
assert got["state"] == "cancelled", got
assert got["finished_at"] is not None
assert client.get(f"/jobs/{job_id}/signal").json()["signal"] is None
# The honored cancel is on the event record; no admission_blocked requeue.
cancels = client.get("/events", params={"event": "job_cancelled"}).json()
assert any(
e["job_id"] == job_id and e["payload"].get("via") == "refuse_admission" for e in cancels
)
# And the job must never be claimable again.
assert _claim(client, "w1") is None
def test_refuse_admission_honors_pending_cancel_cwd_missing(client):
"""A2, cwd branch: same contract when the refusal reason is a missing cwd.
Notably this must hold when NO other eligible worker has the cwd — without
the hoisted cancel check the terminal branch mislabels the cancelled job
as failed/cwd_unreachable. The cancel pre-empts the whole refusal
disposition, so no cwd_refused re-route decision is recorded."""
job_id = _assigned_job_with_pending_cancel(client)
r = client.post(
f"/jobs/{job_id}/refuse-admission",
json={"reason": "cwd_missing", "cwd": "/tmp"},
)
assert r.status_code == 200, r.text
got = client.get(f"/jobs/{job_id}").json()
assert got["state"] == "cancelled", got
assert got["termination_reason"] != "cwd_unreachable"
assert client.get(f"/jobs/{job_id}/signal").json()["signal"] is None
def test_reconcile_requeue_honors_pending_cancel(client):
"""A2, reconcile path: a restarted worker (empty in-flight reports) whose
ASSIGNED job has a pending cancel must see it CANCELLED, not requeued."""
job_id = _assigned_job_with_pending_cancel(client, host="ghost")
# Old enough to reconcile, and missing from 2 consecutive reports.
_db_update(client, job_id, started_at=datetime.now(UTC) - timedelta(minutes=5))
_hb(client, "ghost", in_flight=[])
_hb(client, "ghost", in_flight=[])
got = client.get(f"/jobs/{job_id}").json()
assert got["state"] == "cancelled", got
assert got["finished_at"] is not None
cancels = client.get("/events", params={"event": "job_cancelled"}).json()
assert any(e["job_id"] == job_id and e["payload"].get("via") == "reconcile" for e in cancels)
def test_reconcile_requeue_still_clears_preempt_signal(client):
"""M1 stays intact for non-cancel signals: a pending preempt on a dead
worker's ASSIGNED job is cleared by the requeue, not honored."""
_hb(client, "ghost")
job_id = _submit(client)
assert _claim(client, "ghost")["id"] == job_id
_db_update(
client, job_id, signal="preempt", started_at=datetime.now(UTC) - timedelta(minutes=5)
)
_hb(client, "ghost", in_flight=[])
_hb(client, "ghost", in_flight=[])
got = client.get(f"/jobs/{job_id}").json()
assert got["state"] == "queued", got
assert client.get(f"/jobs/{job_id}/signal").json()["signal"] is None
# ---------------------------------------------------------------------------
# F-3 — cwd_unreachable is a fully-formed terminal
# ---------------------------------------------------------------------------
def test_cwd_unreachable_sets_finished_at(client):
"""F-3: the cwd_unreachable terminal transition must stamp finished_at like
every other terminalizing write — retention pruning filters on it, so a
NULL finished_at made these rows immortal."""
_hb(client, "w1")
job_id = _submit(client)
assert _claim(client, "w1")["id"] == job_id
r = client.post(
f"/jobs/{job_id}/refuse-admission",
json={"reason": "cwd_missing", "cwd": "/tmp"},
)
assert r.status_code == 200, r.text
got = client.get(f"/jobs/{job_id}").json()
assert got["state"] == "failed"
assert got["termination_reason"] == "cwd_unreachable"
assert got["finished_at"] is not None
# ---------------------------------------------------------------------------
# F-4 — phase-2 emit-after-commit
# ---------------------------------------------------------------------------
def test_sweeper_phase2_emits_nothing_when_commit_fails(client, tmp_path):
"""F-4: auto_preempt (and sweep_warning) events are emitted only after the
phase-2 commit. A commit failure must not leave events.jsonl claiming a
preempt signal the DB never recorded."""
_hb(client, "w1")
candidate_id = _submit(client)
assert _claim(client, "w1")["id"] == candidate_id
client.post(f"/jobs/{candidate_id}/started")
_db_update(
client,
candidate_id,
preemptible=True,
priority=10,
started_at=datetime.now(UTC) - timedelta(hours=1),
)
queued_id = _submit(client)
_db_update(client, queued_id, priority=90, submitted_at=datetime.now(UTC) - timedelta(hours=1))
engine = client.app.state.engine
class ExplodingSecondCommit(Session):
commits = 0
def commit(self):
type(self).commits += 1
if type(self).commits >= 2: # 1st = phase 1, 2nd = phase 2
raise RuntimeError("injected phase-2 commit failure")
super().commit()
factory = sessionmaker(bind=engine, class_=ExplodingSecondCommit)
logs_dir = tmp_path / "logs" # same logs_path the client fixture built with
with pytest.raises(RuntimeError, match="injected phase-2 commit failure"):
sweeper_mod.sweep_once(factory, logs_dir, lambda: None)
# The signal write was rolled back with the commit — and no event claims it.
assert client.get(f"/jobs/{candidate_id}/signal").json()["signal"] is None
assert client.get("/events", params={"event": "auto_preempt"}).json() == []
# ---------------------------------------------------------------------------
# /log bounded read + /events reserved keys
# ---------------------------------------------------------------------------
def test_log_append_rejects_oversized_chunk(client):
"""The /log size gate must fire on an oversized chunk (and the read is
bounded — the old code buffered the whole body before checking)."""
_hb(client, "w1")
job_id = _submit(client)
assert _claim(client, "w1")["id"] == job_id
r = client.post(f"/jobs/{job_id}/log", content=b"x" * (MAX_LOG_CHUNK_BYTES + 1))
assert r.status_code == 413
ok = client.post(f"/jobs/{job_id}/log", content=b"hello\n")
assert ok.status_code == 200
assert ok.json()["bytes"] == 6
def test_events_payload_reserved_key_is_422_not_500(client):
"""A payload key colliding with an envelope field ('source', 'job_id', …)
used to raise TypeError at the _emit_event call — an uncaught 500. It must
be a 422 at the Pydantic boundary."""
r = client.post(
"/events",
json={"source": "hook", "event": "custom", "payload": {"source": "broker"}},
)
assert r.status_code == 422
assert "reserved" in r.text
ok = client.post(
"/events",
json={"source": "hook", "event": "custom", "payload": {"note": "fine"}},
)
assert ok.status_code == 204