From 4e32ef64847aa8fc15c0ef9dda17fd8c93190821 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 26 Jul 2026 21:41:28 +0000 Subject: [PATCH] Pin the accelerator probes in the DiT family-metadata tests Six tests read family_train_infos() (or a start preflight) without pinning the host probes, so they only held on a machine with a bf16 accelerator: on a GPU-less runner the DiT gate empties precision_modes, turns supports_compile off, and replaces any other preflight message with the no-accelerator note, and all six failed there. A conftest fixture pins both probes for exactly those tests, so they assert the family metadata they are about on every host. The gate's own CPU-only behaviour keeps its dedicated tests. --- studio/backend/tests/conftest.py | 17 +++++++++++++++++ .../tests/test_diffusion_base_precision.py | 4 ++-- .../backend/tests/test_diffusion_dit_trainer.py | 4 ++-- studio/backend/tests/test_diffusion_krea2.py | 2 +- studio/backend/tests/test_diffusion_training.py | 2 +- 5 files changed, 23 insertions(+), 6 deletions(-) diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py index 3226ea2823..9bfe0fc47d 100644 --- a/studio/backend/tests/conftest.py +++ b/studio/backend/tests/conftest.py @@ -182,3 +182,20 @@ def stub_embeddings(monkeypatch): ) monkeypatch.setattr(embeddings, "warm", lambda model_name = None: None) return dim + + +@pytest.fixture +def dit_train_host(monkeypatch): + """Pretend this host can train the DiT families. + + ``family_train_infos()`` and the start preflight both gate on the accelerator / bf16 probes, so + on a GPU-less runner every DiT family reports no precision modes, ``supports_compile`` False, + and a "needs a GPU" note that replaces any other preflight message. Tests about family metadata + or about a different preflight pin the probes here so they assert the same thing on every host; + the gate itself is covered by its own tests in test_diffusion_base_precision.py. + """ + import core.training.diffusion_train_common as dtc + + monkeypatch.setattr(dtc, "dit_accelerator_missing_reason", lambda *_a, **_k: None) + monkeypatch.setattr(dtc, "bf16_unsupported_reason", lambda *_a, **_k: None) + return dtc diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index 5193093986..abaaf2f427 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -92,7 +92,7 @@ def test_base_precision_denies_fp8_for_corrupted_family(): assert flux.base_precision == "fp8" -def test_family_train_infos_drops_denied_fp8_for_qwen(monkeypatch): +def test_family_train_infos_drops_denied_fp8_for_qwen(monkeypatch, dit_train_host): # /info advertises the machine's DiT modes per family, but a family whose DiT the mode corrupts # must not offer it, so the UI never surfaces a mode normalized() would reject. monkeypatch.setattr( @@ -528,7 +528,7 @@ def test_train_precision_modes_gates_dense_on_bf16_support(monkeypatch): # ── family_train_infos precision fields ─────────────────────────────────────── -def test_family_train_infos_carries_precision_fields(monkeypatch): +def test_family_train_infos_carries_precision_fields(monkeypatch, dit_train_host): # Pin the machine probe so the DiT families carry a deterministic mode list, while SDXL (no # precision selector) stays empty regardless. monkeypatch.setattr(common, "train_precision_modes", lambda: (["nf4", "bf16"], "auto")) diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index 4c2848e10a..ba724b28d3 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -194,7 +194,7 @@ def test_gated_access_requires_token(): _assert_gated_access("black-forest-labs/FLUX.2-klein-4B", None) # Klein is open -def test_family_train_infos_lists_dit_families(): +def test_family_train_infos_lists_dit_families(dit_train_host): infos = {i["name"]: i for i in family_train_infos()} for fam in ("sdxl", "flux.1", "qwen-image", "z-image", "flux.2-klein", "flux.2-dev"): assert fam in infos, f"{fam} missing from family_train_infos" @@ -213,7 +213,7 @@ def test_family_train_infos_lists_dit_families(): assert "4bit" in infos["z-image"]["default_base"].lower() -def test_family_train_infos_sdxl_supports_compile_without_precision_modes(monkeypatch): +def test_family_train_infos_sdxl_supports_compile_without_precision_modes(monkeypatch, dit_train_host): # Regional compile now applies to every family (the SDXL trainer compiles its U-Net blocks too), # but base_precision stays DiT-only, so SDXL advertises no precision modes while z-image keeps # its own. Pin the precision list so the assertion holds regardless of the host GPU. diff --git a/studio/backend/tests/test_diffusion_krea2.py b/studio/backend/tests/test_diffusion_krea2.py index 6f1728e672..0981189c3b 100644 --- a/studio/backend/tests/test_diffusion_krea2.py +++ b/studio/backend/tests/test_diffusion_krea2.py @@ -166,7 +166,7 @@ def test_krea2_family_wiring(): # ── training wiring ────────────────────────────────────────────────────────── -def test_krea2_training_registry(): +def test_krea2_training_registry(dit_train_host): from core.inference.diffusion_families import trainable_family_names from core.training.diffusion_train_common import ( family_train_infos, diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 6264f416f6..ed8aa10156 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -1568,7 +1568,7 @@ def test_info_lists_trainable_families(client): assert families["z-image"]["defaults"]["resolution"] == 768 -def test_start_gated_base_without_access_is_400_and_keeps_gpu(client, monkeypatch): +def test_start_gated_base_without_access_is_400_and_keeps_gpu(client, monkeypatch, dit_train_host): # A gated FLUX base with no valid token must 400 from the HEAD preflight BEFORE the GPU residents # are freed, so a doomed start never evicts the user's loaded model. import urllib.error