From cd0a90fcc5643d28b54e0b15822fce600fb8ebd2 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 17 Mar 2026 15:17:12 +0000 Subject: [PATCH] switch to pure polling for training progress --- .../src/features/training/api/train-api.ts | 68 ----------- .../hooks/use-training-runtime-lifecycle.ts | 110 +----------------- 2 files changed, 5 insertions(+), 173 deletions(-) diff --git a/studio/frontend/src/features/training/api/train-api.ts b/studio/frontend/src/features/training/api/train-api.ts index 10f70deaf8..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,71 +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; -} - -export async function streamTrainingProgress(options: { - signal: AbortSignal; - lastEventId?: number | null; - onOpen?: () => void; - onEvent: (event: ParsedSseEvent) => void; -}): Promise { - // Build WebSocket URL from current page location - const protocol = window.location.protocol === "https:" ? "wss:" : "ws:"; - const params = new URLSearchParams(); - - // Pass auth token as query param (WebSocket can't use Authorization header) - const token = localStorage.getItem("unsloth_auth_token"); - if (token) params.set("token", token); - if (typeof options.lastEventId === "number") { - params.set("last_event_id", String(options.lastEventId)); - } - - const url = `${protocol}//${window.location.host}/api/train/progress/ws?${params}`; - - return new Promise((resolve, reject) => { - const ws = new WebSocket(url); - - // Wire up AbortSignal to close the socket - const onAbort = () => ws.close(); - options.signal.addEventListener("abort", onAbort); - - ws.onopen = () => { - options.onOpen?.(); - }; - - ws.onmessage = (messageEvent) => { - try { - const msg = JSON.parse(messageEvent.data) as { - event: ProgressEventName; - id: number | null; - data: TrainingProgressPayload; - }; - options.onEvent({ - event: msg.event, - id: msg.id, - payload: msg.data, - }); - } catch { - // Ignore parse errors for malformed messages - } - }; - - ws.onclose = () => { - options.signal.removeEventListener("abort", onAbort); - resolve(); - }; - - ws.onerror = () => { - options.signal.removeEventListener("abort", onAbort); - reject(new Error("WebSocket connection failed")); - }; - }); -} - 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(); }; }, []); }