diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index f87229d78f..fade5e9800 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -439,11 +439,69 @@ function waitForModelReady(abortSignal?: AbortSignal): Promise { */ // Cap cascade so broken cached repos can't spam /api/inference/load. const MAX_AUTO_LOAD_ATTEMPTS = 3; +// /validate is cheaper but a long trust-blocked cache still produces one POST +// per repo; cap that path separately so the cascade is fully bounded. +const MAX_AUTO_VALIDATE_ATTEMPTS = 10; -async function autoLoadSmallestModel(): Promise<{ +// Singleflight handle: concurrent sends arriving while a cascade is in +// flight share its result instead of fanning out another /load budget. +let autoLoadInflight: Promise<{ + loaded: boolean; + blockedByTrustRemoteCode: boolean; +}> | null = null; + +async function autoLoadSmallestModel(abortSignal?: AbortSignal): Promise<{ loaded: boolean; blockedByTrustRemoteCode: boolean; }> { + // Pre-aborted callers bail before even creating a toast. + if (abortSignal?.aborted) { + return { loaded: false, blockedByTrustRemoteCode: false }; + } + // Followers share the leader's cascade result but race against their own + // abortSignal so a cancelled stream isn't blocked on the leader. + if (autoLoadInflight) return awaitWithAbort(autoLoadInflight, abortSignal); + const work = runAutoLoadCascade(abortSignal); + autoLoadInflight = work; + try { + return await work; + } finally { + autoLoadInflight = null; + } +} + +function awaitWithAbort( + work: Promise<{ loaded: boolean; blockedByTrustRemoteCode: boolean }>, + signal?: AbortSignal, +): Promise<{ loaded: boolean; blockedByTrustRemoteCode: boolean }> { + if (!signal) return work; + if (signal.aborted) { + return Promise.resolve({ loaded: false, blockedByTrustRemoteCode: false }); + } + return new Promise((resolve, reject) => { + const onAbort = () => + resolve({ loaded: false, blockedByTrustRemoteCode: false }); + signal.addEventListener("abort", onAbort, { once: true }); + work.then( + (value) => { + signal.removeEventListener("abort", onAbort); + resolve(value); + }, + (err) => { + signal.removeEventListener("abort", onAbort); + reject(err); + }, + ); + }); +} + +async function runAutoLoadCascade(abortSignal?: AbortSignal): Promise<{ + loaded: boolean; + blockedByTrustRemoteCode: boolean; +}> { + if (abortSignal?.aborted) { + return { loaded: false, blockedByTrustRemoteCode: false }; + } const store = useChatRuntimeStore.getState(); const hfToken = store.hfToken || null; const trustRemoteCode = store.params.trustRemoteCode ?? false; @@ -455,6 +513,7 @@ async function autoLoadSmallestModel(): Promise<{ let blockedByTrustRemoteCode = false; let hadNonTrustFailure = false; let loadAttempts = 0; + let validateAttempts = 0; async function canAutoLoad(payload: { model_path: string; @@ -462,6 +521,8 @@ async function autoLoadSmallestModel(): Promise<{ is_lora: boolean; gguf_variant?: string | null; }): Promise { + if (validateAttempts >= MAX_AUTO_VALIDATE_ATTEMPTS) return false; + validateAttempts += 1; const validation = await validateModel({ ...payload, hf_token: hfToken, @@ -480,11 +541,17 @@ async function autoLoadSmallestModel(): Promise<{ listCachedModels().catch(() => []), ]); + if (abortSignal?.aborted) { + toast.dismiss(toastId); + return { loaded: false, blockedByTrustRemoteCode: false }; + } + // Try GGUF first: pick the repo with the smallest total size, // then pick its smallest downloaded variant. if (ggufRepos.length > 0) { const sorted = [...ggufRepos].sort((a, b) => a.size_bytes - b.size_bytes); for (const repo of sorted) { + if (abortSignal?.aborted) break; if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break; try { const variants = await listGgufVariants(repo.repo_id); @@ -566,6 +633,7 @@ async function autoLoadSmallestModel(): Promise<{ if (modelRepos.length > 0) { const sorted = [...modelRepos].sort((a, b) => a.size_bytes - b.size_bytes); for (const repo of sorted) { + if (abortSignal?.aborted) break; if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break; try { if ( @@ -624,6 +692,10 @@ async function autoLoadSmallestModel(): Promise<{ } } + if (abortSignal?.aborted) { + toast.dismiss(toastId); + return { loaded: false, blockedByTrustRemoteCode: false }; + } // Cap also gates the default download so the total /api/inference/load // budget across cached + fallback is MAX_AUTO_LOAD_ATTEMPTS, not +1. if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) { @@ -736,7 +808,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (!useChatRuntimeStore.getState().params.checkpoint) { // Auto-load the smallest downloaded model const { loaded, blockedByTrustRemoteCode } = - await autoLoadSmallestModel(); + await autoLoadSmallestModel(abortSignal); if (!loaded) { toast.error( blockedByTrustRemoteCode diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 67395a8378..9e67236bed 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -1315,7 +1315,7 @@ def _release_self_exe_lock_windows() -> None: os.replace(exe, stale) except OSError as e: # Not fatal; setup.ps1 retries from a sibling process. - print(f"[update] could not rename {exe.name} -> {stale.name}: {e}") + typer.echo(f"[update] could not rename {exe.name} -> {stale.name}: {e}") def _restore_self_exe_lock_windows() -> None: @@ -1341,7 +1341,7 @@ def _restore_self_exe_lock_windows() -> None: try: os.replace(stale, exe) except OSError as e: - print(f"[update] could not restore {stale.name} -> {exe.name}: {e}") + typer.echo(f"[update] could not restore {stale.name} -> {exe.name}: {e}") def _cleanup_self_exe_lock_windows() -> None: