fix: throttle and cache HuggingFace modelInfo API calls (#4696)
* fix: throttle and cache HuggingFace modelInfo API calls The frontend was firing 40 to 60 parallel modelInfo requests on app startup with zero caching or deduplication, causing HF rate limits. Adds a caching layer (hf-cache.ts) with TTL cache, inflight request dedup, and a concurrency limiter. Also debounces the HF token input so typing a token no longer re-fires all model searches per keystroke. * fix: only fetch VRAM info for visible models in chat selector * Fix cache key isolation and VRAM badge stability for PR #4696 - Cache key now includes a token fingerprint (last 8 chars) instead of a boolean, so switching HF tokens gives separate cache entries instead of serving stale data from the previous token. - Extract token via credentials?.accessToken to match the @huggingface/hub API surface. - Extend CachedResult type with safetensors/tags fields so downstream consumers no longer need unsafe `as` casts. - Merge VRAM param map with previous state on scroll instead of replacing it, preventing a brief flash of missing VRAM badges when new models become visible. * Fix VRAM badges missing for search-filtered recommended models When a user types a search query, filteredRecommendedIds can include models beyond the currently visible page. These models had no VRAM data because useRecommendedModelVram only received visibleRecommendedIds. Now we pass the union of visibleRecommendedIds and filteredRecommendedIds to the VRAM hook, so recommended models surfaced by search also show their VRAM badges. The hf-cache layer ensures no duplicate network calls. * Apply biome formatting to hf-cache.ts and use-recommended-model-vram.ts Auto-formatted with biome check --write to match project lint rules: - Block statements for single-line if/for bodies - Import sorting (type imports first) - Consistent line wrapping * Fix extractToken to handle both current and deprecated HF auth forms The @huggingface/hub CredentialsParams type is a union: - { accessToken: "hf_..." } (current preferred form) - { credentials: { accessToken: "..." } } (deprecated form) Previously only checked params.credentials?.accessToken (deprecated path). Now checks both forms so the cache key is correct regardless of which calling convention is used. * Simplify extractToken, map merge, and set construction - extractToken: remove type assertions, use direct property access with truthiness checks for cleaner union type handling - VRAM map merge: use Map spread constructor instead of manual for loop - idsForVram: use Set spread construction for more concise dedup * Add rationale comment for MAX_CONCURRENT=3 in hf-cache.ts * Skip GGUF repos in VRAM fetch and pre-populate cache from listModels Two changes to reduce redundant HF API calls: 1. Filter GGUF repos from idsForVram before passing to useRecommendedModelVram. GGUF repos have no safetensors metadata and the render layer already shows a static "GGUF" badge -- fetching modelInfo for them is a no-op that wastes a semaphore slot and a network round-trip. 2. Add primeCacheFromListing() to hf-cache.ts and call it from listModels yield sites in mergedModelIterator and priorityThenListingIterator. listModels returns the same type (ModelEntry & Pick<ApiModelInfo, T>) as modelInfo with the same additionalFields, so the data is interchangeable. Priming only writes if the key is not already fresh, so it never overwrites a recent modelInfo response. This means models discovered via listModels are already in cache when useRecommendedModelVram later calls cachedModelInfo for them, eliminating duplicate network requests. * Fix cache key mismatch: prime both token and anonymous slots The VRAM hook calls cachedModelInfo without credentials (anonymous key), but listModels results were primed only under the authenticated key. For authenticated users the priming was a no-op -- cache miss every time. Fix: prime both the token-specific slot and the anonymous slot when an access token is present. Public model metadata (safetensors, tags) is identical regardless of auth so this is safe. Also add a defensive guard in primeCacheFromListing for empty name. * Auto-prime anonymous cache slot from authenticated modelInfo fetches When cachedModelInfo is called with a token, the result was only stored under the token-specific key (e.g. model::abc12345). The VRAM hook calls cachedModelInfo without credentials and reads the anonymous slot (model::anon), causing a cache miss and duplicate fetch for every priority model. Now cachedModelInfo also writes to the anonymous slot on success when a token is present. Public model metadata (safetensors, tags) is identical regardless of auth, so this is safe and eliminates ~10 duplicate API calls on first page load. * Guard anonymous cache priming against gated/private models Only prime the anonymous cache slot for non-gated, non-private models. Previously, authenticated modelInfo responses and listing results were unconditionally copied into the anonymous slot, which could briefly expose gated/private model metadata after clearing the HF token. Now checks result.gated and result.private before writing the anon slot. Public unsloth/ models (the common case) still benefit from the optimization; gated models like meta-llama/* require a fresh fetch per auth context. * Extract primeFromListing helper to deduplicate cache priming logic The cache priming pattern (prime token slot + conditionally prime anon slot for non-gated models) was duplicated in three places. Extracted into a single primeFromListing() function for maintainability. * Export CachedResult type, add isStale helper, simplify primeFromListing - Export CachedResult so consumers can use it directly instead of the indirect Parameters<typeof ...> pattern. - Extract isStale(key) helper to deduplicate the cache freshness check that was repeated in primeCacheFromListing, cachedModelInfo, and the anonymous-slot priming logic. - Simplify primeFromListing to use CachedResult directly for both the data parameter and the gated/private guard, eliminating the double cast. --------- Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Wasim Yousef Said
GitHub
parent
3b5a49776b
commit
28aaf849bf
7 changed files with 225 additions and 20 deletions
|
|
@ -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],
|
||||
|
|
|
|||
|
|
@ -105,6 +105,7 @@ export function ExportPage() {
|
|||
const [isLoadingLocalModels, setIsLoadingLocalModels] = useState(true);
|
||||
const [localModelsError, setLocalModelsError] = useState<string | null>(null);
|
||||
const debouncedModelQuery = useDebouncedValue(modelInput);
|
||||
const debouncedHfToken = useDebouncedValue(hfToken, 500);
|
||||
|
||||
const [exportMethod, setExportMethod] = useState<ExportMethod | null>(null);
|
||||
const [quantLevels, setQuantLevels] = useState<string[]>([]);
|
||||
|
|
@ -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 } =
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
});
|
||||
|
|
|
|||
|
|
@ -119,6 +119,7 @@ export function ModelSection() {
|
|||
const [localModelsError, setLocalModelsError] = useState<string | null>(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,
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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<string>();
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<string, number>
|
||||
>(new Map());
|
||||
const [paramCountById, setParamCountById] = useState<Map<string, number>>(
|
||||
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<string, number>();
|
||||
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);
|
||||
}
|
||||
})();
|
||||
|
|
|
|||
163
studio/frontend/src/lib/hf-cache.ts
Normal file
163
studio/frontend/src/lib/hf-cache.ts
Normal file
|
|
@ -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<string, number> };
|
||||
tags?: string[];
|
||||
};
|
||||
|
||||
interface CacheEntry {
|
||||
data: CachedResult;
|
||||
ts: number;
|
||||
}
|
||||
|
||||
const cache = new Map<string, CacheEntry>();
|
||||
const inflight = new Map<string, Promise<CachedResult>>();
|
||||
|
||||
// ── Concurrency semaphore ───────────────────────────────────────
|
||||
|
||||
let active = 0;
|
||||
const waiting: Array<() => void> = [];
|
||||
|
||||
function acquire(): Promise<void> {
|
||||
if (active < MAX_CONCURRENT) {
|
||||
active++;
|
||||
return Promise.resolve();
|
||||
}
|
||||
return new Promise<void>((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<typeof modelInfo>[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<typeof modelInfo>[0],
|
||||
): Promise<CachedResult> {
|
||||
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;
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue