Reject dense DiT precisions on a CUDA-absent host before eviction; stabilize family-info tests

The start-route preflight caught the bf16-GPU and int8-torchao requirements but not the dense
precisions' CUDA requirement: on a GPU-less host bf16_unsupported_reason exempts CPU-only, so a
bf16/fp8 (or int8-with-torchao) DiT request passed the preflight, evicted resident workloads, then
raised only in the trainer child. Add the dense-mode CUDA gate mirroring _resolve_base_precision so
the doomed run is rejected up front. Also pin bf16_unsupported_reason in the two positive-path
family-info tests so they are deterministic across GPU types (a non-bf16 CUDA box would otherwise
empty every DiT family's advertised modes).

studio/backend/core/training/diffusion_train_common.py

@@ -277,21 +277,36 @@ def bf16_unsupported_reason(resolved_family: str) -> Optional[str]:
 def training_precision_preflight_error(resolved_family: str, base_precision: str) -> Optional[str]:
     """Reason the requested DiT precision cannot run on this host, else None -- checked by the
     start route BEFORE evicting resident GPU workloads (the trainer's own checks fire only in the
-    child, after eviction). Two gates: the bf16-GPU requirement (bf16_unsupported_reason), and an
-    explicit int8 needing a FUNCTIONAL torchao (its _int8_quantize_base has no fallback, so
-    _resolve_base_precision would otherwise raise only after eviction). Never raises."""
+    child, after eviction). Three gates, all mirroring _resolve_base_precision so a doomed run is
+    rejected before teardown: the bf16-GPU requirement (bf16_unsupported_reason); the dense
+    precisions (bf16/int8/fp8) requiring a CUDA GPU; and an explicit int8 needing a FUNCTIONAL
+    torchao (its _int8_quantize_base has no fallback). Never raises."""
     reason = bf16_unsupported_reason(resolved_family)
     if reason:
         return reason
-    if (
-        (resolved_family or "").strip().lower() in _DIT_TRAIN_FAMILIES
-        and (base_precision or "").strip().lower() == "int8"
-        and not has_functional_torchao()
-    ):
-        return (
-            "base_precision='int8' needs a functional torchao install; this host's torchao is "
-            "missing or the non-functional Windows-ROCm stub. Use 'nf4', 'bf16', or 'auto'."
-        )
+    fam = (resolved_family or "").strip().lower()
+    mode = (base_precision or "").strip().lower()
+    if fam in _DIT_TRAIN_FAMILIES and mode in ("bf16", "int8", "fp8"):
+        # The DiT trainer's dense precisions all require CUDA (_resolve_base_precision rejects
+        # bf16/int8/fp8 on device != "cuda"). bf16_unsupported_reason exempts a CPU-only host (the
+        # fp32 fallback for import/unit tests), so without this a dense request on a GPU-less host
+        # would pass the preflight, evict resident workloads, then raise only in the child.
+        try:
+            import torch
+
+            has_cuda = torch.cuda.is_available()
+        except Exception:  # noqa: BLE001 -- no torch / probe failure -> treat as no CUDA
+            has_cuda = False
+        if not has_cuda:
+            return (
+                f"base_precision={mode!r} needs a CUDA GPU; this host has none. "
+                "Use base_precision='nf4' or 'auto'."
+            )
+        if mode == "int8" and not has_functional_torchao():
+            return (
+                "base_precision='int8' needs a functional torchao install; this host's torchao is "
+                "missing or the non-functional Windows-ROCm stub. Use 'nf4', 'bf16', or 'auto'."
+            )
     return None
 
 

studio/backend/tests/test_diffusion_base_precision.py

@@ -100,6 +100,10 @@ def test_family_train_infos_drops_denied_fp8_for_qwen(monkeypatch):
     monkeypatch.setattr(
         common, "train_precision_modes", lambda: (["nf4", "bf16", "int8", "fp8", "auto"], "auto")
     )
+    # family_train_infos reads the live GPU via bf16_unsupported_reason; pin it to "bf16 OK" so
+    # this positive-path assertion is deterministic across GPU types (a non-bf16 CUDA box would
+    # otherwise empty every DiT family's modes). The empty-on-non-bf16 path is covered separately.
+    monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: None)
     infos = {i["name"]: i for i in common.family_train_infos()}
     assert "fp8" not in infos["qwen-image"]["precision_modes"]
     assert "int8" in infos["qwen-image"]["precision_modes"]  # int8 is fine on Qwen
@@ -187,6 +191,18 @@ def test_training_precision_preflight_error(monkeypatch):
     assert training_precision_preflight_error("sdxl", "int8") is None
     assert training_precision_preflight_error("", "int8") is None
 
+    # On a CUDA-ABSENT host, bf16_unsupported_reason exempts CPU-only, but the DiT trainer's dense
+    # precisions still require CUDA (mirroring _resolve_base_precision), so bf16/int8/fp8 for a DiT
+    # family are rejected UP FRONT rather than after eviction. nf4/auto (and SDXL) still pass.
+    monkeypatch.setattr(common, "has_functional_torchao", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    for dense in ("bf16", "int8", "fp8"):
+        reason = training_precision_preflight_error("flux.1", dense)
+        assert reason is not None and "CUDA" in reason
+    assert training_precision_preflight_error("flux.1", "nf4") is None
+    assert training_precision_preflight_error("flux.1", "auto") is None
+    assert training_precision_preflight_error("sdxl", "bf16") is None
+
 
 def test_family_train_infos_empties_dit_modes_on_non_bf16(monkeypatch):
     # On a non-bf16 GPU the start route rejects EVERY DiT family (even nf4), so /info must not
@@ -483,6 +499,9 @@ def test_family_train_infos_carries_precision_fields(monkeypatch):
     # Pin the machine probe so the DiT families carry a deterministic mode list, while SDXL
     # (no precision selector) stays empty regardless of the probe.
     monkeypatch.setattr(common, "train_precision_modes", lambda: (["nf4", "bf16"], "auto"))
+    # Also pin bf16_unsupported_reason (family_train_infos reads the live GPU through it): "bf16 OK"
+    # so this positive-path assertion is deterministic across GPU types, not just on CPU-only CI.
+    monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: None)
     infos = {i["name"]: i for i in common.family_train_infos()}
 
     flux = infos["flux.1"]