36df317293
Comment-only pass over the Python this PR touches: drop what the code already says, collapse multi-line explanations that still read on one line, and keep the reasoning that is not recoverable from the code. No code, docstring semantics or behaviour changes; verified with an AST comparison against the previous revision, and the backend suite is unchanged (same 37 environment failures as before: the API integration tests that need a live keyed server, the flash-attn install hooks, and the GPU memory fields).
328 lines
13 KiB
Python
328 lines
13 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Tests for the diffusion engine router (diffusers vs native sd.cpp selection)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from core.inference import diffusion_engine_router as r
|
|
from core.inference.diffusion_families import detect_family
|
|
from core.inference.sd_cpp_engine import ENGINE_DIFFUSERS, ENGINE_SD_CPP
|
|
|
|
_ENVS = (
|
|
"UNSLOTH_DIFFUSION_ENGINE",
|
|
"UNSLOTH_DIFFUSION_SD_CPP",
|
|
"UNSLOTH_DIFFUSION_SD_CPP_MPS",
|
|
"UNSLOTH_DIFFUSION_SD_CPP_INSTALL",
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse = True)
|
|
def _clean_env_and_state(monkeypatch):
|
|
for e in _ENVS:
|
|
monkeypatch.delenv(e, raising = False)
|
|
# A light status-capable stub so neither selection nor active_status() imports the heavy
|
|
# diffusers/sd.cpp backends; the active engine NAME comes from module state.
|
|
monkeypatch.setattr(
|
|
r,
|
|
"get_active_diffusion_engine",
|
|
lambda: SimpleNamespace(status = lambda: {"loaded": False, "repo_id": None}),
|
|
)
|
|
# Default: no resident sd-server (so existing tests exercise the sd-cli path) and a stubbed
|
|
# runnability probe, so neither reaches the real install/exec path.
|
|
monkeypatch.setattr(r, "ensure_sd_server_binary", lambda **_: None)
|
|
monkeypatch.setattr(r, "_server_binary_runnable", lambda *_a, **_k: True)
|
|
yield
|
|
|
|
|
|
def _set_device(monkeypatch, backend):
|
|
monkeypatch.setattr(
|
|
r,
|
|
"resolve_diffusion_device_target",
|
|
lambda: SimpleNamespace(backend = backend, device = backend),
|
|
)
|
|
|
|
|
|
def _set_binary(monkeypatch, path):
|
|
monkeypatch.setattr(r, "ensure_sd_cpp_binary", lambda **_: path)
|
|
|
|
|
|
def _set_runnable(monkeypatch, version = "sd-cli v0"):
|
|
"""Stub the runnability probe so a stubbed binary path is treated as executable
|
|
(the router now probes ``SdCppEngine(...).version()`` before committing to native)."""
|
|
monkeypatch.setattr(r, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: version))
|
|
|
|
|
|
def _select(fam_name = "z-image"):
|
|
"""Activate the engine for a family and return which engine was chosen."""
|
|
r.select_and_activate_engine(detect_family(fam_name))
|
|
return r.active_engine_name()
|
|
|
|
|
|
# ── core selection matrix ─────────────────────────────────────────────────────
|
|
|
|
|
|
def test_cpu_with_binary_and_supported_family_picks_sd_cpp(monkeypatch):
|
|
_set_device(monkeypatch, "cpu")
|
|
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
|
_set_runnable(monkeypatch)
|
|
assert _select() == ENGINE_SD_CPP
|
|
assert r.active_engine_name() == ENGINE_SD_CPP
|
|
|
|
|
|
def test_cpu_with_only_sd_server_picks_sd_cpp(monkeypatch):
|
|
# An sd-server-only install (no runnable sd-cli) must still route to native: the backend prefers
|
|
# the resident server, so a runnable sd-server is native availability.
|
|
_set_device(monkeypatch, "cpu")
|
|
_set_binary(monkeypatch, None) # no sd-cli
|
|
monkeypatch.setattr(r, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: None))
|
|
monkeypatch.setattr(r, "ensure_sd_server_binary", lambda **_: "/usr/bin/sd-server")
|
|
assert _select() == ENGINE_SD_CPP
|
|
|
|
|
|
def test_present_but_not_runnable_binary_falls_back(monkeypatch):
|
|
# A binary that exists but cannot run must fall back to diffusers at selection, not commit native
|
|
# and fail inside the load.
|
|
_set_device(monkeypatch, "cpu")
|
|
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
|
monkeypatch.setattr(r, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: None))
|
|
assert _select() == ENGINE_DIFFUSERS
|
|
assert "binary unavailable" in (r.active_status()["fallback_reason"] or "")
|
|
|
|
|
|
@pytest.mark.parametrize("gpu", ["cuda", "rocm", "xpu"])
|
|
def test_gpu_backends_use_diffusers(monkeypatch, gpu):
|
|
_set_device(monkeypatch, gpu)
|
|
_set_binary(monkeypatch, "/usr/bin/sd-cli") # even with a binary, GPU stays diffusers
|
|
assert _select() == ENGINE_DIFFUSERS
|
|
assert "uses diffusers" in (r.active_status()["fallback_reason"] or "")
|
|
|
|
|
|
def test_forced_diffusers_overrides_cpu(monkeypatch):
|
|
_set_device(monkeypatch, "cpu")
|
|
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
|
monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "diffusers")
|
|
assert _select() == ENGINE_DIFFUSERS
|
|
assert "forced" in (r.active_status()["fallback_reason"] or "")
|
|
|
|
|
|
def test_sd_cpp_disabled_uses_diffusers(monkeypatch):
|
|
_set_device(monkeypatch, "cpu")
|
|
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
|
monkeypatch.setenv("UNSLOTH_DIFFUSION_SD_CPP", "0")
|
|
assert _select() == ENGINE_DIFFUSERS
|
|
assert "disabled" in (r.active_status()["fallback_reason"] or "")
|
|
|
|
|
|
def test_mps_default_diffusers_but_optin_sd_cpp(monkeypatch):
|
|
_set_device(monkeypatch, "mps")
|
|
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
|
_set_runnable(monkeypatch)
|
|
# Default: MPS is not native-eligible -> diffusers.
|
|
assert _select() == ENGINE_DIFFUSERS
|
|
# Opt in: MPS routes to sd.cpp.
|
|
monkeypatch.setenv("UNSLOTH_DIFFUSION_SD_CPP_MPS", "1")
|
|
assert _select() == ENGINE_SD_CPP
|
|
|
|
|
|
def test_unsupported_family_falls_back(monkeypatch):
|
|
_set_device(monkeypatch, "cpu")
|
|
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
|
monkeypatch.setattr(r, "family_sd_cpp_supported", lambda fam: False)
|
|
assert _select() == ENGINE_DIFFUSERS
|
|
assert "no native sd.cpp asset mapping" in (r.active_status()["fallback_reason"] or "")
|
|
|
|
|
|
def test_missing_binary_falls_back(monkeypatch):
|
|
_set_device(monkeypatch, "cpu")
|
|
_set_binary(monkeypatch, None) # install unavailable
|
|
assert _select() == ENGINE_DIFFUSERS
|
|
assert "binary unavailable" in (r.active_status()["fallback_reason"] or "")
|
|
|
|
|
|
def test_force_sd_cpp_on_gpu_when_binary_present(monkeypatch):
|
|
_set_device(monkeypatch, "cuda")
|
|
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
|
_set_runnable(monkeypatch)
|
|
monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "sd_cpp")
|
|
assert _select() == ENGINE_SD_CPP
|
|
|
|
|
|
def test_force_sd_cpp_without_binary_falls_back(monkeypatch):
|
|
_set_device(monkeypatch, "cuda")
|
|
_set_binary(monkeypatch, None)
|
|
monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "sd_cpp")
|
|
assert _select() == ENGINE_DIFFUSERS
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"backend, expected",
|
|
[("rocm", "rocm"), ("cuda", "cuda"), ("xpu", "vulkan"), ("cpu", "auto"), ("mps", "auto")],
|
|
)
|
|
def test_install_accelerator_maps_backend(backend, expected):
|
|
assert r._install_accelerator_for(backend) == expected
|
|
|
|
|
|
def test_force_native_install_uses_gpu_accelerator(monkeypatch):
|
|
# Forcing sd_cpp on a ROCm host with no binary must install the ROCm build, not the default CPU
|
|
# one, else the forced-native generation silently runs on CPU.
|
|
_set_device(monkeypatch, "rocm")
|
|
_set_runnable(monkeypatch)
|
|
seen = {}
|
|
|
|
def _fake_ensure(**kwargs):
|
|
seen.update(kwargs)
|
|
return "/usr/bin/sd-cli"
|
|
|
|
monkeypatch.setattr(r, "ensure_sd_cpp_binary", _fake_ensure)
|
|
monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "sd_cpp")
|
|
assert _select() == ENGINE_SD_CPP
|
|
assert seen.get("accelerator") == "rocm"
|
|
|
|
|
|
# ── active_status annotation ──────────────────────────────────────────────────
|
|
|
|
|
|
def test_active_status_injects_engine_and_reason(monkeypatch):
|
|
_set_device(monkeypatch, "cpu")
|
|
_set_binary(monkeypatch, None)
|
|
_select() # -> diffusers fallback (no binary)
|
|
st = r.active_status()
|
|
assert st["engine"] == ENGINE_DIFFUSERS
|
|
assert st["fallback_reason"] and "binary unavailable" in st["fallback_reason"]
|
|
|
|
|
|
# ── engine-switch eviction ordering ───────────────────────────────────────────
|
|
|
|
|
|
def test_switch_unloads_old_engine_before_publishing_new(monkeypatch):
|
|
# The arbiter's diffusion evictor unloads get_active_diffusion_engine(); if the router published
|
|
# the new (empty) engine BEFORE the old one finished unloading, a concurrent acquire could evict
|
|
# the empty engine and take the GPU while the old model was still resident. Assert the OLD engine
|
|
# stays the published target until its unload() completes.
|
|
seen = {}
|
|
|
|
def _fake_engine():
|
|
return SimpleNamespace(
|
|
unload = lambda: seen.__setitem__("active_during_unload", r.active_engine_name()),
|
|
status = lambda: {"loaded": False, "repo_id": None},
|
|
)
|
|
|
|
monkeypatch.setattr(r, "get_active_diffusion_engine", lambda: _fake_engine())
|
|
r._active_engine_name = ENGINE_SD_CPP
|
|
r._activate(ENGINE_DIFFUSERS, "switch test")
|
|
assert seen["active_during_unload"] == ENGINE_SD_CPP
|
|
assert r.active_engine_name() == ENGINE_DIFFUSERS
|
|
|
|
|
|
def test_no_switch_keeps_engine_and_refreshes_reason(monkeypatch):
|
|
# When the engine does not change, _activate must not spuriously unload anything and must still
|
|
# refresh the recorded fallback reason.
|
|
calls = {"unload": 0}
|
|
|
|
def _fake_engine():
|
|
return SimpleNamespace(
|
|
unload = lambda: calls.__setitem__("unload", calls["unload"] + 1),
|
|
status = lambda: {"loaded": False, "repo_id": None},
|
|
)
|
|
|
|
monkeypatch.setattr(r, "get_active_diffusion_engine", lambda: _fake_engine())
|
|
r._active_engine_name = ENGINE_DIFFUSERS
|
|
r._activate(ENGINE_DIFFUSERS, "still diffusers")
|
|
assert calls["unload"] == 0
|
|
assert r.active_engine_name() == ENGINE_DIFFUSERS
|
|
assert r.active_status()["fallback_reason"] == "still diffusers"
|
|
|
|
|
|
def test_activate_serializes_switch_and_concurrent_query(monkeypatch):
|
|
# Regression: without the transition lock a second _activate during the slow unload() reads the
|
|
# not-yet-updated active engine and returns it. Assert the query is blocked until the switch ends.
|
|
import threading
|
|
|
|
r._active_engine_name = ENGINE_DIFFUSERS
|
|
r._fallback_reason = None
|
|
|
|
release_unload = threading.Event()
|
|
unload_started = threading.Event()
|
|
|
|
def _slow_unload():
|
|
unload_started.set()
|
|
release_unload.wait(2.0)
|
|
|
|
engine = SimpleNamespace(status = lambda: {"loaded": False, "repo_id": None}, unload = _slow_unload)
|
|
monkeypatch.setattr(r, "get_active_diffusion_engine", lambda: engine)
|
|
|
|
switch_done = threading.Event()
|
|
|
|
def _switch():
|
|
r._activate(ENGINE_SD_CPP, None) # diffusers -> sd_cpp: unloads the old engine (blocks)
|
|
switch_done.set()
|
|
|
|
t = threading.Thread(target = _switch)
|
|
t.start()
|
|
assert unload_started.wait(2.0) # switch is mid-unload, holding the transition lock
|
|
|
|
query_done = threading.Event()
|
|
|
|
def _query():
|
|
r._activate(ENGINE_DIFFUSERS, None) # would hit the "no change" branch pre-fix
|
|
query_done.set()
|
|
|
|
q = threading.Thread(target = _query)
|
|
q.start()
|
|
# Serialized: the query cannot complete while the switch holds the transition lock.
|
|
assert not query_done.wait(0.4)
|
|
|
|
release_unload.set()
|
|
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 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)
|