Studio autoload: hold final attempt for pending scans, stop workers on exit, picker platform gate
The last load attempt is not spent while folder scans are still pending, since a smaller candidate can still enter the pool; earlier attempts keep flowing incrementally. Resolution workers stop claiming scans once autoload reaches a terminal result, so a successful early load no longer leaves background scans contending with inference. Local rows now apply the picker's chat-only format gate (GGUF on any host, MLX on Mac) and the cascade skips cached non-GGUF repos on chat-only installs, matching what the picker offers; the remembered path stays ungated since a recorded load proves the model runs on this install.
This commit is contained in:
Unsloth
parent
43a1edc2ea
commit
1d582a8781
2 changed files with 250 additions and 125 deletions
|
|
@ -15,6 +15,8 @@ import {
|
|||
isHiddenModelId,
|
||||
} from "@/features/hub/lib/hidden-models";
|
||||
import { resolveInitialConfig } from "@/features/model-picker";
|
||||
import { isMlxId } from "@/features/model-picker/components/model-selector/recommended-fit";
|
||||
import { fetchDeviceType, usePlatformStore } from "@/config/env";
|
||||
import { projectHasSources } from "@/features/rag/api/rag-api";
|
||||
import { apiUrl } from "@/lib/api-base";
|
||||
import { parseParamCountB } from "@/lib/model-size";
|
||||
|
|
@ -1479,15 +1481,51 @@ const AUTO_LOAD_LOCAL_SOURCES: ReadonlySet<string> = new Set([
|
|||
"custom",
|
||||
]);
|
||||
|
||||
/** The picker's chat-only platform snapshot, read once per auto-load run. */
|
||||
type AutoLoadPlatform = {
|
||||
chatOnly: boolean;
|
||||
isMac: boolean;
|
||||
};
|
||||
|
||||
// Mirrors the picker's localModelIsGguf / localModelIsMlx checks: the backend
|
||||
// format hint is authoritative for indexed rows, with the same name/path
|
||||
// fallbacks the picker applies.
|
||||
function localRowIsGgufLike(row: LocalModelInfo): boolean {
|
||||
return (
|
||||
row.model_format === "gguf" || row.path.toLowerCase().endsWith(".gguf")
|
||||
);
|
||||
}
|
||||
|
||||
function localRowIsMlxNamed(row: LocalModelInfo): boolean {
|
||||
return (
|
||||
isMlxId(row.id) ||
|
||||
isMlxId(row.display_name ?? "") ||
|
||||
isMlxId(row.model_id ?? "")
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Backend-indexed local rows eligible for background auto-load: same policy
|
||||
* as the on-device picker (complete, chat-capable, not hidden infra), plus
|
||||
* no variant requirement, since a background load cannot ask for a quant.
|
||||
*/
|
||||
function isAutoLoadableLocalRow(row: LocalModelInfo): boolean {
|
||||
function isAutoLoadableLocalRow(
|
||||
row: LocalModelInfo,
|
||||
platform: AutoLoadPlatform,
|
||||
): boolean {
|
||||
if (!AUTO_LOAD_LOCAL_SOURCES.has(row.source)) return false;
|
||||
if (row.capabilities?.can_chat !== true) return false;
|
||||
if (row.partial) return false;
|
||||
// Chat-only installs run GGUF (any host) and MLX (Mac only); the picker
|
||||
// hides other local formats there, so the background load must not pick a
|
||||
// row the user could not have selected (mirrors sortedLocalDir's gate).
|
||||
if (
|
||||
platform.chatOnly &&
|
||||
!localRowIsGgufLike(row) &&
|
||||
!(platform.isMac && localRowIsMlxNamed(row))
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
// Adapters are chat-capable but load by resolving their base model, which
|
||||
// for a Hub-id base can trigger the implicit remote fetch a background
|
||||
// auto-load must never start. Adapters stay interactive-only.
|
||||
|
|
@ -2007,9 +2045,20 @@ export async function autoLoadOnDeviceModel(): Promise<{
|
|||
};
|
||||
}
|
||||
|
||||
// Resolve the platform snapshot the picker's format gates key on. Cached
|
||||
// after the app's boot fetch and never throws (falls back to client-side
|
||||
// detection when the backend is not ready).
|
||||
await fetchDeviceType();
|
||||
const platformState = usePlatformStore.getState();
|
||||
const platform: AutoLoadPlatform = {
|
||||
chatOnly: platformState.isChatOnly(),
|
||||
isMac: platformState.deviceType === "mac",
|
||||
};
|
||||
const ggufRepos = allGgufRepos.filter(isAutoLoadableCachedRepo);
|
||||
const modelRepos = allModelRepos.filter(isAutoLoadableCachedRepo);
|
||||
const localRows = allLocalRows.filter(isAutoLoadableLocalRow);
|
||||
const localRows = allLocalRows.filter((row) =>
|
||||
isAutoLoadableLocalRow(row, platform),
|
||||
);
|
||||
// Dedupe candidates that resolve to the SAME load target (e.g. a custom
|
||||
// scan folder pointing into an HF cache). Keyed on kind + load target /
|
||||
// on-disk path: a shared model_id does not mean the same files (a distinct
|
||||
|
|
@ -2218,8 +2267,12 @@ export async function autoLoadOnDeviceModel(): Promise<{
|
|||
}
|
||||
readyPool.splice(at, 0, entry);
|
||||
};
|
||||
// Non-GGUF cached repos need no scan: their snapshot loads whole.
|
||||
for (const repo of modelRepos) {
|
||||
// Non-GGUF cached repos need no scan: their snapshot loads whole. The
|
||||
// picker hides cached non-GGUF rows entirely on chat-only installs, so
|
||||
// the automatic cascade must not pick one there either; the remembered
|
||||
// path above stays ungated since a recorded load is user precedent that
|
||||
// the model runs on this install (e.g. an MLX repo loaded on a Mac).
|
||||
for (const repo of platform.chatOnly ? [] : modelRepos) {
|
||||
insertReady({
|
||||
type: "cached-model",
|
||||
repo,
|
||||
|
|
@ -2335,6 +2388,11 @@ export async function autoLoadOnDeviceModel(): Promise<{
|
|||
}
|
||||
}
|
||||
let pendingJobs = resolutionJobs.length;
|
||||
// Once auto-load reaches a terminal result (a model loaded, the attempt
|
||||
// cap was hit, or the pool drained), the workers stop claiming jobs so a
|
||||
// successful early load does not leave background scans hammering the
|
||||
// backend while inference is already running. In-flight scans finish.
|
||||
let resolutionStopped = false;
|
||||
let progressWaiters: Array<() => void> = [];
|
||||
const signalProgress = (): void => {
|
||||
const waiters = progressWaiters;
|
||||
|
|
@ -2356,7 +2414,7 @@ export async function autoLoadOnDeviceModel(): Promise<{
|
|||
),
|
||||
},
|
||||
async () => {
|
||||
while (nextJob < resolutionJobs.length) {
|
||||
while (!resolutionStopped && nextJob < resolutionJobs.length) {
|
||||
const job = resolutionJobs[nextJob];
|
||||
nextJob += 1;
|
||||
let entry: FallbackCandidate | null = null;
|
||||
|
|
@ -2385,151 +2443,165 @@ export async function autoLoadOnDeviceModel(): Promise<{
|
|||
});
|
||||
});
|
||||
}
|
||||
while (loadAttempts < MAX_AUTO_LOAD_ATTEMPTS) {
|
||||
const candidate = readyPool[0];
|
||||
if (!candidate) {
|
||||
if (pendingJobs <= 0) {
|
||||
break;
|
||||
try {
|
||||
while (loadAttempts < MAX_AUTO_LOAD_ATTEMPTS) {
|
||||
const candidate = readyPool[0];
|
||||
if (!candidate) {
|
||||
if (pendingJobs <= 0) {
|
||||
break;
|
||||
}
|
||||
await nextProgress();
|
||||
continue;
|
||||
}
|
||||
await nextProgress();
|
||||
continue;
|
||||
}
|
||||
// Every pending scan job can still yield a GGUF candidate (no-scan
|
||||
// rows were seeded upfront), and GGUF outranks safetensors in the
|
||||
// documented order, so the model-kind group stays gated until the
|
||||
// scans settle; resolved GGUF entries keep flowing immediately.
|
||||
if (isModelKindEntry(candidate) && pendingJobs > 0) {
|
||||
await nextProgress();
|
||||
continue;
|
||||
}
|
||||
readyPool.shift();
|
||||
if (candidate.type === "cached-gguf") {
|
||||
const repo = candidate.repo;
|
||||
if (!candidate.retry) {
|
||||
// A shared load target may already have been visited through an
|
||||
// indexed local row (e.g. a scan folder aliasing this cache).
|
||||
if (isSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path)) {
|
||||
// The final attempt is precious: while scans are still pending, a
|
||||
// smaller candidate can still enter the pool, so the last slot is
|
||||
// not spent until resolution settles and the global order is
|
||||
// complete. Earlier attempts keep flowing incrementally.
|
||||
if (pendingJobs > 0 && loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS - 1) {
|
||||
await nextProgress();
|
||||
continue;
|
||||
}
|
||||
// Every pending scan job can still yield a GGUF candidate (no-scan
|
||||
// rows were seeded upfront), and GGUF outranks safetensors in the
|
||||
// documented order, so the model-kind group stays gated until the
|
||||
// scans settle; resolved GGUF entries keep flowing immediately.
|
||||
if (isModelKindEntry(candidate) && pendingJobs > 0) {
|
||||
await nextProgress();
|
||||
continue;
|
||||
}
|
||||
readyPool.shift();
|
||||
if (candidate.type === "cached-gguf") {
|
||||
const repo = candidate.repo;
|
||||
if (!candidate.retry) {
|
||||
// A shared load target may already have been visited through an
|
||||
// indexed local row (e.g. a scan folder aliasing this cache).
|
||||
if (isSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path)) {
|
||||
continue;
|
||||
}
|
||||
markSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path);
|
||||
}
|
||||
const skipKey = autoLoadCandidateKey(
|
||||
"gguf",
|
||||
repo.load_id || repo.repo_id,
|
||||
candidate.variant.quant,
|
||||
);
|
||||
if (skippedAutoLoadCandidates.has(skipKey)) {
|
||||
continue;
|
||||
}
|
||||
markSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path);
|
||||
try {
|
||||
if (
|
||||
await loadAutoLoadCandidate({
|
||||
id: repo.repo_id,
|
||||
loadId: repo.load_id,
|
||||
kind: "gguf",
|
||||
ggufVariant: candidate.variant.quant,
|
||||
maxSeqLength: 0,
|
||||
successLabel: `Loaded ${repo.repo_id} (${candidate.variant.quant})`,
|
||||
inventoryId: repo.inventory_id ?? null,
|
||||
source: "hf_cache",
|
||||
})
|
||||
) {
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
}
|
||||
} catch {
|
||||
hadNonTrustFailure = true;
|
||||
skippedAutoLoadCandidates.add(skipKey);
|
||||
// A quant that passed validation can still fail /load (corrupt
|
||||
// file, llama.cpp startup error). Re-enter the repo's next
|
||||
// complete quant into the global size order, so one repo of
|
||||
// failing quants cannot starve a smaller model elsewhere.
|
||||
// Validation blocks are model-scoped, so they get no requeue.
|
||||
try {
|
||||
const next = await resolveCachedGgufEntry(repo);
|
||||
if (next) {
|
||||
insertReady({ ...next, retry: true });
|
||||
}
|
||||
} catch {
|
||||
hadNonTrustFailure = true;
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
const skipKey = autoLoadCandidateKey(
|
||||
"gguf",
|
||||
repo.load_id || repo.repo_id,
|
||||
candidate.variant.quant,
|
||||
);
|
||||
if (skippedAutoLoadCandidates.has(skipKey)) {
|
||||
if (candidate.type === "cached-model") {
|
||||
const repo = candidate.repo;
|
||||
if (isSeen("model", repo.load_id || repo.repo_id, repo.cache_path)) {
|
||||
continue;
|
||||
}
|
||||
markSeen("model", repo.load_id || repo.repo_id, repo.cache_path);
|
||||
if (
|
||||
skippedAutoLoadCandidates.has(
|
||||
autoLoadCandidateKey("model", repo.load_id || repo.repo_id),
|
||||
)
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
try {
|
||||
if (
|
||||
await loadAutoLoadCandidate({
|
||||
id: repo.repo_id,
|
||||
loadId: repo.load_id,
|
||||
kind: "model",
|
||||
ggufVariant: null,
|
||||
maxSeqLength: 4096,
|
||||
successLabel: `Loaded ${repo.repo_id}`,
|
||||
inventoryId: repo.inventory_id ?? null,
|
||||
source: "hf_cache",
|
||||
})
|
||||
) {
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
}
|
||||
} catch {
|
||||
hadNonTrustFailure = true;
|
||||
skippedAutoLoadCandidates.add(
|
||||
autoLoadCandidateKey("model", repo.load_id || repo.repo_id),
|
||||
);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
const row = candidate.row;
|
||||
const localCandidate = candidate.candidate;
|
||||
if (!candidate.retry) {
|
||||
if (isSeen(localCandidate.kind, row.load_id, row.id, row.path)) {
|
||||
continue;
|
||||
}
|
||||
markSeen(localCandidate.kind, row.load_id, row.id, row.path);
|
||||
}
|
||||
if (isSkippedAutoLoadCandidate(localCandidate)) {
|
||||
continue;
|
||||
}
|
||||
try {
|
||||
if (
|
||||
await loadAutoLoadCandidate({
|
||||
id: repo.repo_id,
|
||||
loadId: repo.load_id,
|
||||
kind: "gguf",
|
||||
ggufVariant: candidate.variant.quant,
|
||||
maxSeqLength: 0,
|
||||
successLabel: `Loaded ${repo.repo_id} (${candidate.variant.quant})`,
|
||||
inventoryId: repo.inventory_id ?? null,
|
||||
source: "hf_cache",
|
||||
})
|
||||
) {
|
||||
if (await loadAutoLoadCandidate(localCandidate)) {
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
}
|
||||
} catch {
|
||||
hadNonTrustFailure = true;
|
||||
skippedAutoLoadCandidates.add(skipKey);
|
||||
// A quant that passed validation can still fail /load (corrupt
|
||||
// file, llama.cpp startup error). Re-enter the repo's next
|
||||
// complete quant into the global size order, so one repo of
|
||||
// failing quants cannot starve a smaller model elsewhere.
|
||||
// Validation blocks are model-scoped, so they get no requeue.
|
||||
skippedAutoLoadCandidates.add(autoLoadSkipKey(localCandidate));
|
||||
// Same requeue as the cached-gguf branch: the folder's next complete
|
||||
// quant re-enters the global size order instead of retrying inline.
|
||||
try {
|
||||
const next = await resolveCachedGgufEntry(repo);
|
||||
const next = await resolveLocalRowCandidate(
|
||||
row,
|
||||
null,
|
||||
isSkippedAutoLoadCandidate,
|
||||
);
|
||||
if (next) {
|
||||
insertReady({ ...next, retry: true });
|
||||
insertReady({
|
||||
type: "local",
|
||||
row,
|
||||
candidate: next.candidate,
|
||||
sizeBytes: next.sizeBytes,
|
||||
retry: true,
|
||||
});
|
||||
}
|
||||
} catch {
|
||||
hadNonTrustFailure = true;
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (candidate.type === "cached-model") {
|
||||
const repo = candidate.repo;
|
||||
if (isSeen("model", repo.load_id || repo.repo_id, repo.cache_path)) {
|
||||
continue;
|
||||
}
|
||||
markSeen("model", repo.load_id || repo.repo_id, repo.cache_path);
|
||||
if (
|
||||
skippedAutoLoadCandidates.has(
|
||||
autoLoadCandidateKey("model", repo.load_id || repo.repo_id),
|
||||
)
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
try {
|
||||
if (
|
||||
await loadAutoLoadCandidate({
|
||||
id: repo.repo_id,
|
||||
loadId: repo.load_id,
|
||||
kind: "model",
|
||||
ggufVariant: null,
|
||||
maxSeqLength: 4096,
|
||||
successLabel: `Loaded ${repo.repo_id}`,
|
||||
inventoryId: repo.inventory_id ?? null,
|
||||
source: "hf_cache",
|
||||
})
|
||||
) {
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
}
|
||||
} catch {
|
||||
hadNonTrustFailure = true;
|
||||
skippedAutoLoadCandidates.add(
|
||||
autoLoadCandidateKey("model", repo.load_id || repo.repo_id),
|
||||
);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
const row = candidate.row;
|
||||
const localCandidate = candidate.candidate;
|
||||
if (!candidate.retry) {
|
||||
if (isSeen(localCandidate.kind, row.load_id, row.id, row.path)) {
|
||||
continue;
|
||||
}
|
||||
markSeen(localCandidate.kind, row.load_id, row.id, row.path);
|
||||
}
|
||||
if (isSkippedAutoLoadCandidate(localCandidate)) {
|
||||
continue;
|
||||
}
|
||||
try {
|
||||
if (await loadAutoLoadCandidate(localCandidate)) {
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
}
|
||||
} catch {
|
||||
hadNonTrustFailure = true;
|
||||
skippedAutoLoadCandidates.add(autoLoadSkipKey(localCandidate));
|
||||
// Same requeue as the cached-gguf branch: the folder's next complete
|
||||
// quant re-enters the global size order instead of retrying inline.
|
||||
try {
|
||||
const next = await resolveLocalRowCandidate(
|
||||
row,
|
||||
null,
|
||||
isSkippedAutoLoadCandidate,
|
||||
);
|
||||
if (next) {
|
||||
insertReady({
|
||||
type: "local",
|
||||
row,
|
||||
candidate: next.candidate,
|
||||
sizeBytes: next.sizeBytes,
|
||||
retry: true,
|
||||
});
|
||||
}
|
||||
} catch {
|
||||
hadNonTrustFailure = true;
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
// Runs on every exit (successful return, cap, drained pool, or a
|
||||
// thrown error) so no worker keeps scanning after the outcome is set.
|
||||
resolutionStopped = true;
|
||||
}
|
||||
|
||||
// No auto-loadable on-device model (or the attempt cap was hit). Never
|
||||
|
|
|
|||
|
|
@ -846,7 +846,9 @@ def test_fallback_orders_by_resolved_quant_size():
|
|||
assert "sizeBytes: sizeOrUnknownBytes(variant.size_bytes)" in auto_load
|
||||
# The all-variant row sum only orders non-GGUF cached repos, whose
|
||||
# snapshot loads whole.
|
||||
seed_block = auto_load.split("for (const repo of modelRepos)", 1)[1]
|
||||
seed_block = auto_load.split(
|
||||
"for (const repo of platform.chatOnly ? [] : modelRepos)", 1
|
||||
)[1]
|
||||
seed_block = seed_block.split("const resolveCachedGgufEntry", 1)[0]
|
||||
assert "sizeOrUnknownBytes(repo.size_bytes)" in seed_block
|
||||
assert auto_load.count("sizeOrUnknownBytes(repo.size_bytes)") == 1
|
||||
|
|
@ -932,6 +934,57 @@ def test_pending_gguf_scans_gate_safetensors_candidates():
|
|||
assert "if (isModelKindEntry(candidate) && pendingJobs > 0) {" in auto_load
|
||||
|
||||
|
||||
def test_final_attempt_waits_for_pending_scans():
|
||||
"""Fast-resolving candidates whose loads fail must not exhaust the attempt
|
||||
cap while pending scans can still yield a smaller loadable quant: the
|
||||
final attempt is only spent once resolution has settled and the global
|
||||
order is complete."""
|
||||
src = _read("features/chat/api/chat-adapter.ts")
|
||||
auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1]
|
||||
assert (
|
||||
"if (pendingJobs > 0 && loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS - 1) {"
|
||||
in auto_load
|
||||
)
|
||||
|
||||
|
||||
def test_resolution_workers_stop_on_terminal_result():
|
||||
"""A successful early load (or any other terminal outcome) must stop the
|
||||
workers from claiming further folder scans, so autoload cannot leave
|
||||
background scans contending with inference for the backend and disk."""
|
||||
src = _read("features/chat/api/chat-adapter.ts")
|
||||
auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1]
|
||||
assert "let resolutionStopped = false;" in auto_load
|
||||
assert (
|
||||
"while (!resolutionStopped && nextJob < resolutionJobs.length)"
|
||||
in auto_load
|
||||
)
|
||||
# The flag is set in a finally so every exit path (return, break, throw)
|
||||
# stops the workers.
|
||||
assert "resolutionStopped = true;" in auto_load
|
||||
assert auto_load.index("} finally {") < auto_load.index(
|
||||
"resolutionStopped = true;"
|
||||
)
|
||||
|
||||
|
||||
def test_local_rows_apply_picker_platform_gate():
|
||||
"""Chat-only installs run GGUF (any host) and MLX (Mac only); the picker
|
||||
hides other local formats and all cached non-GGUF rows there, so the
|
||||
background cascade must not load a row the user could not have picked.
|
||||
The remembered path stays ungated: a recorded load is user precedent."""
|
||||
src = _read("features/chat/api/chat-adapter.ts")
|
||||
local_fn = src.split("function isAutoLoadableLocalRow", 1)[1]
|
||||
local_fn = local_fn.split("\nfunction ", 1)[0]
|
||||
assert "platform.chatOnly" in local_fn
|
||||
assert "localRowIsGgufLike(row)" in local_fn
|
||||
assert "platform.isMac && localRowIsMlxNamed(row)" in local_fn
|
||||
auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1]
|
||||
# The platform snapshot hydrates through the cached, non-throwing fetch.
|
||||
assert "await fetchDeviceType();" in auto_load
|
||||
# Cascade seeding of cached non-GGUF repos mirrors the picker's
|
||||
# chat-only exclusion; the remembered lookup above it stays unfiltered.
|
||||
assert "platform.chatOnly ? [] : modelRepos" in auto_load
|
||||
|
||||
|
||||
def test_autoload_keys_preserve_posix_path_case():
|
||||
"""Linux filesystems distinguish /models/Foo from /models/foo, so seen
|
||||
keys and remembered-model matching must not fold case on POSIX paths;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue