Serialize the GPU handoffs, gate DiT training on a GPU, and keep 3.9 installable

Six review findings, three of them evict-then-fail orderings:

- The chat load reclaimed the GPU without telling the arbiter it existed. A
  chat load holds no llama-server process until its GGUF has downloaded,
  which is minutes, so a competing Images/Video acquire in that window
  found nothing to cancel, took the GPU, and the chat load then spawned
  onto the same device. It now registers an in-flight marker through
  acquire_for's register hook (under the arbiter lock, as the image and
  video loads do), the evictor cancels a marked load, and the route undoes
  itself if ownership moved while it loaded.
- The Hub-download conflict check ran after that handoff, so a GGUF the
  download manager already owns destroyed the resident Images/Video
  pipeline and then 409'd, having loaded nothing. It moves above the
  handoff, together with the marker it handshakes with.
- The image load released the engine router's transition lock before
  registering the load, so a second load choosing the other engine could
  unload the still-idle engine this one captured; the load then landed on a
  deactivated engine, where generate, status, unload and the arbiter's
  evictor can no longer reach it. Registration now happens under that lock
  and refuses if the engine changed.
- Training a DiT family on a host with no GPU was accepted: nf4 is not a
  CPU fallback, its 4-bit load goes through bitsandbytes, which requires
  CUDA, XPU or MPS. The start unloaded the working Images pipeline, pulled
  the text encoders, and only then died in the child. Rejected before the
  teardown now, and /info stops advertising a precision that always 400s.
  SDXL keeps its documented fp32-on-CPU path.
- Both diffusion pages kept the routed-pick marker forever, so re-picking
  the same checkpoint (after chat evicted it) neither loaded nor cleared
  the query string. The marker is released once the query is gone. The
  Images key also carried a stray NUL byte, which made the file read as
  binary to grep and other tooling.
- diffusers dropped Python 3.9 in 0.38, so the unconditional >=0.39.0 pin
  left pip no candidate at all on 3.9 and made every install that composes
  the huggingface extras unresolvable there. The floor is conditional now.

Also fixes tests that were already red on the branch: two hand-built
request fakes had gone stale against fields this branch added, and the
handoff-ordering test only failed on a host with fewer than two GPUs.
This commit is contained in:
Daniel Han 2026-07-26 14:46:02 +00:00
commit bc00a8e797
14 changed files with 520 additions and 69 deletions

View file

@ -283,3 +283,48 @@ def test_activate_serializes_switch_and_concurrent_query(monkeypatch):
t.join(2.0)
q.join(2.0)
assert switch_done.is_set() and query_done.is_set()
def test_begin_load_on_refuses_an_engine_that_was_switched_away(monkeypatch):
# A load selects its engine, then yields (device probe, arbiter acquire) before registering.
# A second load choosing the OTHER engine transitions in that gap and unloads the still-idle
# engine this one captured, so registering there would leave a model that generate / status /
# unload and the arbiter's evictor can no longer reach (they all resolve the ACTIVE engine).
diffusers = SimpleNamespace(name = "diffusers")
sd_cpp = SimpleNamespace(name = "sd_cpp")
active = {"engine": diffusers}
monkeypatch.setattr(r, "get_active_diffusion_engine", lambda: active["engine"])
started: list[str] = []
assert r.begin_load_on(diffusers, lambda: started.append("ok") or "status") == "status"
assert started == ["ok"]
# A competing request switched the active engine after this one captured `diffusers`.
active["engine"] = sd_cpp
with pytest.raises(RuntimeError, match = "engine changed"):
r.begin_load_on(diffusers, lambda: started.append("leaked"))
assert started == ["ok"]
def test_begin_load_on_holds_the_transition_lock_while_registering(monkeypatch):
# The check and the registration must be one operation: taken under the same lock a switch
# takes, so no _activate can slip between them.
import threading
engine = SimpleNamespace(name = "diffusers")
monkeypatch.setattr(r, "get_active_diffusion_engine", lambda: engine)
inside = threading.Event()
release = threading.Event()
def _slow_start():
inside.set()
release.wait(2.0)
return "status"
t = threading.Thread(target = lambda: r.begin_load_on(engine, _slow_start))
t.start()
assert inside.wait(2.0)
assert r._transition_lock.locked()
release.set()
t.join(2.0)