ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

switch to pure polling for training progress

Browse source
This commit is contained in:
Roland Tannous 2026-03-17 15:17:12 +00:00
commit cd0a90fcc5
2 changed files with 5 additions and 173 deletions
Download patch file Download diff file Expand all files Collapse all files

68
studio/frontend/src/features/training/api/train-api.ts
View file

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

110
studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts
View file

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