From 89e5f69d90887c3ca80bde62ca00afc3ea89a9df Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 11:52:18 +0000 Subject: [PATCH] Don't prefetch dense shards for a prequant load; surface resolved provenance - _dense_quant_prefetch_needed widened the transformer/ prefetch to pull the base repo's full dense bf16 shards even when a prequant checkpoint is configured (candidate.prequant), contradicting its own docstring. That both defeats the prequant download savings and can hard-fail begin_load on a disk-full (no GGUF fallback there). Only widen for a real dense build (candidate is not None and not candidate.prequant). - DiffusionStatusResponse declared no 'resolved' field, so Pydantic's default extra='ignore' silently dropped the per-control auto-policy provenance the backend records (build_resolved_record / state.resolved) -- the plumbing never reached any client. Declare the field so it round-trips. --- studio/backend/core/inference/diffusion.py | 8 ++++- studio/backend/models/inference.py | 6 ++++ .../backend/tests/test_diffusion_backend.py | 32 ++++++++++++++++--- 3 files changed, 41 insertions(+), 5 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 1ae3ce2668..ac61efccf8 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -456,7 +456,13 @@ class DiffusionBackend: prequant_path = kwargs.get("transformer_prequant_path"), logger = None, ) - return candidate is not None + # A prequant candidate loads from the small pre-quantized checkpoint (+ config / + # companions), NOT the base repo's full dense transformer/ shards, so widening the + # prefetch to pull those shards both defeats the prequant download savings and can + # hard-fail the load: the widened pull runs in begin_load, where a disk-full has no + # GGUF fallback (unlike the in-load_pipeline dense failure). Only widen for a real + # dense build. + return candidate is not None and not candidate.prequant except Exception: # noqa: BLE001 — widening the prefetch is best-effort only return False diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 8b5f3602d6..93e8dd96a5 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -2131,6 +2131,12 @@ class DiffusionStatusResponse(BaseModel): "picker's enabled state). Diffusers only, for families with a ControlNet pipeline; False " "for the native engine, GGUF-via-diffusers, and torchao fp8/int8 dense.", ) + resolved: Optional[dict] = Field( + None, + description = "Per-control auto-policy provenance (value/source/reason for each resolved " + "setting), or null. Declared explicitly so the field is not dropped by the default " + "extra='ignore', which would silently discard the backend's resolved record.", + ) # ── OpenAI-compatible images API (POST /v1/images/generations) ── diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index f210c73b9e..0fbfe9f32b 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -2115,7 +2115,9 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): logger = None, ): seen.append(requested) - return object() # a viable dense-quant candidate (scheme resolves AND disk fits) + # A real (non-prequant) dense-quant candidate: scheme resolves AND disk fits, so the + # loader takes the dense build that needs the base repo's bf16 transformer/ shards. + return types.SimpleNamespace(prequant = False) monkeypatch.setattr(dmod, "resolve_dense_quant_candidate", fake_candidate) @@ -2132,13 +2134,35 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8", "speed_mode": "off"}) is False ) - # No viable candidate (unsupported scheme / no disk room / a prequant checkpoint shortcut) - # -> never widen. The disk guard here is exactly what averts filling the cache volume and - # hard-failing the load instead of falling back to the GGUF. + # A PREQUANT candidate loads the small pre-quantized checkpoint (+ config / companions), + # NOT the base repo's dense transformer/ shards, so the widened prefetch must NOT fire -- + # otherwise it defeats the prequant download savings and can hard-fail begin_load (no GGUF + # fallback there) on a disk-full. + monkeypatch.setattr( + dmod, "resolve_dense_quant_candidate", lambda **kw: types.SimpleNamespace(prequant = True) + ) + assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False + # No viable candidate at all (unsupported scheme / no disk room) -> never widen. The disk + # guard here is exactly what averts filling the cache volume and hard-failing the load + # instead of falling back to the GGUF. monkeypatch.setattr(dmod, "resolve_dense_quant_candidate", lambda **kw: None) assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False +def test_diffusion_status_response_carries_resolved(): + # The backend records per-control auto-policy provenance (build_resolved_record) on + # state.resolved; the response model must DECLARE the field or Pydantic's default + # extra='ignore' silently drops it, leaving that plumbing dead (never reaching a client). + from models.inference import DiffusionStatusResponse + + rec = {"transformer_quant": {"value": "fp8", "source": "auto", "reason": "blackwell"}} + resp = DiffusionStatusResponse(loaded = True, resolved = rec) + assert resp.resolved == rec + assert resp.model_dump()["resolved"] == rec + # Absent by default (nothing resolved / native engine). + assert DiffusionStatusResponse(loaded = False).resolved is None + + def test_companion_cache_bytes_local_dir_excludes_transformer(tmp_path): # A LOCAL diffusers base: sum the on-disk VAE / text-encoder weights so auto memory # planning sees the resident companions, but exclude transformer/ (the GGUF supplies