From 5f902af456dd479197279122ffdc3bd94dca7c23 Mon Sep 17 00:00:00 2001 From: Samit Date: Fri, 6 Mar 2026 22:01:27 -0800 Subject: [PATCH] fixed model unload before load --- studio/backend/models/inference.py | 28 +++++++ studio/backend/routes/inference.py | 46 +++++++++++ .../src/features/chat/api/chat-api.ts | 16 ++++ .../chat/hooks/use-chat-model-runtime.ts | 77 +++++++++++++++---- .../frontend/src/features/chat/types/api.ts | 10 +++ 5 files changed, 162 insertions(+), 15 deletions(-) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 3a908dcfc3..951167a25b 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -25,6 +25,34 @@ class UnloadRequest(BaseModel): model_path: str = Field(..., description="Model identifier to unload") +class ValidateModelRequest(BaseModel): + """ + Lightweight validation request to check whether a model identifier + *can be resolved* into a ModelConfig. + + This does NOT actually load weights into GPU memory. + """ + model_path: str = Field(..., description="Model identifier or local path") + hf_token: Optional[str] = Field(None, description="HuggingFace token for gated models") + gguf_variant: Optional[str] = Field(None, description="GGUF quantization variant (e.g. 'Q4_K_M')") + + +class ValidateModelResponse(BaseModel): + """ + Result of model validation. + + valid == True means ModelConfig.from_identifier() succeeded and basic + introspection (GGUF / LoRA / vision flags) is available. + """ + valid: bool = Field(..., description="Whether the model identifier looks valid") + message: str = Field(..., description="Human-readable validation message") + identifier: Optional[str] = Field(None, description="Resolved model identifier") + display_name: Optional[str] = Field(None, description="Display name derived from identifier") + is_gguf: bool = Field(False, description="Whether this is a GGUF model (llama.cpp)") + is_lora: bool = Field(False, description="Whether this is a LoRA adapter") + is_vision: bool = Field(False, description="Whether this is a vision-capable model") + + class GenerateRequest(BaseModel): """Request for text generation (legacy /generate/stream endpoint)""" messages: List[dict] = Field(..., description="Chat messages in OpenAI format") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 512eb05526..c4a0cdedd9 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -49,6 +49,8 @@ from models.inference import ( ChoiceDelta, CompletionChoice, CompletionMessage, + ValidateModelRequest, + ValidateModelResponse, ) from auth.authentication import get_current_subject @@ -198,6 +200,50 @@ async def load_model( ) +@router.post("/validate", response_model=ValidateModelResponse) +async def validate_model( + request: ValidateModelRequest, + current_subject: str = Depends(get_current_subject), +): + """ + Lightweight validation endpoint for model identifiers. + + This checks that ModelConfig.from_identifier() can resolve the given + model_path, but it does NOT actually load model weights into GPU memory. + """ + try: + config = ModelConfig.from_identifier( + model_id=request.model_path, + hf_token=request.hf_token, + gguf_variant=request.gguf_variant, + ) + + if not config: + raise HTTPException( + status_code=400, + detail=f"Invalid model identifier: {request.model_path}", + ) + + return ValidateModelResponse( + valid=True, + message="Model identifier is valid.", + identifier=config.identifier, + display_name=getattr(config, "display_name", config.identifier), + is_gguf=getattr(config, "is_gguf", False), + is_lora=getattr(config, "is_lora", False), + is_vision=getattr(config, "is_vision", False), + ) + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error validating model identifier '{request.model_path}': {e}", exc_info=True) + raise HTTPException( + status_code=400, + detail=f"Invalid model: {str(e)}", + ) + + @router.post("/unload", response_model=UnloadResponse) async def unload_model( request: UnloadRequest, diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 5d5a9551ef..d134aa203a 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -9,6 +9,7 @@ import type { OpenAIChatChunk, OpenAIChatCompletionsRequest, UnloadModelRequest, + ValidateModelResponse, } from "../types/api"; function parseErrorText(status: number, body: unknown): string { @@ -66,6 +67,21 @@ export async function loadModel( return parseJsonOrThrow(response); } +export async function validateModel( + payload: LoadModelRequest, +): Promise { + const response = await authFetch("/api/inference/validate", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model_path: payload.model_path, + hf_token: payload.hf_token, + gguf_variant: payload.gguf_variant ?? null, + }), + }); + return parseJsonOrThrow(response); +} + export async function unloadModel(payload: UnloadModelRequest): Promise { const response = await authFetch("/api/inference/unload", { method: "POST", diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index fece047cd8..bb942678b2 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -6,6 +6,7 @@ import { listModels, loadModel, unloadModel, + validateModel, } from "../api/chat-api"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; import type { LoadModelResponse } from "../types/api"; @@ -177,6 +178,17 @@ export function useChatModelRuntime() { const displayName = model?.name || lora?.name || modelId; const currentCheckpoint = useChatRuntimeStore.getState().params.checkpoint; + const previousCheckpoint = currentCheckpoint; + const previousVariant = + useChatRuntimeStore.getState().activeGgufVariant ?? null; + const previousModel = previousCheckpoint + ? models.find((entry) => entry.id === previousCheckpoint) + : undefined; + const previousLora = previousCheckpoint + ? loras.find((entry) => entry.id === previousCheckpoint) + : undefined; + const previousIsLora = + previousModel?.isLora ?? (previousLora ? true : false); const loadingDescription = [ currentCheckpoint ? "Unloading previous model first." : null, extraLoadingDescription ?? null, @@ -189,24 +201,59 @@ export function useChatModelRuntime() { setLoadingModel({ id: modelId, displayName }); try { async function performLoad(): Promise { + let previousWasUnloaded = false; const currentCheckpoint = useChatRuntimeStore.getState().params.checkpoint; - if (currentCheckpoint) { - await unloadModel({ model_path: currentCheckpoint }); + try { + // Lightweight pre-flight validation: avoid unloading a working model + // if the new identifier is clearly invalid (e.g. bad HF id / path). + await validateModel({ + model_path: modelId, + hf_token: null, + max_seq_length: DEFAULT_MODEL_MAX_SEQ_LENGTH, + load_in_4bit: true, + is_lora: isLora, + gguf_variant: ggufVariant ?? null, + }); + + if (currentCheckpoint) { + await unloadModel({ model_path: currentCheckpoint }); + previousWasUnloaded = true; + } + + const loadResponse = await loadModel({ + model_path: modelId, + hf_token: null, + max_seq_length: DEFAULT_MODEL_MAX_SEQ_LENGTH, + load_in_4bit: true, + is_lora: isLora, + gguf_variant: ggufVariant ?? null, + }); + + const currentParams = useChatRuntimeStore.getState().params; + setParams( + mergeRecommendedInference(currentParams, loadResponse, modelId), + ); + await refresh(); + } catch (error) { + // If we unloaded a previous model and the new load failed, attempt a rollback. + if (previousWasUnloaded && previousCheckpoint) { + try { + await loadModel({ + model_path: previousCheckpoint, + hf_token: null, + max_seq_length: DEFAULT_MODEL_MAX_SEQ_LENGTH, + load_in_4bit: true, + is_lora: previousIsLora, + gguf_variant: previousVariant, + }); + await refresh(); + } catch { + // If rollback also fails, surface the original error. + } + } + throw error; } - - const loadResponse = await loadModel({ - model_path: modelId, - hf_token: null, - max_seq_length: DEFAULT_MODEL_MAX_SEQ_LENGTH, - load_in_4bit: true, - is_lora: isLora, - gguf_variant: ggufVariant ?? null, - }); - - const currentParams = useChatRuntimeStore.getState().params; - setParams(mergeRecommendedInference(currentParams, loadResponse, modelId)); - await refresh(); } const loadPromise = performLoad().finally(() => { diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index edf8fecae3..e86b3b622e 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -33,6 +33,16 @@ export interface LoadModelRequest { gguf_variant?: string | null; } +export interface ValidateModelResponse { + valid: boolean; + message: string; + identifier?: string | null; + display_name?: string | null; + is_gguf?: boolean; + is_lora?: boolean; + is_vision?: boolean; +} + export interface GgufVariantDetail { filename: string; quant: string;