switch to pure polling for training progress

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 <token>
+security = HTTPBearer(auto_error = False)  # Reads Authorization: Bearer <token>
+
+
+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:

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):

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,
     )
 
 

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

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<AuthStatus | null> {
+  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<boolean> {
   if (hasAuthToken()) return true;
   if (!hasRefreshToken()) return false;
@@ -17,28 +33,22 @@ async function hasActiveSession(): Promise<boolean> {
 }
 
 async function checkAuthInitialized(): Promise<boolean> {
-  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<boolean> {
-  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<void> {
+  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<void> {
 }
 
 export async function requireGuest(): Promise<void> {
+  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<void> {
+  const status = await getAuthStatus();
+  if (status?.auth_disabled) {
+    throw redirect({ to: getPostAuthRoute() });
+  }
+
   const requiresPasswordChange = await checkPasswordChangeRequired();
 
   if (requiresPasswordChange) return;

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<TrainingMetricsResponse> {
   return parseJson<TrainingMetricsResponse>(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<void> {
-  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 };

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<typeof setTimeout> | 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();
     };
   }, []);
 }

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