mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
fix(studio): adopt server-loaded model before chat auto-load (#5900)
* fix(studio): adopt server-loaded model before chat auto-load When the user starts Studio via `studio run -m`, the web UI could still auto-load a different cached GGUF on the first message because the chat checkpoint was empty. Sync from /api/inference/status before falling back to autoLoadSmallestModel so CLI-loaded models are not replaced. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(studio): hydrate adopted CLI model and harden auto-load errors Extract shared inference-status hydration for refresh() and CLI adopt paths so the first chat turn gets reasoning/tools flags. Wrap auto-load (including adopt) in try/catch for image-edit cleanup, and drop the redundant adopt call in run(). Co-authored-by: Cursor <cursoragent@cursor.com> * Guard model adoption against status failures and mid-flight selection for PR #5900 * ci: trigger pre-commit.ci after main merge Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
James Dawdy
GitHub
parent
f22e890ab8
commit
515abca84e
3 changed files with 257 additions and 162 deletions
|
|
@ -19,6 +19,7 @@ import {
|
|||
toExternalBackendProviderType,
|
||||
} from "../external-providers";
|
||||
import { pickFriendlyContainerName } from "../lib/friendly-names";
|
||||
import { tryAdoptServerActiveModel } from "../lib/apply-inference-status-to-store";
|
||||
import {
|
||||
clampReasoningEffortToLevels,
|
||||
getExternalMaxOutputTokens,
|
||||
|
|
@ -1129,6 +1130,10 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
loaded: boolean;
|
||||
blockedByTrustRemoteCode: boolean;
|
||||
}> {
|
||||
if (await tryAdoptServerActiveModel()) {
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
}
|
||||
|
||||
const store = useChatRuntimeStore.getState();
|
||||
const hfToken = store.hfToken || null;
|
||||
const trustRemoteCode = store.params.trustRemoteCode ?? false;
|
||||
|
|
@ -1472,6 +1477,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
}
|
||||
|
||||
if (!useChatRuntimeStore.getState().params.checkpoint) {
|
||||
// Prefer a model already loaded by the CLI/API before auto-loading.
|
||||
let loaded: boolean;
|
||||
let blockedByTrustRemoteCode: boolean;
|
||||
try {
|
||||
|
|
|
|||
|
|
@ -24,12 +24,15 @@ import {
|
|||
} from "../api/chat-api";
|
||||
import { formatEta, formatRate } from "../utils/format-transfer";
|
||||
import {
|
||||
CHAT_REASONING_ENABLED_KEY,
|
||||
loadOptionalBool,
|
||||
type ReasoningEffort,
|
||||
resolveToolsEnabledOnLoad,
|
||||
useChatRuntimeStore,
|
||||
} from "../stores/chat-runtime-store";
|
||||
import {
|
||||
applyActiveModelStatusToStore,
|
||||
clampLocalReasoningEffort,
|
||||
normalizeSpeculativeType,
|
||||
resolveInferenceCheckpointId,
|
||||
} from "../lib/apply-inference-status-to-store";
|
||||
import {
|
||||
mergeBackendRecommendedInference,
|
||||
resolveLoadMaxSeqLength,
|
||||
|
|
@ -211,39 +214,6 @@ function getTrustRemoteCodeRequiredMessage(modelName: string): string {
|
|||
return `${modelName} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`;
|
||||
}
|
||||
|
||||
// Canonicalises any backend/persisted value onto the Speculative Decoding
|
||||
// dropdown's modes ("auto"/"mtp"/"ngram"/"mtp+ngram"/"off"/null). Mirrors
|
||||
// backend _canonicalize_spec_mode so legacy persisted values round-trip.
|
||||
function normalizeSpeculativeType(v: string | null | undefined): string | null {
|
||||
if (v == null) return null;
|
||||
const s = String(v).trim().toLowerCase();
|
||||
if (!s) return null;
|
||||
if (s === "auto" || s === "default") return "auto";
|
||||
if (s === "off") return "off";
|
||||
if (s === "ngram-simple") return "ngram-simple";
|
||||
if (s === "mtp" || s === "draft-mtp") return "mtp";
|
||||
if (s === "ngram" || s === "ngram-mod") return "ngram";
|
||||
if (s === "mtp+ngram") return "mtp+ngram";
|
||||
// Comma-chained legacy values (e.g. from older persisted state).
|
||||
const parts = s.split(",").map((p) => p.trim()).filter(Boolean);
|
||||
const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp");
|
||||
const hasNgram = parts.some((p) => p === "ngram" || p === "ngram-mod");
|
||||
if (hasMtp && hasNgram) return "mtp+ngram";
|
||||
if (hasMtp) return "mtp";
|
||||
if (hasNgram) return "ngram";
|
||||
// Unknown -> safe fallback to Auto so the dropdown stays controlled.
|
||||
return "auto";
|
||||
}
|
||||
|
||||
type LocalReasoningEffort = Extract<ReasoningEffort, "low" | "medium" | "high">;
|
||||
|
||||
function clampLocalReasoningEffort(value: ReasoningEffort): LocalReasoningEffort {
|
||||
if (value === "low" || value === "medium" || value === "high") {
|
||||
return value;
|
||||
}
|
||||
return "low";
|
||||
}
|
||||
|
||||
export function useChatModelRuntime() {
|
||||
const params = useChatRuntimeStore((state) => state.params);
|
||||
const models = useChatRuntimeStore((state) => state.models);
|
||||
|
|
@ -328,132 +298,15 @@ export function useChatModelRuntime() {
|
|||
const selectedCheckpoint = useChatRuntimeStore.getState().params.checkpoint;
|
||||
const isExternalSelectionActive = isExternalModelId(selectedCheckpoint);
|
||||
if (statusRes.active_model && !isExternalSelectionActive) {
|
||||
setCheckpoint(statusRes.active_model, statusRes.gguf_variant);
|
||||
|
||||
// Apply inference defaults on reconnect (page refresh with model already loaded)
|
||||
if (statusRes.inference) {
|
||||
const currentParams = useChatRuntimeStore.getState().params;
|
||||
setParams(
|
||||
mergeBackendRecommendedInference({
|
||||
current: currentParams,
|
||||
response: statusRes,
|
||||
modelId: statusRes.active_model,
|
||||
presetSource: useChatRuntimeStore.getState().activePresetSource,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
// Restore reasoning/tools support flags and context length
|
||||
const hydratingExistingModel =
|
||||
selectedCheckpoint !== statusRes.active_model ||
|
||||
useChatRuntimeStore.getState().activeGgufVariant !==
|
||||
(statusRes.gguf_variant ?? null);
|
||||
const supportsReasoning = statusRes.supports_reasoning ?? false;
|
||||
const reasoningAlwaysOn = statusRes.reasoning_always_on ?? false;
|
||||
const reasoningStyle = statusRes.reasoning_style ?? "enable_thinking";
|
||||
const reasoningEffortLevels =
|
||||
reasoningStyle === "reasoning_effort"
|
||||
? (["low", "medium", "high"] as const)
|
||||
: (["low", "medium", "high"] as const);
|
||||
const supportsPreserveThinking = statusRes.supports_preserve_thinking ?? false;
|
||||
const supportsTools = statusRes.supports_tools ?? false;
|
||||
const storedReasoningEnabled = loadOptionalBool(
|
||||
CHAT_REASONING_ENABLED_KEY,
|
||||
);
|
||||
const currentGgufContextLength = statusRes.is_gguf
|
||||
? (statusRes.context_length ?? null)
|
||||
: null;
|
||||
const ggufMaxContextLength = statusRes.is_gguf
|
||||
? (statusRes.max_context_length ?? null)
|
||||
: null;
|
||||
const ggufNativeContextLength = statusRes.is_gguf
|
||||
? (statusRes.native_context_length ?? null)
|
||||
: null;
|
||||
const currentSpecType = normalizeSpeculativeType(
|
||||
statusRes.speculative_type,
|
||||
);
|
||||
// Refresh runs on F5 (needs hydration) and right after a load (store
|
||||
// already set). For user-configurable params, only hydrate when the
|
||||
// shadow `loaded*` field is null ("not yet hydrated"); otherwise we'd
|
||||
// clobber what the load path just applied and revert the user.
|
||||
const prevState = useChatRuntimeStore.getState();
|
||||
const clampedReasoningEffort = clampLocalReasoningEffort(
|
||||
prevState.reasoningEffort,
|
||||
);
|
||||
const nextDefaultChatTemplate =
|
||||
statusRes.chat_template === undefined
|
||||
? prevState.defaultChatTemplate
|
||||
: statusRes.chat_template;
|
||||
useChatRuntimeStore.setState({
|
||||
supportsReasoning,
|
||||
reasoningAlwaysOn,
|
||||
reasoningStyle,
|
||||
supportsReasoningOff: reasoningStyle !== "reasoning_effort",
|
||||
reasoningEffortLevels,
|
||||
reasoningEffort: clampedReasoningEffort,
|
||||
supportsPreserveThinking,
|
||||
supportsTools,
|
||||
// Reset per-turn reasoning flag so:
|
||||
// 1. non-reasoning models don't inherit a stale off state, and
|
||||
// 2. local reasoning-effort models (Off hidden via
|
||||
// supportsReasoningOff=false) don't carry reasoningEnabled=false
|
||||
// from an external model where Off was selected -- the composer
|
||||
// would still show "Think: <level>" but the adapter would omit
|
||||
// the kwarg, so Harmony falls back to its default effort.
|
||||
reasoningEnabled: supportsReasoning
|
||||
? reasoningStyle === "reasoning_effort"
|
||||
? true
|
||||
: useChatRuntimeStore.getState().reasoningEnabled
|
||||
: true,
|
||||
ggufContextLength: currentGgufContextLength,
|
||||
ggufMaxContextLength,
|
||||
ggufNativeContextLength,
|
||||
modelRequiresTrustRemoteCode:
|
||||
statusRes.requires_trust_remote_code ?? false,
|
||||
defaultChatTemplate: nextDefaultChatTemplate,
|
||||
loadedIsMultimodal: isMultimodalResponse(statusRes),
|
||||
specFallbackReason: statusRes.spec_fallback_reason ?? null,
|
||||
...(prevState.loadedSpeculativeType === null && {
|
||||
speculativeType: currentSpecType,
|
||||
loadedSpeculativeType: currentSpecType,
|
||||
}),
|
||||
...(statusRes.spec_draft_n_max !== undefined &&
|
||||
prevState.loadedSpecDraftNMax === null &&
|
||||
prevState.specDraftNMax === null && {
|
||||
specDraftNMax: statusRes.spec_draft_n_max ?? null,
|
||||
loadedSpecDraftNMax: statusRes.spec_draft_n_max ?? null,
|
||||
}),
|
||||
...(statusRes.cache_type_kv !== undefined &&
|
||||
prevState.loadedKvCacheDtype === null && {
|
||||
kvCacheDtype: statusRes.cache_type_kv,
|
||||
loadedKvCacheDtype: statusRes.cache_type_kv,
|
||||
}),
|
||||
...(statusRes.chat_template_override !== undefined &&
|
||||
prevState.loadedChatTemplateOverride === null &&
|
||||
prevState.chatTemplateOverride === null && {
|
||||
chatTemplateOverride: statusRes.chat_template_override,
|
||||
loadedChatTemplateOverride: statusRes.chat_template_override,
|
||||
}),
|
||||
});
|
||||
// setModels(listRes...) above used catalog data, which omits audio
|
||||
// capability. Re-apply live status so attach gates survive a refresh.
|
||||
syncModelCapabilities(statusRes.active_model, statusRes);
|
||||
|
||||
// Set reasoning default for Qwen3.5/3.6 small models
|
||||
if (
|
||||
supportsReasoning &&
|
||||
hydratingExistingModel &&
|
||||
storedReasoningEnabled === null
|
||||
) {
|
||||
let reasoningDefault = true;
|
||||
const mid = statusRes.active_model.toLowerCase();
|
||||
if (mid.includes("qwen3.5") || mid.includes("qwen3.6")) {
|
||||
const sizeMatch = mid.match(/(\d+\.?\d*)\s*b/);
|
||||
if (sizeMatch && parseFloat(sizeMatch[1]) < 9) {
|
||||
reasoningDefault = false;
|
||||
}
|
||||
}
|
||||
useChatRuntimeStore.setState({ reasoningEnabled: reasoningDefault });
|
||||
const checkpointId = resolveInferenceCheckpointId(statusRes);
|
||||
if (checkpointId) {
|
||||
setCheckpoint(checkpointId, statusRes.gguf_variant);
|
||||
applyActiveModelStatusToStore(statusRes, {
|
||||
previousCheckpoint: selectedCheckpoint,
|
||||
});
|
||||
// setModels(listRes...) above used catalog data, which omits audio
|
||||
// capability. Re-apply live status so attach gates survive a refresh.
|
||||
syncModelCapabilities(checkpointId, statusRes);
|
||||
}
|
||||
} else if (!statusRes.active_model && !isExternalSelectionActive) {
|
||||
useChatRuntimeStore.setState({
|
||||
|
|
|
|||
|
|
@ -0,0 +1,236 @@
|
|||
// 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 { getInferenceStatus } from "../api/chat-api";
|
||||
import { mergeBackendRecommendedInference } from "../presets/preset-policy";
|
||||
import {
|
||||
CHAT_REASONING_ENABLED_KEY,
|
||||
loadOptionalBool,
|
||||
type ReasoningEffort,
|
||||
resolveToolsEnabledOnLoad,
|
||||
useChatRuntimeStore,
|
||||
} from "../stores/chat-runtime-store";
|
||||
import { isMultimodalResponse, type InferenceStatusResponse } from "../types/api";
|
||||
import type { ChatModelSummary } from "../types/runtime";
|
||||
|
||||
type LocalReasoningEffort = Extract<ReasoningEffort, "low" | "medium" | "high">;
|
||||
|
||||
// Canonicalises backend / persisted speculative mode values onto the UI modes.
|
||||
export function normalizeSpeculativeType(
|
||||
v: string | null | undefined,
|
||||
): string | null {
|
||||
if (v == null) return null;
|
||||
const s = String(v).trim().toLowerCase();
|
||||
if (!s) return null;
|
||||
if (s === "auto" || s === "default") return "auto";
|
||||
if (s === "off") return "off";
|
||||
if (s === "ngram-simple") return "ngram-simple";
|
||||
if (s === "mtp" || s === "draft-mtp") return "mtp";
|
||||
if (s === "ngram" || s === "ngram-mod") return "ngram";
|
||||
if (s === "mtp+ngram") return "mtp+ngram";
|
||||
const parts = s.split(",").map((p) => p.trim()).filter(Boolean);
|
||||
const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp");
|
||||
const hasNgram = parts.some((p) => p === "ngram" || p === "ngram-mod");
|
||||
if (hasMtp && hasNgram) return "mtp+ngram";
|
||||
if (hasMtp) return "mtp";
|
||||
if (hasNgram) return "ngram";
|
||||
return "auto";
|
||||
}
|
||||
|
||||
export function clampLocalReasoningEffort(
|
||||
value: ReasoningEffort,
|
||||
): LocalReasoningEffort {
|
||||
if (value === "low" || value === "medium" || value === "high") {
|
||||
return value;
|
||||
}
|
||||
return "low";
|
||||
}
|
||||
|
||||
export function resolveInferenceCheckpointId(
|
||||
status: InferenceStatusResponse,
|
||||
): string | null {
|
||||
if (!status.active_model) return null;
|
||||
return status.model_identifier ?? status.active_model;
|
||||
}
|
||||
|
||||
function ensureActiveModelInStoreList(
|
||||
status: InferenceStatusResponse,
|
||||
checkpointId: string,
|
||||
): void {
|
||||
const store = useChatRuntimeStore.getState();
|
||||
if (store.models.some((model) => model.id === checkpointId)) {
|
||||
return;
|
||||
}
|
||||
const summary: ChatModelSummary = {
|
||||
id: checkpointId,
|
||||
name: status.active_model ?? checkpointId,
|
||||
isVision: status.is_vision ?? false,
|
||||
isLora: false,
|
||||
isGguf: status.is_gguf ?? false,
|
||||
isAudio: status.is_audio ?? false,
|
||||
audioType: status.audio_type ?? null,
|
||||
hasAudioInput: status.has_audio_input ?? false,
|
||||
};
|
||||
store.setModels([...store.models, summary]);
|
||||
}
|
||||
|
||||
export type ApplyInferenceStatusOptions = {
|
||||
previousCheckpoint?: string;
|
||||
};
|
||||
|
||||
/** Mirror refresh() hydration so adopted CLI models get reasoning/tools flags. */
|
||||
export function applyActiveModelStatusToStore(
|
||||
status: InferenceStatusResponse,
|
||||
options: ApplyInferenceStatusOptions = {},
|
||||
): void {
|
||||
const checkpointId = resolveInferenceCheckpointId(status);
|
||||
if (!checkpointId) return;
|
||||
|
||||
const store = useChatRuntimeStore.getState();
|
||||
const previousCheckpoint =
|
||||
options.previousCheckpoint ?? store.params.checkpoint;
|
||||
|
||||
if (status.inference) {
|
||||
store.setParams(
|
||||
mergeBackendRecommendedInference({
|
||||
current: store.params,
|
||||
response: status,
|
||||
modelId: checkpointId,
|
||||
presetSource: store.activePresetSource,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
const hydratingExistingModel =
|
||||
previousCheckpoint !== checkpointId ||
|
||||
store.activeGgufVariant !== (status.gguf_variant ?? null);
|
||||
const supportsReasoning = status.supports_reasoning ?? false;
|
||||
const reasoningAlwaysOn = status.reasoning_always_on ?? false;
|
||||
const reasoningStyle = status.reasoning_style ?? "enable_thinking";
|
||||
const reasoningEffortLevels =
|
||||
reasoningStyle === "reasoning_effort"
|
||||
? (["low", "medium", "high"] as const)
|
||||
: (["low", "medium", "high"] as const);
|
||||
const supportsPreserveThinking = status.supports_preserve_thinking ?? false;
|
||||
const supportsTools = status.supports_tools ?? false;
|
||||
const storedReasoningEnabled = loadOptionalBool(CHAT_REASONING_ENABLED_KEY);
|
||||
const currentGgufContextLength = status.is_gguf
|
||||
? (status.context_length ?? null)
|
||||
: null;
|
||||
const ggufMaxContextLength = status.is_gguf
|
||||
? (status.max_context_length ?? null)
|
||||
: null;
|
||||
const ggufNativeContextLength = status.is_gguf
|
||||
? (status.native_context_length ?? null)
|
||||
: null;
|
||||
const currentSpecType = normalizeSpeculativeType(status.speculative_type);
|
||||
const prevState = useChatRuntimeStore.getState();
|
||||
const clampedReasoningEffort = clampLocalReasoningEffort(
|
||||
prevState.reasoningEffort,
|
||||
);
|
||||
const nextDefaultChatTemplate =
|
||||
status.chat_template === undefined
|
||||
? prevState.defaultChatTemplate
|
||||
: status.chat_template;
|
||||
|
||||
useChatRuntimeStore.setState({
|
||||
supportsReasoning,
|
||||
reasoningAlwaysOn,
|
||||
reasoningStyle,
|
||||
supportsReasoningOff: reasoningStyle !== "reasoning_effort",
|
||||
reasoningEffortLevels,
|
||||
reasoningEffort: clampedReasoningEffort,
|
||||
supportsPreserveThinking,
|
||||
supportsTools,
|
||||
...resolveToolsEnabledOnLoad(supportsTools),
|
||||
reasoningEnabled: supportsReasoning
|
||||
? reasoningStyle === "reasoning_effort"
|
||||
? true
|
||||
: useChatRuntimeStore.getState().reasoningEnabled
|
||||
: true,
|
||||
ggufContextLength: currentGgufContextLength,
|
||||
ggufMaxContextLength,
|
||||
ggufNativeContextLength,
|
||||
modelRequiresTrustRemoteCode: status.requires_trust_remote_code ?? false,
|
||||
defaultChatTemplate: nextDefaultChatTemplate,
|
||||
loadedIsMultimodal: isMultimodalResponse(status),
|
||||
specFallbackReason: status.spec_fallback_reason ?? null,
|
||||
...(prevState.loadedSpeculativeType === null && {
|
||||
speculativeType: currentSpecType,
|
||||
loadedSpeculativeType: currentSpecType,
|
||||
}),
|
||||
...(status.spec_draft_n_max !== undefined &&
|
||||
prevState.loadedSpecDraftNMax === null &&
|
||||
prevState.specDraftNMax === null && {
|
||||
specDraftNMax: status.spec_draft_n_max ?? null,
|
||||
loadedSpecDraftNMax: status.spec_draft_n_max ?? null,
|
||||
}),
|
||||
...(status.cache_type_kv !== undefined &&
|
||||
prevState.loadedKvCacheDtype === null && {
|
||||
kvCacheDtype: status.cache_type_kv,
|
||||
loadedKvCacheDtype: status.cache_type_kv,
|
||||
}),
|
||||
...(status.chat_template_override !== undefined &&
|
||||
prevState.loadedChatTemplateOverride === null &&
|
||||
prevState.chatTemplateOverride === null && {
|
||||
chatTemplateOverride: status.chat_template_override,
|
||||
loadedChatTemplateOverride: status.chat_template_override,
|
||||
}),
|
||||
});
|
||||
|
||||
ensureActiveModelInStoreList(status, checkpointId);
|
||||
|
||||
if (
|
||||
supportsReasoning &&
|
||||
hydratingExistingModel &&
|
||||
storedReasoningEnabled === null
|
||||
) {
|
||||
let reasoningDefault = true;
|
||||
const mid = checkpointId.toLowerCase();
|
||||
if (mid.includes("qwen3.5") || mid.includes("qwen3.6")) {
|
||||
const sizeMatch = mid.match(/(\d+\.?\d*)\s*b/);
|
||||
if (sizeMatch && parseFloat(sizeMatch[1]) < 9) {
|
||||
reasoningDefault = false;
|
||||
}
|
||||
}
|
||||
useChatRuntimeStore.setState({ reasoningEnabled: reasoningDefault });
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Adopt the model already loaded on the inference server (e.g. via
|
||||
* ``unsloth studio run -m``) into the chat UI checkpoint without
|
||||
* triggering a new /api/inference/load.
|
||||
*/
|
||||
export async function tryAdoptServerActiveModel(): Promise<boolean> {
|
||||
const store = useChatRuntimeStore.getState();
|
||||
if (store.params.checkpoint) {
|
||||
return true;
|
||||
}
|
||||
|
||||
let status: InferenceStatusResponse;
|
||||
try {
|
||||
status = await getInferenceStatus();
|
||||
} catch {
|
||||
// Status endpoint unavailable: fall back to the normal auto-load path.
|
||||
return false;
|
||||
}
|
||||
if (!status.active_model) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const checkpointId = resolveInferenceCheckpointId(status);
|
||||
if (!checkpointId) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Re-check after the await: keep a checkpoint the user picked meanwhile.
|
||||
const previousCheckpoint =
|
||||
useChatRuntimeStore.getState().params.checkpoint;
|
||||
if (previousCheckpoint) {
|
||||
return true;
|
||||
}
|
||||
store.setCheckpoint(checkpointId, status.gguf_variant);
|
||||
applyActiveModelStatusToStore(status, { previousCheckpoint });
|
||||
return true;
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue