diff --git a/studio/backend/hub/services/models/downloads.py b/studio/backend/hub/services/models/downloads.py index d0db367d96..19460950b9 100644 --- a/studio/backend/hub/services/models/downloads.py +++ b/studio/backend/hub/services/models/downloads.py @@ -6,7 +6,6 @@ from __future__ import annotations import asyncio -import hashlib from typing import Optional, Sequence, TYPE_CHECKING from fastapi import HTTPException @@ -52,22 +51,9 @@ def _download_job_key(repo_id: str, variant: Optional[str]) -> str: _SCOPE_PREFIX = "@" -def _scope_variant(scope_id: Optional[str], files: Optional[Sequence[str]] = None) -> Optional[str]: - """The variant slot for a scoped job: ``@scope`` plus the requested FILE SET. - - The file set has to be part of the identity. Every image download for a repo would - otherwise share ``@diffusion`` (and every video one ``@video``), so switching quant - while the first scoped job is still running made the second request adopt it: the UI - waited on the first file set and then loaded a file that was never fetched. - """ +def _scope_variant(scope_id: Optional[str]) -> Optional[str]: scope = (scope_id or "").strip() - if not scope: - return None - wanted = sorted({f.strip() for f in (files or []) if f and f.strip()}) - if not wanted: - return f"{_SCOPE_PREFIX}{scope}" - digest = hashlib.sha256("\n".join(wanted).encode()).hexdigest()[:12] - return f"{_SCOPE_PREFIX}{scope}-{digest}" + return f"{_SCOPE_PREFIX}{scope}" if scope else None def scoped_file_blob_hashes( @@ -170,7 +156,7 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional ) # A scoped job fetches only `files` and keys itself apart from the repo's full snapshot. scoped_files = [f for f in (body.files or []) if f and f.strip()] - scope_variant = _scope_variant(body.scope_id, scoped_files) + scope_variant = _scope_variant(body.scope_id) if scope_variant is not None: if variant is not None: raise HTTPException( @@ -254,6 +240,15 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional if not claimed: if claim_state == "admission_blocked": raise _load_in_flight_error(repo_id) + if claim_state == "scope_file_mismatch": + raise HTTPException( + status_code = 409, + detail = ( + f"Another download for '{repo_id}' is already fetching a different " + "set of files. Wait for it to finish (or cancel it), then start " + "this one." + ), + ) # claim_state is the blocking job's state. The client can attach only # when the blocker is this key's own in-flight job (adoptable); a # cross-variant conflict or in-progress delete is not accepted. diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index fa5862a13c..a732c69fc0 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -51,6 +51,24 @@ def _denylist_inert(monkeypatch): monkeypatch.setattr(folder_browser, "is_denied_system_path", lambda _p: False) +def _download_body(**over) -> SimpleNamespace: + """A download request with every field the route reads. + + Built by hand rather than through the schema so these tests stay cheap, so it lists + the newer scoping fields too: leaving one off reads as AttributeError inside the + route, not as a missing default. + """ + body = { + "repo_id": "Org/Model", + "gguf_variant": None, + "use_xet": False, + "scope_id": None, + "files": None, + } + body.update(over) + return SimpleNamespace(**body) + + def _repo(repo_id: str, files: list[SimpleNamespace], repo_path: Path): return SimpleNamespace( repo_id = repo_id, @@ -3042,7 +3060,7 @@ def test_model_download_records_completed_baseline_for_new_gguf_variant(monkeypa asyncio.run( downloads.download_model_response( - SimpleNamespace(repo_id = "Org/Model", gguf_variant = "Q4_K_M", use_xet = False) + _download_body(gguf_variant = "Q4_K_M") ) ) @@ -3128,7 +3146,7 @@ def test_gguf_model_download_skips_completed_baseline_for_variant_resume_state( asyncio.run( downloads.download_model_response( - SimpleNamespace(repo_id = "Org/Model", gguf_variant = "Q4_K_M", use_xet = False) + _download_body(gguf_variant = "Q4_K_M") ) ) @@ -3329,7 +3347,7 @@ def test_model_claim_register_cancel_uses_registry_marker_owner(monkeypatch): result = asyncio.run( downloads.download_model_response( - SimpleNamespace(repo_id = "Org/Model", gguf_variant = None, use_xet = False) + _download_body() ) ) @@ -3445,7 +3463,7 @@ def test_model_download_watcher_invalidates_hf_cache_scan(monkeypatch): result = asyncio.run( downloads.download_model_response( - SimpleNamespace(repo_id = "Org/Model", gguf_variant = None, use_xet = False) + _download_body() ) ) @@ -3513,18 +3531,10 @@ def test_two_concurrent_same_repo_variants_both_complete(monkeypatch, tmp_path): async def _run_both(): return await asyncio.gather( downloads.download_model_response( - SimpleNamespace( - repo_id = "Org/Model", - gguf_variant = "Q4_K_M", - use_xet = False, - ) + _download_body(gguf_variant = "Q4_K_M") ), downloads.download_model_response( - SimpleNamespace( - repo_id = "Org/Model", - gguf_variant = "Q8_0", - use_xet = False, - ) + _download_body(gguf_variant = "Q8_0") ), ) diff --git a/studio/backend/hub/utils/download_registry.py b/studio/backend/hub/utils/download_registry.py index 2f8a91b99c..10ec3aacbc 100644 --- a/studio/backend/hub/utils/download_registry.py +++ b/studio/backend/hub/utils/download_registry.py @@ -1155,6 +1155,17 @@ class DownloadRegistry: return False, conflict_state current = self._jobs.get(key, DownloadState("idle")).state if current in _ACTIVE_STATES and not replace_active: + # A scope slot is shared by every file set that rides it (two quants of + # one repo both key as "@diffusion"), so adopting the live job would let + # the caller wait on files it never asked for. Reject instead; checked + # here, under the lock, so a concurrent claim cannot slip past it. + live = self._metadata.get(key) + if ( + scoped_files is not None + and live is not None + and sorted(set(live.scoped_files)) != sorted(set(scoped_files)) + ): + return False, "scope_file_mismatch" return False, current if generation is None: self._generation_seq += 1 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 43d8257839..1d8cd4a17f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -16001,6 +16001,15 @@ async def diffusion_download_plan( hf_token = request.hf_token, transformer_quant = request.transformer_quant, speed_mode = request.speed_mode, + # The dense-quant prefetch decision reads the memory policy, the prequant path + # and the adapter selection too (an offload policy never runs the dense build; + # a baked LoRA always does), so the plan has to see the same values the load + # will. Without them it stages the base transformer/ shards for a low-VRAM + # load that never opens them, or omits them for a baked-LoRA load that does. + memory_mode = request.memory_mode, + cpu_offload = request.cpu_offload, + transformer_prequant_path = request.transformer_prequant_path, + loras = request.loras, ) return DiffusionDownloadPlanResponse(**plan) except (ValueError, FileNotFoundError) as exc: diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index a7ac02d1f2..9e7306e577 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -1002,3 +1002,43 @@ def test_status_resolved_defaults_to_null(client): # backends and the unloaded state). body = client.get("/api/inference/images/status").json() assert body["resolved"] is None + + +def test_download_plan_forwards_the_load_time_controls(client, monkeypatch): + # The plan drives the staged download, so it has to be computed from the SAME + # configuration the load will run with. The dense-quant prefetch decision reads the + # memory policy, the prequant path and the adapter selection as well as speed/quant: + # dropping them stages the base transformer/ shards for a low-VRAM load that never + # opens them, or omits them for a baked-LoRA load that does. + backend = diffusion_module.get_diffusion_backend() + seen: dict = {} + + def _plan(model_path, **kwargs): + seen["model_path"] = model_path + seen.update(kwargs) + return {"entries": [], "total_bytes": 0} + + monkeypatch.setattr(backend, "download_plan", _plan, raising = False) + + resp = client.post( + "/api/inference/images/download-plan", + json = { + "model_path": "unsloth/FLUX.1-dev-GGUF", + "gguf_filename": "flux1-dev-Q4_K_M.gguf", + "model_kind": "gguf", + "hf_token": "hf_secret", + "speed_mode": "off", + "transformer_quant": "int8", + "memory_mode": "low_vram", + "cpu_offload": True, + "loras": [{"id": "unsloth/some-lora", "weight": 0.8}], + }, + ) + + assert resp.status_code == 200 + assert seen["hf_token"] == "hf_secret" + assert seen["speed_mode"] == "off" + assert seen["transformer_quant"] == "int8" + assert seen["memory_mode"] == "low_vram" + assert seen["cpu_offload"] is True + assert len(seen["loras"] or []) == 1 diff --git a/studio/backend/tests/test_scoped_download_job.py b/studio/backend/tests/test_scoped_download_job.py index d72489154a..dff692ace6 100644 --- a/studio/backend/tests/test_scoped_download_job.py +++ b/studio/backend/tests/test_scoped_download_job.py @@ -22,6 +22,8 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) +from fastapi import HTTPException + from hub.schemas.downloads import DownloadModelRequest from hub.services import download_lifecycle from hub.services.models import downloads as dl @@ -91,8 +93,7 @@ def test_scoped_start_spawns_a_file_scoped_worker(monkeypatch): result = asyncio.run(dl.download_model_response(_request())) assert result["accepted"] is True - # The scope key carries a digest of the requested file set (see _scope_variant). - scope_variant = dl._scope_variant("diffusion", FILES) + scope_variant = dl._scope_variant("diffusion") assert result["job_key"].endswith(scope_variant) args = spawned["args"] @@ -123,7 +124,7 @@ def test_scoped_files_survive_into_the_registry(monkeypatch): assert captured["scoped_files"] == FILES metadata = dl._registry.get_job_metadata( - dl._download_job_key("black-forest-labs/FLUX.1-dev", dl._scope_variant("diffusion", FILES)) + dl._download_job_key("black-forest-labs/FLUX.1-dev", dl._scope_variant("diffusion")) ) assert metadata is not None and list(metadata.scoped_files) == FILES @@ -136,20 +137,43 @@ def test_files_manifest_round_trips(): Path(path).unlink(missing_ok = True) -def test_scope_keys_differ_per_requested_file_set(): - # Two quants of one repo are two different downloads. Sharing "@diffusion" made the - # second request adopt the first job, so the UI waited on the wrong file set and then - # loaded a file that had never been fetched. - a = dl._scope_variant("diffusion", ["flux1-dev-Q4_K_M.gguf", "ae.safetensors"]) - b = dl._scope_variant("diffusion", ["flux1-dev-Q2_K.gguf", "ae.safetensors"]) - assert a != b - assert a.startswith("@diffusion-") and b.startswith("@diffusion-") - # Order and duplicates must not change the identity (the same set is the same job). - assert a == dl._scope_variant( - "diffusion", ["ae.safetensors", "flux1-dev-Q4_K_M.gguf", "flux1-dev-Q4_K_M.gguf"] - ) - # A scope with no files keeps the bare form, and stays valid as a variant slot. - from hub.utils.paths import is_valid_gguf_variant +def test_a_different_file_set_is_not_adopted(monkeypatch): + # Two quants of one repo are two different downloads that share the "@diffusion" slot. + # Adopting the running one made the UI wait on the wrong file set and then load a file + # that had never been fetched, so the second request is refused while the first runs. + monkeypatch.setattr(dl, "_reject_if_load_in_flight", lambda repo_id: None) + monkeypatch.setattr(dl, "resolve_cached_repo_id_case", lambda repo, **k: repo) + monkeypatch.setattr(dl, "scoped_file_blob_hashes", lambda *a, **k: frozenset()) + monkeypatch.setattr(download_lifecycle, "launch_worker", lambda *a, **k: "running") - assert dl._scope_variant("video", []) == "@video" - assert is_valid_gguf_variant(a) and is_valid_gguf_variant("@video") + key = dl._download_job_key("black-forest-labs/FLUX.1-dev", dl._scope_variant("diffusion")) + try: + first = asyncio.run(dl.download_model_response(_request())) + assert first["accepted"] is True + + with pytest.raises(HTTPException) as other_files: + asyncio.run( + dl.download_model_response( + _request(files = ["model_index.json", "flux1-dev-Q2_K.gguf"]) + ) + ) + assert other_files.value.status_code == 409 + assert "different" in other_files.value.detail + + # The same file set is still the same download: it adopts the live job as before, + # in any order and with duplicates collapsed. + same = asyncio.run( + dl.download_model_response(_request(files = [FILES[1], FILES[0], FILES[0]])) + ) + assert same["accepted"] is True and same["job_key"] == key + finally: + dl._registry.set_job(key, "complete") + + +def test_scope_key_stays_derivable_from_the_scope_alone(): + # The download manager builds this key client-side (it polls and cancels before any + # server round-trip tells it a key), so the scope name alone must produce it. + assert dl._scope_variant("diffusion") == "@diffusion" + assert dl._scope_variant("video") == "@video" + assert dl._scope_variant(" ") is None + assert is_valid_gguf_variant("@diffusion") diff --git a/studio/frontend/src/features/hub/download-manager/download-manager-state.ts b/studio/frontend/src/features/hub/download-manager/download-manager-state.ts index c24bb1e154..c78839b30f 100644 --- a/studio/frontend/src/features/hub/download-manager/download-manager-state.ts +++ b/studio/frontend/src/features/hub/download-manager/download-manager-state.ts @@ -71,6 +71,10 @@ function sanitizePersistedJob(value: unknown): ManagedDownload | null { ...(Number.isSafeInteger(value.serverGeneration) ? { serverGeneration: Number(value.serverGeneration) } : {}), + ...(Array.isArray(value.scopedFiles) && + value.scopedFiles.every((f) => typeof f === "string") + ? { scopedFiles: value.scopedFiles as string[] } + : {}), }; } @@ -109,6 +113,7 @@ function toPersistedJob( ...(job.serverGeneration !== undefined ? { serverGeneration: job.serverGeneration } : {}), + ...(job.scopedFiles !== undefined ? { scopedFiles: job.scopedFiles } : {}), }; } diff --git a/studio/frontend/src/features/hub/download-manager/download-manager-types.ts b/studio/frontend/src/features/hub/download-manager/download-manager-types.ts index 8f45320c83..4eed53f8e4 100644 --- a/studio/frontend/src/features/hub/download-manager/download-manager-types.ts +++ b/studio/frontend/src/features/hub/download-manager/download-manager-types.ts @@ -25,6 +25,14 @@ export interface ManagedDownload { error: string | null; startedAt: number; serverGeneration?: number; + /** + * Files a scoped job is fetching, when known. Every file set of one repo rides the same + * scope slot (see `scopedVariant`), so this is what separates "my transfer is already + * running" from "a different quant of this repo is running": adopting the latter would + * report ready for files nobody fetched. Absent means unknown (an unscoped job, or one + * hydrated by an older build), and unknown stays adoptable as before. + */ + scopedFiles?: string[]; } export interface DownloadRequest { diff --git a/studio/frontend/src/features/hub/download-manager/poll-loop.ts b/studio/frontend/src/features/hub/download-manager/poll-loop.ts index c01abf6840..bb911e4a37 100644 --- a/studio/frontend/src/features/hub/download-manager/poll-loop.ts +++ b/studio/frontend/src/features/hub/download-manager/poll-loop.ts @@ -677,6 +677,14 @@ export async function startJob( ...(Number.isSafeInteger(seedGeneration) ? { serverGeneration: seedGeneration } : {}), + // Recorded so a later start for the same scope slot can tell whether this running + // job is fetching its files or a different quant's. An adopted job keeps whatever + // the existing record knew. + ...(req.files && req.files.length > 0 + ? { scopedFiles: [...req.files] } + : opts.adopt && existing?.scopedFiles + ? { scopedFiles: existing.scopedFiles } + : {}), }); if (!opts.adopt) { diff --git a/studio/frontend/src/features/hub/download-manager/transport-conflict.ts b/studio/frontend/src/features/hub/download-manager/transport-conflict.ts index 3267a6ae05..36d4b62d6e 100644 --- a/studio/frontend/src/features/hub/download-manager/transport-conflict.ts +++ b/studio/frontend/src/features/hub/download-manager/transport-conflict.ts @@ -9,7 +9,10 @@ import { apiTransportStatusWithRetry, effectiveTransportMode, } from "./download-api-adapter"; -import type { DownloadRequest } from "./download-manager-types"; +import type { + DownloadRequest, + ManagedDownload, +} from "./download-manager-types"; import { findActiveJobForRepo, getState, @@ -96,7 +99,24 @@ export type DownloadStartOutcome = "started" | "conflict" | "busy" | "error"; // callers never claim a download began when it did not. function isJobActiveFor(req: DownloadRequest): boolean { const job = getState().jobs[jobKeyOf(req.kind, req.repoId, req.variant)]; - return Boolean(job && ACTIVE_STATES.has(job.state)); + if (!job || !ACTIVE_STATES.has(job.state)) return false; + return !scopedFileSetDiffers(job, req); +} + +// Every file set of one repo rides the same scope slot, so a live job on this key counts +// as this request's transfer only when it is fetching the same files: adopting a sibling +// quant's job would report ready for files nobody fetched. A job with no recorded list +// (unscoped, or hydrated by an older build) stays adoptable as before. +function scopedFileSetDiffers( + job: ManagedDownload, + req: DownloadRequest, +): boolean { + if (!job.scopedFiles || !req.files || req.files.length === 0) return false; + const live = [...new Set(job.scopedFiles)].sort(); + const wanted = [...new Set(req.files)].sort(); + return ( + live.length !== wanted.length || live.some((f, i) => f !== wanted[i]) + ); } async function runWithPendingStartGuard( diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 0fb17014c7..251c2d0fdf 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1791,6 +1791,17 @@ export function ImagesPage({ active = true }: { active?: boolean }) { model_path: repoId, gguf_filename: opts.filename, model_kind: opts.kind, + // The same token and Advanced values handleLoad sends, so the plan describes the + // load that will actually run. Without the token the backend's metadata lookup + // fails on a gated base and quietly plans no companion entry, leaving the load + // to pull those multi-GB files inline, outside this manager; without the memory + // and quant controls an explicit Speed/Precision Off or low-VRAM pick stages + // base transformer/ shards the load never opens. + hf_token: hfApiToken(getHfToken()), + cpu_offload: cpuOffload, + speed_mode: speedMode === "auto" ? undefined : speedMode, + transformer_quant: transformerQuant === "auto" ? undefined : transformerQuant, + memory_mode: memoryMode === "auto" ? undefined : memoryMode, }); if (plan.entries.length > 0) { pendingStagedLoad.current = { repoId, opts }; @@ -1809,7 +1820,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { } return handleLoadRef.current(repoId, opts); }, - [stage], + [stage, cpuOffload, speedMode, transformerQuant, memoryMode], ); // A diffusion model picked from the chat picker arrives as ?model= on this route. Load it @@ -1821,6 +1832,11 @@ export function ImagesPage({ active = true }: { active?: boolean }) { const navigateSelf = useNavigate(); const handledRouteModel = useRef(null); useEffect(() => { + // Only the page being shown consumes the query. This hook is loose (`strict: false`) + // and both diffusion pages stay mounted once visited, so the hidden one saw + // /video?model= too and raced this one: it navigated back to its own route and tried + // to load the other page's checkpoint as its own kind of model. + if (!active) return; const wanted = routeSearch.model; // Key on the model AND the quant: this page stays mounted across navigation, so a // model-only marker made every later pick of the same repo a no-op -- including a @@ -1837,7 +1853,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { : { kind: "pipeline" }, false, ); - }, [routeSearch.model, routeSearch.quant, loadOrStage, navigateSelf]); + }, [active, routeSearch.model, routeSearch.quant, loadOrStage, navigateSelf]); // Reload the current model with the current advanced options. const handleReapply = useCallback(() => { diff --git a/studio/frontend/src/features/video/video-page.tsx b/studio/frontend/src/features/video/video-page.tsx index 71636b38ea..372b698fe0 100644 --- a/studio/frontend/src/features/video/video-page.tsx +++ b/studio/frontend/src/features/video/video-page.tsx @@ -551,6 +551,17 @@ export function VideoPage({ active = true }: { active?: boolean }) { useEffect(() => { playCountRef.current = 0; }, [selectedId]); + // Pause the preview when this page stops being the visible one. The keep-alive layout only + // hides it, and display:none does not pause a media element, so a clip the user unmuted + // would keep decoding and playing its audio over whatever page they opened next. + const previewRef = useRef(null); + // The media element's own handlers fire while the page is hidden, so they read `active` + // through a ref rather than closing over a stale render's value. + const activeRef = useRef(active); + useEffect(() => { + activeRef.current = active; + if (!active) previewRef.current?.pause(); + }, [active]); const [srcById, setSrcById] = useState>(() => Object.fromEntries(galleryCache.srcById), ); @@ -1119,6 +1130,10 @@ export function VideoPage({ active = true }: { active?: boolean }) { model_path: repoId, gguf_filename: opts.filename, model_kind: opts.kind, + // Same token handleLoad sends: without it the backend's metadata lookup fails on + // a gated base and the plan quietly comes back without the companion entry, so + // the load pulls those multi-GB files inline, outside this manager. + hf_token: hfApiToken(getHfToken()), }); if (plan.entries.length > 0) { pendingStagedLoad.current = { repoId, opts }; @@ -1149,6 +1164,11 @@ export function VideoPage({ active = true }: { active?: boolean }) { const navigateSelf = useNavigate(); const handledRouteModel = useRef(null); useEffect(() => { + // Only the page being shown consumes the query. This hook is loose (`strict: false`) + // and both diffusion pages stay mounted once visited, so the hidden one saw + // /images?model= too and raced that page: it navigated back to /video and tried to + // load an image checkpoint as a video model. + if (!active) return; const wanted = routeSearch.model; // Model AND quant, as on the Images page: this page stays mounted, so a model-only // marker turned every later pick of the same repo into a silent no-op. @@ -1163,7 +1183,7 @@ export function VideoPage({ active = true }: { active?: boolean }) { : { kind: "pipeline" }, false, ); - }, [routeSearch.model, routeSearch.quant, loadOrStage, navigateSelf]); + }, [active, routeSearch.model, routeSearch.quant, loadOrStage, navigateSelf]); // Reload the current model with the current advanced options. const handleReapply = useCallback(() => { @@ -1661,6 +1681,7 @@ export function VideoPage({ active = true }: { active?: boolean }) { it. The counter resets per selection (`key` remounts the element). */}