diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index b39f915764..409dee29a9 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -1,6 +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 os import secrets from datetime import datetime, timedelta, timezone from typing import Optional, Tuple @@ -21,7 +22,28 @@ ALGORITHM = "HS256" ACCESS_TOKEN_EXPIRE_MINUTES = 60 REFRESH_TOKEN_EXPIRE_DAYS = 7 -security = HTTPBearer() # Reads Authorization: Bearer +security = HTTPBearer(auto_error = False) # Reads Authorization: Bearer + + +def is_auth_disabled() -> bool: + """ + Return True when auth should be bypassed. + + For the HF Spaces branch, auth is disabled by default so users can access + Studio without login/signup. You can explicitly re-enable auth with: + UNSLOTH_STUDIO_ENABLE_AUTH=1 + """ + raw_enable = os.getenv("UNSLOTH_STUDIO_ENABLE_AUTH", "").strip().lower() + if raw_enable in {"1", "true", "yes", "on"}: + return False + + raw = os.getenv("UNSLOTH_STUDIO_DISABLE_AUTH", "").strip().lower() + if raw in {"1", "true", "yes", "on"}: + return True + if raw in {"0", "false", "no", "off"}: + return False + + return True def _get_secret_for_subject(subject: str) -> str: @@ -103,7 +125,7 @@ def reload_secret() -> None: async def get_current_subject( - credentials: HTTPAuthorizationCredentials = Depends(security), + credentials: Optional[HTTPAuthorizationCredentials] = Depends(security), ) -> str: """Validate JWT and require the password-change flow to be completed.""" return await _get_current_subject( @@ -113,7 +135,7 @@ async def get_current_subject( async def get_current_subject_allow_password_change( - credentials: HTTPAuthorizationCredentials = Depends(security), + credentials: Optional[HTTPAuthorizationCredentials] = Depends(security), ) -> str: """Validate JWT but allow access to the password-change endpoint.""" return await _get_current_subject( @@ -123,7 +145,7 @@ async def get_current_subject_allow_password_change( async def _get_current_subject( - credentials: HTTPAuthorizationCredentials, + credentials: Optional[HTTPAuthorizationCredentials], *, allow_password_change: bool, ) -> str: @@ -136,6 +158,15 @@ async def _get_current_subject( async def secure_endpoint(current_subject: str = Depends(get_current_subject)): ... """ + if is_auth_disabled(): + return "hf-space-user" + + if credentials is None: + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Missing bearer token", + ) + token = credentials.credentials subject = _decode_subject_without_verification(token) if subject is None: diff --git a/studio/backend/models/auth.py b/studio/backend/models/auth.py index 73d21130ae..edba6ac250 100644 --- a/studio/backend/models/auth.py +++ b/studio/backend/models/auth.py @@ -34,6 +34,10 @@ class AuthStatusResponse(BaseModel): ..., description = "True if the seeded admin must still change the default password", ) + auth_disabled: bool = Field( + default = False, + description = "True when auth is bypassed (HF Spaces mode).", + ) class ChangePasswordRequest(BaseModel): diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index db37ed837d..92b6df08c6 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -20,6 +20,7 @@ from auth.authentication import ( create_refresh_token, get_current_subject, get_current_subject_allow_password_change, + is_auth_disabled, refresh_access_token, ) @@ -34,6 +35,14 @@ async def auth_status() -> AuthStatusResponse: - initialized = False -> frontend should wait for the seeded admin bootstrap. - initialized = True -> frontend should show login or force the first password change. """ + if is_auth_disabled(): + return AuthStatusResponse( + initialized = True, + default_username = storage.DEFAULT_ADMIN_USERNAME, + requires_password_change = False, + auth_disabled = True, + ) + return AuthStatusResponse( initialized = storage.is_initialized(), default_username = storage.DEFAULT_ADMIN_USERNAME, @@ -42,6 +51,7 @@ async def auth_status() -> AuthStatusResponse: ) if storage.is_initialized() else True, + auth_disabled = False, ) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 4f8054f80e..b8ed1d1610 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -7,8 +7,9 @@ Training API routes import sys from pathlib import Path -from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket, WebSocketDisconnect from fastapi.responses import StreamingResponse +from starlette.websockets import WebSocketState from typing import Dict, Optional, Any import structlog from loggers import get_logger @@ -494,101 +495,108 @@ async def get_training_metrics( ) -@router.get("/progress") -async def stream_training_progress( - request: Request, - current_subject: str = Depends(get_current_subject), -): +@router.websocket("/progress/ws") +async def ws_training_progress(websocket: WebSocket): """ - Stream training progress updates using Server-Sent Events (SSE). + Stream training progress updates over WebSocket. This endpoint provides real-time updates on training progress. - Supports reconnection via the SSE spec: - - Sends `id:` with each event so the browser tracks position. - - Sends `retry:` to control reconnection interval. - - Sends named `event:` types (progress, heartbeat, complete, error). - - Reads `Last-Event-ID` header on reconnect to replay missed steps. + Supports reconnection via query params: + - `last_event_id`: resume from a specific step on reconnect. + - `token`: JWT auth token (WebSocket can't use Authorization header). """ - # Read Last-Event-ID header for reconnection resume - last_event_id = request.headers.get("last-event-id") + # Auth: WebSocket can't use Authorization header, so accept token as query param + from auth.authentication import is_auth_disabled + + if not is_auth_disabled(): + token = websocket.query_params.get("token") + if not token: + await websocket.close(code = 4001, reason = "Missing auth token") + return + from auth.authentication import _decode_subject_without_verification + from auth.storage import get_user_and_secret + + subject = _decode_subject_without_verification(token) + if subject is None: + await websocket.close(code = 4001, reason = "Invalid token") + return + record = get_user_and_secret(subject) + if record is None: + await websocket.close(code = 4001, reason = "Invalid or expired token") + return + + await websocket.accept() + + # Read optional last_event_id from query params for reconnection resume + last_event_id = websocket.query_params.get("last_event_id") resume_from_step: Optional[int] = None if last_event_id is not None: try: resume_from_step = int(last_event_id) - logger.info(f"SSE reconnect: resuming from step {resume_from_step}") + logger.info(f"WebSocket reconnect: resuming from step {resume_from_step}") except ValueError: - logger.warning(f"Invalid Last-Event-ID: {last_event_id}") + logger.warning(f"Invalid last_event_id: {last_event_id}") - async def event_generator(): - backend = get_training_backend() - job_id: str = getattr(backend, "current_job_id", "") or "" + backend = get_training_backend() + job_id: str = getattr(backend, "current_job_id", "") or "" - # ── Helpers ────────────────────────────────────────────── - def build_progress( - step: int, - loss: float, - learning_rate: float, - total_steps: int, - epoch: Optional[float] = None, - progress: Optional[Any] = None, - grad_norm_override: Optional[float] = None, - eval_loss_override: Optional[float] = None, - ) -> TrainingProgress: - total = max(total_steps, 0) - if step < 0 or total == 0: - progress_percent = 0.0 - else: - progress_percent = ( - float(step) / float(total) * 100.0 if total > 0 else 0.0 - ) - - # Get actual values from progress object if available - elapsed_seconds = ( - getattr(progress, "elapsed_seconds", None) if progress else None - ) - eta_seconds = getattr(progress, "eta_seconds", None) if progress else None - grad_norm = grad_norm_override - if grad_norm is None and progress: - grad_norm = getattr(progress, "grad_norm", None) - num_tokens = getattr(progress, "num_tokens", None) if progress else None - eval_loss = eval_loss_override - if eval_loss is None and progress: - eval_loss = getattr(progress, "eval_loss", None) - - return TrainingProgress( - job_id = job_id, - step = step, - total_steps = total, - loss = loss, - learning_rate = learning_rate, - progress_percent = progress_percent, - epoch = epoch, - elapsed_seconds = elapsed_seconds, - eta_seconds = eta_seconds, - grad_norm = grad_norm, - num_tokens = num_tokens, - eval_loss = eval_loss, + # ── Helpers ────────────────────────────────────────────── + def build_progress( + step: int, + loss: float, + learning_rate: float, + total_steps: int, + epoch: Optional[float] = None, + progress: Optional[Any] = None, + grad_norm_override: Optional[float] = None, + eval_loss_override: Optional[float] = None, + ) -> TrainingProgress: + total = max(total_steps, 0) + if step < 0 or total == 0: + progress_percent = 0.0 + else: + progress_percent = ( + float(step) / float(total) * 100.0 if total > 0 else 0.0 ) - def format_sse( - data: str, - event: str = "progress", - event_id: Optional[int] = None, - ) -> str: - """Format a single SSE message with id/event/data fields.""" - lines = [] - if event_id is not None: - lines.append(f"id: {event_id}") - lines.append(f"event: {event}") - lines.append(f"data: {data}") - lines.append("") # trailing blank line - lines.append("") # double newline terminates the event - return "\n".join(lines) + elapsed_seconds = ( + getattr(progress, "elapsed_seconds", None) if progress else None + ) + eta_seconds = getattr(progress, "eta_seconds", None) if progress else None + grad_norm = grad_norm_override + if grad_norm is None and progress: + grad_norm = getattr(progress, "grad_norm", None) + num_tokens = getattr(progress, "num_tokens", None) if progress else None + eval_loss = eval_loss_override + if eval_loss is None and progress: + eval_loss = getattr(progress, "eval_loss", None) - # ── Retry directive ────────────────────────────────────── - # Tell the browser to reconnect after 3 seconds if the connection drops - yield "retry: 3000\n\n" + return TrainingProgress( + job_id = job_id, + step = step, + total_steps = total, + loss = loss, + learning_rate = learning_rate, + progress_percent = progress_percent, + epoch = epoch, + elapsed_seconds = elapsed_seconds, + eta_seconds = eta_seconds, + grad_norm = grad_norm, + num_tokens = num_tokens, + eval_loss = eval_loss, + ) + async def send_event(event: str, event_id: Optional[int], payload: TrainingProgress): + """Send a typed event over WebSocket as JSON.""" + if websocket.client_state != WebSocketState.CONNECTED: + return + await websocket.send_json({ + "event": event, + "id": event_id, + "data": payload.model_dump(), + }) + + try: # ── Replay missed steps on reconnect ───────────────────── if resume_from_step is not None and backend.step_history: replayed = 0 @@ -629,12 +637,10 @@ async def stream_training_progress( progress = tp_replay, grad_norm_override = grad_norm_by_step.get(step_val), ) - yield format_sse( - payload.model_dump_json(), event = "progress", event_id = step_val - ) + await send_event("progress", step_val, payload) replayed += 1 if replayed: - logger.info(f"SSE reconnect: replayed {replayed} missed steps") + logger.info(f"WebSocket reconnect: replayed {replayed} missed steps") # ── Initial status (only on fresh connections) ─────────── if resume_from_step is None: @@ -651,11 +657,9 @@ async def stream_training_progress( epoch = initial_epoch, progress = tp, ) - yield format_sse( - initial_progress.model_dump_json(), event = "progress", event_id = 0 - ) + await send_event("progress", 0, initial_progress) - # If not active, send final state and exit + # If not active, send final state and close if not is_active: if backend.step_history: final_step = backend.step_history[-1] @@ -675,23 +679,18 @@ async def stream_training_progress( final_epoch, progress = tp, ) - yield format_sse( - payload.model_dump_json(), event = "complete", event_id = final_step - ) + await send_event("complete", final_step, payload) else: - yield format_sse( - build_progress(-1, 0.0, 0.0, 0, progress = tp).model_dump_json(), - event = "complete", - event_id = 0, + await send_event( + "complete", 0, + build_progress(-1, 0.0, 0.0, 0, progress = tp), ) return # ── Live polling loop ──────────────────────────────────── last_step = resume_from_step if resume_from_step is not None else -1 no_update_count = 0 - max_no_updates = ( - 1800 # Timeout after 30 minutes (large models need time for compilation) - ) + max_no_updates = 1800 # Timeout after 30 min while backend.is_training_active(): try: @@ -723,11 +722,7 @@ async def stream_training_progress( current_epoch, progress = tp_inner, ) - yield format_sse( - progress_payload.model_dump_json(), - event = "progress", - event_id = current_step, - ) + await send_event("progress", current_step, progress_payload) last_step = current_step no_update_count = 0 else: @@ -742,17 +737,11 @@ async def stream_training_progress( current_epoch, progress = tp_inner, ) - yield format_sse( - heartbeat_payload.model_dump_json(), - event = "heartbeat", - event_id = current_step, - ) + await send_event("heartbeat", current_step, heartbeat_payload) else: # No steps yet, but training is active (model loading, etc.) no_update_count += 1 if no_update_count % 5 == 0: - # Pull total_steps and status from trainer so - # the frontend can show "Tokenizing…" etc. tp_prep = getattr( getattr(backend, "trainer", None), "training_progress", @@ -762,17 +751,9 @@ async def stream_training_progress( getattr(tp_prep, "total_steps", 0) if tp_prep else 0 ) preparing_payload = build_progress( - 0, - 0.0, - 0.0, - prep_total, - progress = tp_prep, - ) - yield format_sse( - preparing_payload.model_dump_json(), - event = "heartbeat", - event_id = 0, + 0, 0.0, 0.0, prep_total, progress = tp_prep, ) + await send_event("heartbeat", 0, preparing_payload) # Timeout check if no_update_count > max_no_updates: @@ -783,10 +764,8 @@ async def stream_training_progress( timeout_payload = build_progress( last_step, 0.0, 0.0, 0, progress = tp_timeout ) - yield format_sse( - timeout_payload.model_dump_json(), - event = "error", - event_id = last_step if last_step >= 0 else 0, + await send_event( + "error", last_step if last_step >= 0 else 0, timeout_payload ) break @@ -798,10 +777,8 @@ async def stream_training_progress( getattr(backend, "trainer", None), "training_progress", None ) error_payload = build_progress(0, 0.0, 0.0, 0, progress = tp_error) - yield format_sse( - error_payload.model_dump_json(), - event = "error", - event_id = last_step if last_step >= 0 else 0, + await send_event( + "error", last_step if last_step >= 0 else 0, error_payload ) break @@ -822,18 +799,16 @@ async def stream_training_progress( final_epoch, progress = final_tp, ) - yield format_sse( - final_payload.model_dump_json(), - event = "complete", - event_id = final_step if final_step >= 0 else 0, + await send_event( + "complete", final_step if final_step >= 0 else 0, final_payload ) - return StreamingResponse( - event_generator(), - media_type = "text/event-stream", - headers = { - "Cache-Control": "no-cache", - "Connection": "keep-alive", - "X-Accel-Buffering": "no", - }, - ) + except WebSocketDisconnect: + logger.info("WebSocket client disconnected") + except Exception as e: + logger.error(f"WebSocket error: {e}", exc_info = True) + try: + if websocket.client_state == WebSocketState.CONNECTED: + await websocket.send_json({"event": "error", "id": None, "data": {"error": str(e)}}) + except Exception: + pass diff --git a/studio/frontend/src/app/auth-guards.ts b/studio/frontend/src/app/auth-guards.ts index 1dcdfcb143..27453e2bfe 100644 --- a/studio/frontend/src/app/auth-guards.ts +++ b/studio/frontend/src/app/auth-guards.ts @@ -10,6 +10,22 @@ import { refreshSession, } from "@/features/auth"; +type AuthStatus = { + initialized: boolean; + requires_password_change: boolean; + auth_disabled?: boolean; +}; + +async function getAuthStatus(): Promise { + try { + const res = await fetch("/api/auth/status"); + if (!res.ok) return null; + return (await res.json()) as AuthStatus; + } catch { + return null; + } +} + async function hasActiveSession(): Promise { if (hasAuthToken()) return true; if (!hasRefreshToken()) return false; @@ -17,28 +33,22 @@ async function hasActiveSession(): Promise { } async function checkAuthInitialized(): Promise { - try { - const res = await fetch("/api/auth/status"); - if (!res.ok) return true; // fallback to login on error - const data = (await res.json()) as { initialized: boolean }; - return data.initialized; - } catch { - return true; // fallback to login on error - } + const status = await getAuthStatus(); + if (status?.auth_disabled) return true; + return status?.initialized ?? true; // fallback to login on error } async function checkPasswordChangeRequired(): Promise { - try { - const res = await fetch("/api/auth/status"); - if (!res.ok) return mustChangePassword(); - const data = (await res.json()) as { requires_password_change: boolean }; - return data.requires_password_change || mustChangePassword(); - } catch { - return mustChangePassword(); - } + const status = await getAuthStatus(); + if (status?.auth_disabled) return false; + if (!status) return mustChangePassword(); + return status.requires_password_change || mustChangePassword(); } export async function requireAuth(): Promise { + const status = await getAuthStatus(); + if (status?.auth_disabled) return; + if (await hasActiveSession()) { if (await checkPasswordChangeRequired()) { throw redirect({ to: "/change-password" }); @@ -52,11 +62,21 @@ export async function requireAuth(): Promise { } export async function requireGuest(): Promise { + const status = await getAuthStatus(); + if (status?.auth_disabled) { + throw redirect({ to: getPostAuthRoute() }); + } + if (!(await hasActiveSession())) return; throw redirect({ to: getPostAuthRoute() }); } export async function requirePasswordChangeFlow(): Promise { + const status = await getAuthStatus(); + if (status?.auth_disabled) { + throw redirect({ to: getPostAuthRoute() }); + } + const requiresPasswordChange = await checkPasswordChangeRequired(); if (requiresPasswordChange) return; diff --git a/studio/frontend/src/features/training/api/train-api.ts b/studio/frontend/src/features/training/api/train-api.ts index bb18fb34aa..6664e5e89b 100644 --- a/studio/frontend/src/features/training/api/train-api.ts +++ b/studio/frontend/src/features/training/api/train-api.ts @@ -9,7 +9,6 @@ import type { } from "../types/api"; import type { TrainingMetricsResponse, - TrainingProgressPayload, TrainingStatusResponse, } from "../types/runtime"; @@ -70,118 +69,4 @@ export async function getTrainingMetrics(): Promise { return parseJson(response); } -type ProgressEventName = "progress" | "heartbeat" | "complete" | "error"; - -interface ParsedSseEvent { - event: ProgressEventName; - payload: TrainingProgressPayload; - id: number | null; -} - -function parseSseEvent(rawEvent: string): ParsedSseEvent | null { - const lines = rawEvent.split(/\r?\n/); - let eventName: ProgressEventName = "progress"; - let id: number | null = null; - const dataLines: string[] = []; - - for (const line of lines) { - if (!line) { - continue; - } - if (line.startsWith("event:")) { - const value = line.slice(6).trim(); - if ( - value === "progress" || - value === "heartbeat" || - value === "complete" || - value === "error" - ) { - eventName = value; - } - continue; - } - if (line.startsWith("id:")) { - const value = Number(line.slice(3).trim()); - id = Number.isFinite(value) ? value : null; - continue; - } - if (line.startsWith("data:")) { - dataLines.push(line.slice(5).trimStart()); - } - } - - if (dataLines.length === 0) { - return null; - } - - const parsed = JSON.parse(dataLines.join("\n")) as TrainingProgressPayload; - return { event: eventName, payload: parsed, id }; -} - -export async function streamTrainingProgress(options: { - signal: AbortSignal; - lastEventId?: number | null; - onOpen?: () => void; - onEvent: (event: ParsedSseEvent) => void; -}): Promise { - const headers = new Headers(); - if (typeof options.lastEventId === "number") { - headers.set("Last-Event-ID", String(options.lastEventId)); - } - - const response = await authFetch("/api/train/progress", { - method: "GET", - headers, - signal: options.signal, - }); - - if (!response.ok) { - throw new Error(await readError(response)); - } - - if (!response.body) { - throw new Error("Progress stream unavailable"); - } - - options.onOpen?.(); - - const reader = response.body.getReader(); - const decoder = new TextDecoder(); - let buffer = ""; - - while (true) { - const { value, done } = await reader.read(); - if (done) { - break; - } - - buffer += decoder.decode(value, { stream: true }); - - let separatorIndex = buffer.search(/\r?\n\r?\n/); - while (separatorIndex >= 0) { - const rawEvent = buffer.slice(0, separatorIndex); - const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2; - buffer = buffer.slice(separatorIndex + separatorLength); - - if (rawEvent.startsWith("retry:")) { - separatorIndex = buffer.search(/\r?\n\r?\n/); - continue; - } - - try { - const event = parseSseEvent(rawEvent); - if (event) { - options.onEvent(event); - } - } catch (error) { - if (!isAbortError(error)) { - throw error; - } - } - - separatorIndex = buffer.search(/\r?\n\r?\n/); - } - } -} - export { isAbortError }; diff --git a/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts b/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts index 5965e07eaa..71d392238e 100644 --- a/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts +++ b/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts @@ -7,44 +7,18 @@ import { getTrainingMetrics, getTrainingStatus, isAbortError, - streamTrainingProgress, } from "../api/train-api"; import { useTrainingRuntimeStore } from "../stores/training-runtime-store"; -import type { TrainingRuntimeStore } from "../types/runtime"; -const STATUS_POLL_INTERVAL_MS = 3000; -const METRICS_POLL_INTERVAL_MS = 5000; -const STREAM_RECONNECT_DELAY_MS = 1500; - -function shouldUseLiveSync(state: TrainingRuntimeStore): boolean { - return state.isTrainingRunning || state.phase === "training"; -} +const STATUS_POLL_INTERVAL_MS = 2000; +const METRICS_POLL_INTERVAL_MS = 3000; export function useTrainingRuntimeLifecycle(): void { useEffect(() => { let disposed = false; - let openingStream = false; - let streamController: AbortController | null = null; - let reconnectTimer: ReturnType | null = null; const runtimeStore = useTrainingRuntimeStore; - const clearReconnect = () => { - if (reconnectTimer) { - clearTimeout(reconnectTimer); - reconnectTimer = null; - } - }; - - const stopStream = () => { - clearReconnect(); - if (streamController) { - streamController.abort(); - streamController = null; - } - runtimeStore.getState().setSseConnected(false); - }; - const pollMetrics = async () => { if (!hasAuthToken()) return; const gen = runtimeStore.getState().resetGeneration; @@ -56,7 +30,7 @@ export function useTrainingRuntimeLifecycle(): void { runtimeStore.getState().applyMetrics(metrics); } catch (error) { if (!isAbortError(error) && !disposed && hasAuthToken()) { - runtimeStore.getState().setSseConnected(false); + // silent — next poll will retry } } }; @@ -69,83 +43,10 @@ export function useTrainingRuntimeLifecycle(): void { if (disposed || runtimeStore.getState().resetGeneration !== gen) { return; } - runtimeStore.getState().applyStatus(status); - - const nextState = runtimeStore.getState(); - if (shouldUseLiveSync(nextState)) { - void ensureStream(); - } else { - stopStream(); - } } catch (error) { if (!isAbortError(error) && !disposed && hasAuthToken()) { - runtimeStore.getState().setSseConnected(false); - } - } - }; - - const ensureStream = async () => { - const state = runtimeStore.getState(); - if ( - disposed || - openingStream || - streamController || - !shouldUseLiveSync(state) - ) { - return; - } - - clearReconnect(); - openingStream = true; - const controller = new AbortController(); - streamController = controller; - - try { - await streamTrainingProgress({ - signal: controller.signal, - lastEventId: state.lastEventId, - onOpen: () => { - runtimeStore.getState().setSseConnected(true); - }, - onEvent: (event) => { - const liveStore = runtimeStore.getState(); - if (typeof event.id === "number") { - liveStore.setLastEventId(event.id); - } - - liveStore.applyProgress(event.payload, event.id ?? undefined); - - if (event.event === "complete") { - void pollStatus(); - void pollMetrics(); - stopStream(); - } - - if (event.event === "error") { - liveStore.setRuntimeError("Training stream error"); - stopStream(); - } - }, - }); - } catch (error) { - if (!disposed && !controller.signal.aborted && !isAbortError(error)) { - runtimeStore.getState().setSseConnected(false); - } - } finally { - openingStream = false; - if (streamController === controller) { - streamController = null; - } - runtimeStore.getState().setSseConnected(false); - - if (!disposed && !controller.signal.aborted) { - const liveState = runtimeStore.getState(); - if (shouldUseLiveSync(liveState)) { - reconnectTimer = setTimeout(() => { - void ensureStream(); - }, STREAM_RECONNECT_DELAY_MS); - } + // silent — next poll will retry } } }; @@ -170,7 +71,7 @@ export function useTrainingRuntimeLifecycle(): void { const metricsTimer = setInterval(() => { const state = runtimeStore.getState(); - if (shouldUseLiveSync(state) || state.currentStep > 0) { + if (state.isTrainingRunning || state.phase === "training" || state.currentStep > 0) { void pollMetrics(); } }, METRICS_POLL_INTERVAL_MS); @@ -179,7 +80,6 @@ export function useTrainingRuntimeLifecycle(): void { disposed = true; clearInterval(statusTimer); clearInterval(metricsTimer); - stopStream(); }; }, []); } diff --git a/studio/setup.sh b/studio/setup.sh index 4c8a6c7dde..0e5121aa20 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -325,19 +325,34 @@ rm -rf "$LLAMA_CPP_DIR" # Detect GPU compute capability and limit CUDA architectures # Without this, cmake builds for ALL default archs (very slow) + # CUDA_ARCHS can be pre-set via env var (e.g. Docker ARG) to + # include targets not detectable at build time (no GPU access). + # We always include 86 (A10) and merge with detected GPUs. + _add_arch() { + local _a="$1" + case ";$CUDA_ARCHS;" in + *";$_a;"*) ;; + *) CUDA_ARCHS="${CUDA_ARCHS:+$CUDA_ARCHS;}$_a" ;; + esac + } CUDA_ARCHS="" + # Merge any pre-set CUDA_ARCHS from environment + if [ -n "${CUDA_ARCHS_EXTRA:-}" ]; then + IFS=';' read -ra _env_archs <<< "$CUDA_ARCHS_EXTRA" + for _ea in "${_env_archs[@]}"; do + _ea=$(echo "$_ea" | tr -d '[:space:]') + [ -n "$_ea" ] && _add_arch "$_ea" + done + fi + # Always include sm_86 (A10) + _add_arch "86" if command -v nvidia-smi &>/dev/null; then # Read all GPUs, deduplicate (handles mixed-GPU hosts) _raw_caps=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true) while IFS= read -r _cap; do _cap=$(echo "$_cap" | tr -d '[:space:]') if [[ "$_cap" =~ ^([0-9]+)\.([0-9]+)$ ]]; then - _arch="${BASH_REMATCH[1]}${BASH_REMATCH[2]}" - # Append if not already present - case ";$CUDA_ARCHS;" in - *";$_arch;"*) ;; - *) CUDA_ARCHS="${CUDA_ARCHS:+$CUDA_ARCHS;}$_arch" ;; - esac + _add_arch "${BASH_REMATCH[1]}${BASH_REMATCH[2]}" fi done <<< "$_raw_caps" fi