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

@@ -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

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"))

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.

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,

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