tests/test_worker_concurrency.py
Rama main · réplica de lectura del repositorio original.
"""`max_concurrent > 1` must actually run N jobs at once, and never N+1.
The 2026-07-25 adverse-conditions audit left this open as a LOW: the setting was
covered only by *parsing* tests (`_max_concurrent_jobs()` returns 1 by default,
the API round-trips the field) and by nothing that ran two jobs at the same
time. So the advertised behaviour — the thing the README's pueue/task-spooler
migration table points migrants at as its answer to "parallelism limits" — had
no test at all. Every worker in the author's own fleet runs `max_concurrent: 1`,
so production did not exercise it either.
How the bound actually works, because the tests below depend on it:
- ONE poll thread runs the gate. It reads ``len(_in_flight)`` and, if that is
below max_concurrent, claims a job and calls `_reserve_and_dispatch`.
- `_reserve_and_dispatch` calls `_register_in_flight` **synchronously, in the
poll thread, before it returns** — then starts a daemon thread to do the
work. That ordering is the whole bound. Its docstring records the bug it
replaced: registering inside the worker thread left a window between
``t.start()`` and the thread being scheduled in which the gate could
re-poll against a stale count and oversubscribe.
- Job threads only ever *unregister*, concurrently with the poller.
So the invariant is not "many threads race to reserve" — it is "one poller
reserves while N workers release, and the live count never exceeds the cap."
These tests model exactly that, with real threads.
"""
from __future__ import annotations
import threading
import time
import pytest
import jobd.worker.job_worker as jw
@pytest.fixture(autouse=True)
def _clean_in_flight():
"""Never leak reservations between tests — the registry is module state."""
with jw._in_flight_lock:
jw._in_flight.clear()
yield
with jw._in_flight_lock:
jw._in_flight.clear()
def _job(job_id: int) -> dict:
return {"id": job_id, "vram_gb": 0, "ram_gb": 0, "cpus": 1}
# --- the ordering that bounds concurrency ---------------------------------
class _ManualThread:
"""A `threading.Thread` stand-in whose `start()` deliberately does nothing.
The invariant under test is an *ordering* one — the reservation must be
recorded before `_reserve_and_dispatch` returns, not merely soon after. A
real thread makes that untestable: `Thread.start()` frequently lets the
worker register before the main thread looks, so an assertion on the count
passes whether the code is correct or not. That is not hypothetical — the
first version of this file passed with the historical bug reinstated, and
was only caught by mutation-testing it.
With a thread that never runs on its own, the two implementations are
distinguishable with certainty: correct code has already registered, buggy
code has registered nothing.
"""
def __init__(self, target=None, args=(), daemon=None, **kw):
self._target, self._args = target, args
self.daemon = daemon
self.started = False
def start(self):
self.started = True
def run_now(self):
if self._target is not None:
self._target(*self._args)
def join(self, timeout=None):
return None
def is_alive(self):
return False
def test_reservation_is_recorded_before_dispatch_returns(monkeypatch):
"""The regression guard for the bug `_reserve_and_dispatch` was written to fix.
Deterministic: the worker thread is prevented from running at all, so the
count can only be 1 if the reservation happened synchronously in the
caller. Moving `_register_in_flight` back inside the thread fails this.
"""
monkeypatch.setattr(jw.threading, "Thread", _ManualThread)
threads: list = []
jw._reserve_and_dispatch(_job(1), job_threads=threads, run_in_thread=lambda job: None)
assert threads and threads[0].started, "dispatch did not start a worker thread"
assert jw._running_count() == 1, (
"capacity was not reserved before dispatch returned — the poll loop's "
"gate would read a stale count and oversubscribe"
)
assert jw._in_flight_ids() == [1]
def test_gate_stops_claiming_at_the_cap_even_if_no_job_has_started(monkeypatch):
"""Oversubscription guard, deterministic.
Worker threads never run, so nothing completes and nothing can unregister.
A correct worker reserves synchronously, so the gate sees the cap after
exactly `cap` claims and stops. A worker that registers inside the thread
sees 0 forever and would claim without bound — caught here by the iteration
ceiling rather than by hanging the suite.
"""
monkeypatch.setattr(jw.threading, "Thread", _ManualThread)
cap = 3
ceiling = cap * 10
threads: list = []
dispatched = 0
for jid in range(ceiling):
with jw._in_flight_lock:
live = len(jw._in_flight)
if live >= cap:
break
jw._reserve_and_dispatch(_job(jid), job_threads=threads, run_in_thread=lambda job: None)
dispatched += 1
assert dispatched == cap, (
f"gate admitted {dispatched} jobs against a cap of {cap} — "
"reservations are not visible to the poll loop in time"
)
assert jw._running_count() == cap
def test_the_mirrored_gate_still_matches_the_real_poll_loop():
"""These tests re-implement the poll loop's capacity check rather than
driving the loop itself, which is wrapped in HTTP. Pin the real expression
so the mirror cannot silently drift out of agreement with the source."""
import inspect
source = inspect.getsource(jw)
assert "live >= max_concurrent" in source, (
"the poll loop's capacity gate changed shape; the mirrored gate in "
"these tests must be updated to match"
)
def test_completion_frees_the_slot():
done = threading.Event()
threads: list = []
def run_in_thread(job):
jw._unregister_in_flight(int(job["id"]))
done.set()
jw._reserve_and_dispatch(_job(7), job_threads=threads, run_in_thread=run_in_thread)
assert done.wait(timeout=5), "job thread never ran"
for t in threads:
t.join(timeout=5)
assert jw._running_count() == 0
assert jw._in_flight_ids() == []
# --- N at once, never N+1 --------------------------------------------------
@pytest.mark.parametrize("cap", [2, 3, 5])
def test_jobs_really_do_run_simultaneously_up_to_the_cap(cap: int):
"""The positive direction: at max_concurrent=N, N jobs are in flight AT ONCE.
A serialising worker would keep the peak at 1 and time out on the barrier,
so this fails loudly rather than passing vacuously.
"""
barrier = threading.Barrier(cap, timeout=10)
release = threading.Event()
peak = 0
peak_lock = threading.Lock()
threads: list = []
barrier_broke = threading.Event()
def run_in_thread(job):
nonlocal peak
with peak_lock:
peak = max(peak, jw._running_count())
try:
barrier.wait() # only passes if `cap` jobs are alive together
except threading.BrokenBarrierError:
barrier_broke.set()
release.wait(timeout=5)
jw._unregister_in_flight(int(job["id"]))
for i in range(cap):
assert jw._running_count() < cap, "gate would have refused this claim"
jw._reserve_and_dispatch(_job(i), job_threads=threads, run_in_thread=run_in_thread)
assert jw._running_count() == cap
release.set()
for t in threads:
t.join(timeout=10)
assert not barrier_broke.is_set(), f"jobs did not overlap: only {peak} ran together"
assert peak == cap, f"expected {cap} concurrent jobs, peak was {peak}"
assert jw._running_count() == 0
def test_poller_never_exceeds_the_cap_while_workers_finish_concurrently():
"""The real concurrency check: one poll thread reserving, many threads
releasing, for many rounds. Models the live arrangement rather than an
artificial all-threads-reserve race, which the design never performs.
Any observation above the cap is a hard failure — recorded rather than
asserted inside the thread so the error surfaces on the main thread.
"""
cap = 4
total_jobs = 60
violations: list[int] = []
threads: list = []
next_id = iter(range(total_jobs))
def run_in_thread(job):
observed = jw._running_count()
if observed > cap:
violations.append(observed)
jw._unregister_in_flight(int(job["id"]))
dispatched = 0
deadline = time.monotonic() + 30
while dispatched < total_jobs:
assert time.monotonic() < deadline, (
f"dispatch loop stalled at {dispatched}/{total_jobs} — slots never freed"
)
# This mirrors the poll loop's gate verbatim (job_worker.py ~1647).
with jw._in_flight_lock:
live = len(jw._in_flight)
if live >= cap:
# Yield rather than busy-spin: the real loop sleeps 1.0s here, and
# a hot spin pins a core and starves the very threads we are
# waiting on (it pushed the full suite past its time budget).
time.sleep(0.001)
continue
if live > cap:
violations.append(live)
jid = next(next_id)
jw._reserve_and_dispatch(_job(jid), job_threads=threads, run_in_thread=run_in_thread)
dispatched += 1
for t in threads:
t.join(timeout=10)
assert not violations, f"capacity exceeded {cap}: observed {sorted(set(violations))}"
assert dispatched == total_jobs
assert jw._running_count() == 0
# --- what the broker is told ----------------------------------------------
def test_heartbeat_running_count_tracks_real_reservations():
"""`running` is what the broker stores and shows in `job workers`; it must
be the live in-flight count, not a constant."""
release = threading.Event()
threads: list = []
def run_in_thread(job):
release.wait(timeout=5)
jw._unregister_in_flight(int(job["id"]))
assert jw._running_count() == 0
for i in range(3):
jw._reserve_and_dispatch(_job(i), job_threads=threads, run_in_thread=run_in_thread)
assert jw._running_count() == i + 1
assert jw._in_flight_ids() == [0, 1, 2]
release.set()
for t in threads:
t.join(timeout=5)
assert jw._running_count() == 0
def test_allocations_total_sums_across_concurrent_jobs(monkeypatch):
"""At max_concurrent>1 the heartbeat must subtract EVERY pending claim, not
just the newest — otherwise the broker matches against VRAM the worker has
already promised away."""
monkeypatch.setattr(jw, "effective_vram_request_gb_from_job", lambda j: float(j["vram_gb"]))
for i, vram in enumerate([4.0, 6.0, 2.0]):
jw._register_in_flight({"id": i, "vram_gb": vram, "ram_gb": 1.5, "cpus": 2})
vram_total, ram_total, cpu_total = jw._allocations_total()
assert vram_total == pytest.approx(12.0)
assert ram_total == pytest.approx(4.5)
assert cpu_total == 6
# --- the setting itself ----------------------------------------------------
@pytest.mark.parametrize(
("value", "expected"),
[("1", 1), ("2", 2), ("16", 16)],
)
def test_max_concurrent_jobs_reads_the_env_var(monkeypatch, value, expected):
monkeypatch.setenv("JOBD_WORKER_MAX_CONCURRENT_JOBS", value)
assert jw._max_concurrent_jobs() == expected
@pytest.mark.parametrize("value", ["0", "-3", "", " ", "banana", "2.5"])
def test_max_concurrent_jobs_never_returns_a_useless_cap(monkeypatch, value):
"""A cap below 1 would wedge the poll loop forever (`live >= max` is true at
live=0), so garbage must fall back to a working worker, not a dead one."""
monkeypatch.setenv("JOBD_WORKER_MAX_CONCURRENT_JOBS", value)
assert jw._max_concurrent_jobs() >= 1