ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: keep model downloads running across navigation and loads (#6573)

Browse source 
* Studio: keep model downloads running across navigation and loads

Downloads started from the chat model selector were tied to the staged
pick lifecycle, so they were cancelled in cases where Hub downloads keep
going. This makes the chat download flow behave like the Hub.

- Leaving the chat route or switching thread/project/new chat now detaches
  the staging UI but keeps the in-flight transfer running in the global
  download manager (new keepDownload option on abandonStagedModel).
- Staging a second pick no longer cancels the previous pick's download, so
  multiple models/variants can download at once.
- Picking a model to download while another model is loading now starts the
  download in the background instead of refusing, since a download is
  independent of a load.

* Studio: also background-download remote GGUF quants while a model loads

isDownloadableHubRepo (wantManagerDownload) excludes GGUF sources, so an
uncached remote GGUF quant picked from the chat selector while another model
was loading fell through to the 'Another model is already loading' toast
instead of downloading in the background. Treat an uncached remote hub GGUF as
a background download too, matching the staged-pick download path.

Addresses review feedback from gemini-code-assist and codex on PR #6573.

* Studio: only toast a background download once it actually starts

The chat background-download path (used when a model is already loading)
fired the "Downloading in the background" toast unconditionally, but
requestStart can return without starting a job: a cross-transport partial
records a conflict that is only resolvable from the Hub download card, and
a busy sibling variant returns after its own toast. So the user could be
told a download started when none did, with no way to resolve the conflict
from chat.

requestStart now reports an outcome (started/conflict/busy/error). The
chat path only shows the success toast on an actual start and points the
user to the Hub when a transport conflict needs resolving. The Hub card
surface keeps its existing behavior (it renders the conflict resolver, so
it ignores the outcome).

* Studio: report background-download outcome from real job state

The chat background-download toast trusted requestStart's optimistic
"started", but a start can no-op without throwing: startJob finalizes the
job as "error" when the backend refuses or fails apiStart, its peer guard
skips a fresh start, and hasActiveOrPendingStart trips on a snapshot, peer
variant, or pending preflight that is not this request. So the user could
be told a download started when none did.

Derive the outcome from the actual job state of the exact key
(running/cancelling = started, otherwise error/busy), so the toast only
fires for a transfer that is really live.

Also guard against re-downloading the model that is already loading: the
/load flow downloads before it sets the checkpoint, and that fetch is not
a download-manager job, so picking the same id+variant again would start a
second transfer against the same cache. Detect that pick and surface a
"this model is already loading" toast instead.

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Michael Han 2026-06-22 08:51:14 -07:00 • committed by GitHub
commit 0689bd3842
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 138 additions and 54 deletions
Download patch file Download diff file Expand all files Collapse all files

9
studio/frontend/src/app/routes/__root.tsx
View file

@ -182,7 +182,9 @@ function RootLayout() {
chatRuntime.setActiveThreadId(null);
chatRuntime.setActiveProjectId(null);
chatRuntime.setIncognito(false);
if (chatRuntime.pendingSelection) chatRuntime.abandonStagedModel();
// Detach the staging UI but keep any in-flight download running, like Hub.
if (chatRuntime.pendingSelection)
chatRuntime.abandonStagedModel({ keepDownload: true });
void navigate({
to: "/chat",
search: { new: crypto.randomUUID() },
@ -205,7 +207,10 @@ function RootLayout() {
chatRuntime.setActiveProjectId(null);
chatRuntime.setActiveThreadId(null);
chatRuntime.setIncognito(false);
if (chatRuntime.pendingSelection) chatRuntime.abandonStagedModel();
// Leaving chat must not kill an in-flight download: detach the staging UI
// but keep the transfer running in the manager, like a Hub download.
if (chatRuntime.pendingSelection)
chatRuntime.abandonStagedModel({ keepDownload: true });
}, [isChatRoute]);
return (

86
studio/frontend/src/features/chat/chat-page.tsx
View file

@ -18,6 +18,10 @@ import {
import { useSidebar } from "@/components/ui/sidebar";
import { Tooltip, TooltipContent } from "@/components/ui/tooltip";
import { useLatestRef } from "@/features/hub/hooks/use-latest-ref";
import {
DOWNLOAD_KIND,
downloadManager,
} from "@/features/hub/download-manager";
import {
type NativeIntent,
NativeModelChip,
@ -1093,6 +1097,11 @@ export function ChatPage({
const abandonStaged = useCallback(() => {
useChatRuntimeStore.getState().abandonStagedModel();
}, []);
// Detach a staged pick on navigation without cancelling its download: the
// transfer keeps running in the manager and lands in cache, like Hub.
const detachStaged = useCallback(() => {
useChatRuntimeStore.getState().abandonStagedModel({ keepDownload: true });
}, []);
// Tracks whether the chat page is still mounted, so a staged-load failure that
// resolves after the user left chat doesn't resurrect the abandoned pick.
const mountedRef = useRef(true);
@ -1620,8 +1629,8 @@ export function ChatPage({
const prev = prevChatContextRef.current;
prevChatContextRef.current = chatContextKey;
if (prev === null || prev === chatContextKey) return;
abandonStaged();
}, [chatContextKey, abandonStaged]);
detachStaged();
}, [chatContextKey, detachStaged]);
const hasActiveModel = Boolean(inferenceParams.checkpoint);
// Load immediately, or — when "Load on selection" is off — stage the pick so
@ -1639,25 +1648,70 @@ export function ChatPage({
(!hasGgufSource(selection) && !wantManagerDownload) ||
(store.loadOnSelection && selection.isDownloaded)
) {
// Abandon any staged pick first so its edited knobs (e.g. a custom
// context length) don't leak into this immediate load -- resolveLoad
// reads customContextLength before checking the target is GGUF.
abandonStaged();
// Detach any staged pick first so its edited knobs don't leak into this
// immediate load. Detach (not abandon) keeps its download running.
detachStaged();
await selectModel(selection);
return;
}
// Refuse staging while a load is in flight (it would be silently dropped);
// the immediate-load branch above is already guarded in selectModel.
// Loads can't queue behind each other, but a download is independent: if
// the pick needs downloading, start it in the manager so it runs alongside
// the load. Nothing to download (already on device) just waits.
if (store.modelLoading) {
toast.info("Another model is already loading", {
description: "Wait for it to finish or cancel it first.",
});
// Both an uncached non-GGUF snapshot (wantManagerDownload) and an
// uncached remote GGUF quant download through the manager, so either can
// run in the background while another model loads. wantManagerDownload
// excludes GGUF by design, so the GGUF case is checked separately.
const wantBackgroundDownload =
wantManagerDownload ||
(selection.source === "hub" &&
hasGgufSource(selection) &&
!selection.isDownloaded);
// The model currently loading already downloads as part of its own load
// (the /load flow fetches before setting the checkpoint), so re-picking
// it must not kick off a second transfer against the same cache.
const isLoadingThisPick =
!!loadingModel &&
normalizeModelRef(loadingModel.id) ===
normalizeModelRef(selection.id) &&
(loadingModel.ggufVariant ?? null) === (selection.ggufVariant ?? null);
if (isLoadingThisPick) {
toast.info("This model is already loading", {
description: "It's downloading as part of the load in progress.",
});
} else if (wantBackgroundDownload) {
// Only claim the download started once a job is actually created. A
// transport conflict records state that is only resolvable from the
// Hub download card, so point the user there instead of showing a
// success toast for a transfer that never began; "busy" and "error"
// already surface their own toasts.
const outcome = await downloadManager.requestStart({
kind: DOWNLOAD_KIND.MODEL,
repoId: selection.id,
variant: selection.ggufVariant ?? null,
expectedBytes: selection.expectedBytes ?? 0,
});
if (outcome === "started") {
toast.info("Downloading in the background", {
description:
"It'll be ready to load once the current model finishes.",
});
} else if (outcome === "conflict") {
toast.info("Resume this download from the Hub", {
description:
"An earlier partial download used a different transport. Open the Hub tab to resume or restart it.",
});
}
} else {
toast.info("Another model is already loading", {
description: "Wait for it to finish or cancel it first.",
});
}
return;
}
// Tear down any existing staged pick first so its in-flight download is
// cancelled, not left running after we rebind to the new pick. With the
// toggle on, autoLoad downloads silently then loads; off stages for the sheet.
abandonStaged();
// Detach the prior staged pick (keeping its download) before rebinding, so
// a second pick downloads alongside the first instead of cancelling it.
detachStaged();
store.stageModel({
id: selection.id,
isLora: selection.isLora,
@ -1670,7 +1724,7 @@ export function ChatPage({
autoLoad: store.loadOnSelection,
});
},
[abandonStaged, selectModel],
[detachStaged, selectModel, loadingModel],
);
const loadNativeModelIntent = useCallback(
async (intent: NativeIntent, loadingDescription: string) => {

28
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -778,9 +778,10 @@ type ChatRuntimeStore = {
/** Stage a pick for a deferred load: revert knobs to the loaded baseline,
* record the selection, and open the settings sheet. */
stageModel: (selection: PendingModelSelection) => void;
/** Abandon a staged pick without loading: revert the knobs to the loaded
* baseline and clear the pending selection. */
abandonStagedModel: () => void;
/** Abandon a staged pick without loading: revert knobs to the loaded baseline
* and clear the pending selection. Cancels its in-flight download too, unless
* `keepDownload` is set (navigation keeps the transfer running, like Hub). */
abandonStagedModel: (opts?: { keepDownload?: boolean }) => void;
setCustomContextLength: (v: number | null) => void;
setChatTemplateOverride: (template: string | null) => void;
setPendingAudio: (base64: string, name: string) => void;
@ -1544,15 +1545,9 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
// Refuse staging mid-load: post-load cleanup would silently drop the queued
// pick. stageOrLoad toasts first for callers that can.
if (get().modelLoading) return;
// Rebinding to a new pick keeps the prior pick's download running so the
// user can queue multiple downloads at once (Hub-style).
set((s) => {
if (
s.pendingSelection &&
(s.pendingSelection.id !== selection.id ||
(s.pendingSelection.ggufVariant ?? null) !==
(selection.ggufVariant ?? null))
) {
cancelStagedModelDownload(s.pendingSelection);
}
return {
...loadedBaselineSettings(s),
pendingSelection: selection,
@ -1566,14 +1561,13 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
};
});
},
abandonStagedModel: () => {
abandonStagedModel: (opts) => {
const { pendingSelection } = get();
if (!pendingSelection) return;
// Cancel the staged pick's in-flight download so it doesn't keep running
// after the staging UI is gone. Centralized here so every abandon path
// (sheet close, thread switch, route exit, new chat) cancels it, including
// root-level callers that have no access to the useRepoDownload hook.
cancelStagedModelDownload(pendingSelection);
// Cancel the staged pick's in-flight download (centralized for every abandon
// path: sheet close, thread switch, route exit, new chat). `keepDownload`
// opts out so navigation leaves the transfer running, like a Hub download.
if (!opts?.keepDownload) cancelStagedModelDownload(pendingSelection);
set((s) => ({ ...loadedBaselineSettings(s), pendingSelection: null }));
},
setCustomContextLength: (customContextLength) => set({ customContextLength }),

63
studio/frontend/src/features/hub/download-manager/transport-conflict.ts
View file

@ -80,24 +80,50 @@ async function activeSiblingTransport(
return null;
}
// Outcome of a start request so callers can tell whether a transfer for this
// exact request is actually live before telling the user it began. "started"
// means a running/cancelling job exists for this key (a fresh start or an
// already-active one). "conflict" means a transport partial conflict was
// recorded and must be resolved from the Hub download card; "busy" means the
// repo is occupied by a sibling variant/snapshot/pending start that is not this
// transfer; "error" means the start failed or was refused.
export type DownloadStartOutcome = "started" | "conflict" | "busy" | "error";
// A start can no-op without throwing: the backend can refuse it (startJob
// finalizes "error"), startJob's peer guard can skip it, or
// hasActiveOrPendingStart can trip on a snapshot/peer/pending that is not this
// request. Derive the outcome from the actual job state of this exact key so
// 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));
}
async function runWithPendingStartGuard(
req: DownloadRequest,
action: () => Promise<void>,
): Promise<void> {
action: () => Promise<DownloadStartOutcome>,
): Promise<DownloadStartOutcome> {
const startKey = pendingStartKey(req);
if (hasActiveOrPendingStart(req)) return;
// Already active or pending for the repo: only report "started" when this
// exact request is the live transfer; a peer/snapshot/pending start has not.
if (hasActiveOrPendingStart(req)) {
return isJobActiveFor(req) ? "started" : "busy";
}
runtimeRegistry.pendingStartRepoKeys.add(startKey);
try {
await action();
return await action();
} catch (error) {
reportConflictStartError(error);
return "error";
} finally {
runtimeRegistry.pendingStartRepoKeys.delete(startKey);
}
}
export async function requestStart(req: DownloadRequest): Promise<void> {
await runWithPendingStartGuard(req, async () => {
export async function requestStart(
req: DownloadRequest,
): Promise<DownloadStartOutcome> {
return runWithPendingStartGuard(req, async () => {
let mode: TransportMode = getTransportMode();
try {
mode = await effectiveTransportMode(mode);
@ -119,7 +145,7 @@ export async function requestStart(req: DownloadRequest): Promise<void> {
? "This repository is currently downloading with Xet. Switch to Xet or wait for it to finish."
: "This repository is currently downloading with HTTP. Switch to HTTP or wait for it to finish.",
});
return;
return "busy";
}
} catch (err) {
console.warn("Active download transport check failed.", err);
@ -139,7 +165,7 @@ export async function requestStart(req: DownloadRequest): Promise<void> {
},
pending: req,
});
return;
return "conflict";
}
if (status.has_partial && !status.last_transport) {
toast.info("Restarting this download", {
@ -163,7 +189,7 @@ export async function requestStart(req: DownloadRequest): Promise<void> {
"Starting with HTTP so an existing partial is not discarded. Switch transport to retry with Xet.",
});
await startJob(req, { useXet: false });
return;
return isJobActiveFor(req) ? "started" : "error";
}
toast.warning("Couldn't verify existing partial download", {
description:
@ -171,6 +197,7 @@ export async function requestStart(req: DownloadRequest): Promise<void> {
});
}
await startJob(req, { useXet: mode === TRANSPORT.XET });
return isJobActiveFor(req) ? "started" : "error";
});
}
@ -178,22 +205,24 @@ export function resumeConflict(conflictKey: string): void {
const entry = getState().conflicts[conflictKey];
if (!entry) return;
setConflict(conflictKey, null);
void runWithPendingStartGuard(entry.pending, () =>
startJob(entry.pending, {
void runWithPendingStartGuard(entry.pending, async () => {
await startJob(entry.pending, {
useXet: entry.info.previous === TRANSPORT.XET,
}),
);
});
return "started";
});
}
export function restartConflict(conflictKey: string): void {
const entry = getState().conflicts[conflictKey];
if (!entry) return;
setConflict(conflictKey, null);
void runWithPendingStartGuard(entry.pending, () =>
startJob(entry.pending, {
void runWithPendingStartGuard(entry.pending, async () => {
await startJob(entry.pending, {
useXet: entry.info.next === TRANSPORT.XET,
}),
);
});
return "started";
});
}
export function cancelConflict(conflictKey: string): void {

6
studio/frontend/src/features/hub/download-manager/use-repo-download.ts
View file

@ -120,8 +120,10 @@ export function useRepoDownload(config: RepoDownloadConfig): DownloadJob {
);
const requestStartDownload = useCallback(
(variant: string | null, expectedBytes: number) => {
return downloadManager.requestStart({
async (variant: string | null, expectedBytes: number) => {
// This surface renders the conflict resolver (transportConflict), so the
// start outcome is handled by the card UI; the awaited result is ignored.
await downloadManager.requestStart({
kind,
repoId,
variant,