diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index e4acc83a42..b64f850f10 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -521,11 +521,6 @@ export function HubModelPicker({ const hasMoreRecommended = visibleRecommendedIds.length < recommendedIds.length; - // Fetch VRAM info for the full pool once (recommendedIds is stable across - // page increments) so we don't re-fetch on every scroll. - const { paramCountById: recommendedParamCountById } = - useRecommendedModelVram(recommendedIds); - const showHfSection = debouncedQuery.trim().length > 0; // Recommended models that match the current search query @@ -535,6 +530,19 @@ export function HubModelPicker({ return recommendedIds.filter((id) => normalizeForSearch(id).includes(q)); }, [showHfSection, debouncedQuery, recommendedIds]); + // Fetch VRAM info for visible models, plus any models surfaced by a search + // query so that filtered recommended models also show VRAM badges. + // Skip GGUF repos: they have no safetensors metadata and the render layer + // already shows a static "GGUF" badge instead of VRAM data. + const idsForVram = useMemo(() => { + const ids = showHfSection + ? [...new Set([...visibleRecommendedIds, ...filteredRecommendedIds])] + : visibleRecommendedIds; + return ids.filter((id) => !isGgufRepo(id)); + }, [visibleRecommendedIds, showHfSection, filteredRecommendedIds]); + const { paramCountById: recommendedParamCountById } = + useRecommendedModelVram(idsForVram); + const recommendedSet = useMemo( () => new Set(showHfSection ? filteredRecommendedIds : visibleRecommendedIds), [showHfSection, filteredRecommendedIds, visibleRecommendedIds], diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index fd43fdb90d..98634c40ff 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -105,6 +105,7 @@ export function ExportPage() { const [isLoadingLocalModels, setIsLoadingLocalModels] = useState(true); const [localModelsError, setLocalModelsError] = useState(null); const debouncedModelQuery = useDebouncedValue(modelInput); + const debouncedHfToken = useDebouncedValue(hfToken, 500); const [exportMethod, setExportMethod] = useState(null); const [quantLevels, setQuantLevels] = useState([]); @@ -205,7 +206,7 @@ export function ExportPage() { isLoading: isLoadingHfModels, error: hfSearchError, } = useHfModelSearch(debouncedModelQuery, { - accessToken: hfToken || undefined, + accessToken: debouncedHfToken || undefined, excludeGguf: true, }); const { error: tokenValidationError, isChecking: isCheckingToken } = diff --git a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx index 4e7454e20f..f2a4796c54 100644 --- a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx @@ -92,6 +92,7 @@ export function ModelSelectionStep() { const [inputValue, setInputValue] = useState(""); const selectingRef = useRef(false); const debouncedQuery = useDebouncedValue(inputValue); + const debouncedHfToken = useDebouncedValue(hfToken, 500); const task = modelType ? MODEL_TYPE_TO_HF_TASK[modelType] : undefined; const { results: hfResults, @@ -101,7 +102,7 @@ export function ModelSelectionStep() { error: hfSearchError, } = useHfModelSearch(debouncedQuery, { task, - accessToken: hfToken || undefined, + accessToken: debouncedHfToken || undefined, excludeGguf: true, priorityIds: PRIORITY_TRAINING_MODELS, }); diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx index 0d2b3e074d..63b7678842 100644 --- a/studio/frontend/src/features/studio/sections/model-section.tsx +++ b/studio/frontend/src/features/studio/sections/model-section.tsx @@ -119,6 +119,7 @@ export function ModelSection() { const [localModelsError, setLocalModelsError] = useState(null); const selectingRef = useRef(false); const debouncedQuery = useDebouncedValue(inputValue); + const debouncedHfToken = useDebouncedValue(hfToken, 500); function handleModelSelect(id: string | null) { selectingRef.current = true; @@ -167,7 +168,7 @@ export function ModelSection() { error: hfSearchError, } = useHfModelSearch(debouncedQuery, { task, - accessToken: hfToken || undefined, + accessToken: debouncedHfToken || undefined, excludeGguf: true, priorityIds: PRIORITY_TRAINING_MODELS, }); diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts index 69ea4d3b83..32b261956b 100644 --- a/studio/frontend/src/hooks/use-hf-model-search.ts +++ b/studio/frontend/src/hooks/use-hf-model-search.ts @@ -2,7 +2,8 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import type { PipelineType } from "@huggingface/hub"; -import { listModels, modelInfo } from "@huggingface/hub"; +import { listModels } from "@huggingface/hub"; +import { type CachedResult, cachedModelInfo, primeCacheFromListing } from "@/lib/hf-cache"; import { useCallback, useMemo } from "react"; import { useHfPaginatedSearch } from "./use-hf-paginated-search"; @@ -104,6 +105,24 @@ function makeMapModel(excludeGguf: boolean) { /** Number of unsloth results to pull up-front before yielding general results. */ const UNSLOTH_PREFETCH = 20; +/** + * Prime the hf-cache from a listModels result. For public (non-gated, + * non-private) models, also prime the anonymous slot so the VRAM hook + * gets cache hits without re-fetching. Gated/private models are only + * cached under the caller's token to avoid auth leakage. + */ +function primeFromListing( + name: string, + accessToken: string | undefined, + model: unknown, +): void { + const data = model as CachedResult; + primeCacheFromListing(name, accessToken, data); + if (accessToken && !data.private && !data.gated) { + primeCacheFromListing(name, undefined, data); + } +} + /** * Creates a merged async generator that yields unsloth-owned models first, * then general results (with deduplication). @@ -134,7 +153,10 @@ async function* mergedModelIterator( let count = 0; for await (const model of unslothIter) { const m = model as { name?: string }; - if (m.name) seen.add(m.name); + if (m.name) { + seen.add(m.name); + primeFromListing(m.name, accessToken, model); + } yield model; count++; if (count >= UNSLOTH_PREFETCH) break; @@ -144,6 +166,9 @@ async function* mergedModelIterator( for await (const model of generalIter) { const m = model as { name?: string }; if (m.name && seen.has(m.name)) continue; + if (m.name) { + primeFromListing(m.name, accessToken, model); + } yield model; } } @@ -167,7 +192,7 @@ async function* priorityThenListingIterator( const seen = new Set(); const settled = await Promise.allSettled( priorityIds.map((id) => - modelInfo({ + cachedModelInfo({ name: id, additionalFields: ["safetensors", "tags"], ...(accessToken ? { credentials: { accessToken } } : {}), @@ -192,6 +217,9 @@ async function* priorityThenListingIterator( for await (const model of generalIter) { const m = model as { name?: string }; if (m.name && seen.has(m.name)) continue; + if (m.name) { + primeFromListing(m.name, accessToken, model); + } yield model; } } diff --git a/studio/frontend/src/hooks/use-recommended-model-vram.ts b/studio/frontend/src/hooks/use-recommended-model-vram.ts index cf464365d5..8826b38da9 100644 --- a/studio/frontend/src/hooks/use-recommended-model-vram.ts +++ b/studio/frontend/src/hooks/use-recommended-model-vram.ts @@ -1,7 +1,7 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { modelInfo } from "@huggingface/hub"; +import { cachedModelInfo } from "@/lib/hf-cache"; import { useEffect, useState } from "react"; /** @@ -10,9 +10,9 @@ import { useEffect, useState } from "react"; * models in the chat model dropdown. */ export function useRecommendedModelVram(ids: string[]) { - const [paramCountById, setParamCountById] = useState< - Map - >(new Map()); + const [paramCountById, setParamCountById] = useState>( + new Map(), + ); const [isLoading, setIsLoading] = useState(false); const stableKey = [...ids].filter(Boolean).sort().join(","); @@ -30,14 +30,15 @@ export function useRecommendedModelVram(ids: string[]) { const next = new Map(); await Promise.all( stableIds.map(async (id) => { - if (canceled) return; + if (canceled) { + return; + } try { - const info = await modelInfo({ + const info = await cachedModelInfo({ name: id, additionalFields: ["safetensors"], }); - const raw = info as { safetensors?: { total?: number } }; - const total = raw.safetensors?.total; + const total = info.safetensors?.total; if (typeof total === "number" && total > 0) { next.set(id, total); } @@ -47,7 +48,9 @@ export function useRecommendedModelVram(ids: string[]) { }), ); if (!canceled) { - setParamCountById(next); + // Merge with previous state so that VRAM badges for already-visible + // models are preserved while newly-visible models are still loading. + setParamCountById((prev) => new Map([...prev, ...next])); setIsLoading(false); } })(); diff --git a/studio/frontend/src/lib/hf-cache.ts b/studio/frontend/src/lib/hf-cache.ts new file mode 100644 index 0000000000..d87345eb61 --- /dev/null +++ b/studio/frontend/src/lib/hf-cache.ts @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { type ModelEntry, modelInfo } from "@huggingface/hub"; + +/** + * Thin caching + throttling layer over `modelInfo()` from @huggingface/hub. + * + * - TTL cache: avoids re-fetching the same model within CACHE_TTL_MS + * - In-flight dedup: concurrent callers for the same key share one request + * - Concurrency limiter: at most MAX_CONCURRENT requests in parallel; + * the rest queue and fire as slots free up + */ + +const CACHE_TTL_MS = 5 * 60 * 1000; // 5 minutes +// HF API allows bursts but rate-limits sustained traffic; 3 parallel requests +// keeps startup snappy while staying well under the observed throttle threshold. +const MAX_CONCURRENT = 3; + +// ── Cache & in-flight maps ────────────────────────────────────── + +// Extend ModelEntry with the additional fields we always request so callers +// do not need unsafe casts to access safetensors/tags. +export type CachedResult = ModelEntry & { + safetensors?: { total?: number; parameters?: Record }; + tags?: string[]; +}; + +interface CacheEntry { + data: CachedResult; + ts: number; +} + +const cache = new Map(); +const inflight = new Map>(); + +// ── Concurrency semaphore ─────────────────────────────────────── + +let active = 0; +const waiting: Array<() => void> = []; + +function acquire(): Promise { + if (active < MAX_CONCURRENT) { + active++; + return Promise.resolve(); + } + return new Promise((resolve) => + waiting.push(() => { + active++; + resolve(); + }), + ); +} + +function release() { + active--; + const next = waiting.shift(); + if (next) { + next(); + } +} + +// ── Public API ────────────────────────────────────────────────── + +// Always request the superset of fields any consumer needs so a single +// cache entry covers all callers (e.g. ["safetensors"] and ["safetensors","tags"]). +const ALL_FIELDS: ("safetensors" | "tags")[] = ["safetensors", "tags"]; + +function isStale(key: string): boolean { + const hit = cache.get(key); + if (!hit) return true; + return Date.now() - hit.ts >= CACHE_TTL_MS; +} + +function cacheKey(name: string, token: string | undefined): string { + if (!token) { + return `${name}::anon`; + } + // Use last 8 chars as a lightweight fingerprint so different tokens get + // separate cache entries without storing the full secret in memory. + return `${name}::${token.slice(-8)}`; +} + +function extractToken( + params: Parameters[0], +): string | undefined { + // The @huggingface/hub CredentialsParams union supports two forms: + // { accessToken: "hf_..." } -- current preferred form + // { credentials: { accessToken: "..." }} -- deprecated form + // Check both so the cache key is correct regardless of which form callers use. + if (params.accessToken) { + return params.accessToken; + } + if (params.credentials && "accessToken" in params.credentials) { + return params.credentials.accessToken; + } + return undefined; +} + +/** + * Pre-populate the cache with data from a listModels result. + * Only writes if the key is not already fresh -- never overwrites a recent + * modelInfo response with a listing response. + */ +export function primeCacheFromListing( + name: string, + token: string | undefined, + data: CachedResult, +): void { + if (!name) return; + const key = cacheKey(name, token); + if (!isStale(key)) return; // already fresh, don't overwrite + cache.set(key, { data, ts: Date.now() }); +} + +export async function cachedModelInfo( + params: Parameters[0], +): Promise { + const token = extractToken(params); + const key = cacheKey(params.name, token); + + // 1. Return from cache if fresh + if (!isStale(key)) { + return cache.get(key)!.data; + } + + // 2. Share in-flight request if one exists + const flying = inflight.get(key); + if (flying) { + return flying; + } + + // 3. New request, gated by concurrency semaphore + const promise = (async () => { + await acquire(); + try { + const result = await modelInfo({ + ...params, + additionalFields: ALL_FIELDS, + }); + const entry = { data: result as CachedResult, ts: Date.now() }; + cache.set(key, entry); + // For public (non-gated, non-private) models, also prime the anonymous + // cache slot so the VRAM hook (which reads without credentials) gets a + // cache hit. We skip gated/private models to avoid leaking auth-scoped + // metadata into the anonymous slot. + const r = result as CachedResult & { gated?: false | "auto" | "manual"; private?: boolean }; + if (token && !r.private && !r.gated) { + const anonKey = cacheKey(params.name, undefined); + if (isStale(anonKey)) { + cache.set(anonKey, entry); + } + } + return result as CachedResult; + } finally { + release(); + inflight.delete(key); + } + })(); + + inflight.set(key, promise); + return promise; +}