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

Merge pull request #33 from unslothai/feature/sse-connection-resilience

Browse source 
feat: SSE Connection Resilience
This commit is contained in:
Roland Tannous 2026-02-12 22:03:20 +04:00 • committed by GitHub
commit 78d2fe5ee3
4 changed files with 467 additions and 50 deletions
Download patch file Download diff file Expand all files Collapse all files

5
studio/backend/models/training.py
View file

@ -84,6 +84,11 @@ class TrainingStatus(BaseModel):
message: str = Field(..., description="Human-readable status message")
error: Optional[str] = Field(None, description="Error details if phase is 'error'")
details: Optional[dict] = Field(None, description="Phase-specific info, e.g. {'model_size': '8B'}")
metric_history: Optional[dict] = Field(
None,
description="Full metric history arrays for chart recovery after SSE reconnection. "
"Keys: 'steps', 'loss', 'lr' — each a list of numeric values.",
)
class TrainingProgress(BaseModel):

188
studio/backend/routes/training.py
View file

@ -3,7 +3,7 @@ Training API routes
"""
import sys
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import StreamingResponse
from typing import Dict, Optional
import logging
@ -337,6 +337,15 @@ async def get_training_status(
"learning_rate": getattr(progress, "learning_rate", 0.0),
}
# Build metric history for chart recovery after SSE reconnection
metric_history = None
if backend.step_history:
metric_history = {
"steps": list(backend.step_history),
"loss": list(backend.loss_history),
"lr": list(backend.lr_history),
}
return TrainingStatus(
job_id=job_id,
phase=phase,
@ -344,6 +353,7 @@ async def get_training_status(
message=status_message,
error=error_message,
details=details,
metric_history=metric_history,
)
except Exception as e:
@ -393,18 +403,34 @@ async def get_training_metrics(
@router.get("/progress")
async def stream_training_progress(
request: Request,
current_subject: str = Depends(get_current_subject),
):
"""
Stream training progress updates using Server-Sent Events (SSE).
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.
"""
# Read Last-Event-ID header for reconnection resume
last_event_id = request.headers.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}")
except ValueError:
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", "")
# Helper to build a TrainingProgress payload from raw values
# ── Helpers ──────────────────────────────────────────────
def build_progress(
step: int,
loss: float,
@ -434,45 +460,86 @@ async def stream_training_progress(
num_tokens=None,
)
# Send initial status
is_active = backend.is_training_active()
tp = getattr(getattr(backend, "trainer", None), "training_progress", None)
initial_total_steps = getattr(tp, "total_steps", 0) if tp else 0
initial_epoch = getattr(tp, "epoch", None) if tp else None
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)
initial_progress = build_progress(
step=0,
loss=0.0,
learning_rate=0.0,
total_steps=initial_total_steps,
epoch=initial_epoch,
)
yield f"data: {initial_progress.model_dump_json()}\n\n"
# ── Retry directive ──────────────────────────────────────
# Tell the browser to reconnect after 3 seconds if the connection drops
yield "retry: 3000\n\n"
# If not active, check if there's any history
if not is_active:
if backend.step_history:
# Training completed - send final metrics
final_step = backend.step_history[-1]
final_loss = backend.loss_history[-1] if backend.loss_history else 0.0
final_lr = backend.lr_history[-1] if backend.lr_history else 0.0
final_total_steps = (
getattr(tp, "total_steps", final_step) if tp else final_step
)
final_epoch = getattr(tp, "epoch", None) if tp else None
yield f"data: {build_progress(final_step, final_loss, final_lr, final_total_steps, final_epoch).model_dump_json()}\n\n"
else:
yield f"data: {build_progress(-1, 0.0, 0.0, 0).model_dump_json()}\n\n"
return
# Poll for updates while training is active
last_step = -1
# ── Replay missed steps on reconnect ─────────────────────
if resume_from_step is not None and backend.step_history:
replayed = 0
for i, step_val in enumerate(backend.step_history):
if step_val > resume_from_step:
loss_val = backend.loss_history[i] if i < len(backend.loss_history) else 0.0
lr_val = backend.lr_history[i] if i < len(backend.lr_history) else 0.0
tp_replay = getattr(
getattr(backend, "trainer", None), "training_progress", None
)
total_replay = getattr(tp_replay, "total_steps", step_val) if tp_replay else step_val
epoch_replay = getattr(tp_replay, "epoch", None) if tp_replay else None
payload = build_progress(step_val, loss_val, lr_val, total_replay, epoch_replay)
yield format_sse(payload.model_dump_json(), event="progress", event_id=step_val)
replayed += 1
if replayed:
logger.info(f"SSE reconnect: replayed {replayed} missed steps")
# ── Initial status (only on fresh connections) ───────────
if resume_from_step is None:
is_active = backend.is_training_active()
tp = getattr(getattr(backend, "trainer", None), "training_progress", None)
initial_total_steps = getattr(tp, "total_steps", 0) if tp else 0
initial_epoch = getattr(tp, "epoch", None) if tp else None
initial_progress = build_progress(
step=0,
loss=0.0,
learning_rate=0.0,
total_steps=initial_total_steps,
epoch=initial_epoch,
)
yield format_sse(initial_progress.model_dump_json(), event="progress", event_id=0)
# If not active, send final state and exit
if not is_active:
if backend.step_history:
final_step = backend.step_history[-1]
final_loss = backend.loss_history[-1] if backend.loss_history else 0.0
final_lr = backend.lr_history[-1] if backend.lr_history else 0.0
final_total_steps = (
getattr(tp, "total_steps", final_step) if tp else final_step
)
final_epoch = getattr(tp, "epoch", None) if tp else None
payload = build_progress(final_step, final_loss, final_lr, final_total_steps, final_epoch)
yield format_sse(payload.model_dump_json(), event="complete", event_id=final_step)
else:
yield format_sse(
build_progress(-1, 0.0, 0.0, 0).model_dump_json(),
event="complete",
event_id=0,
)
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 = 300 # Timeout after 5 minutes
while backend.is_training_active():
try:
# Get current metrics
if backend.step_history:
current_step = backend.step_history[-1]
current_loss = backend.loss_history[-1] if backend.loss_history else 0.0
@ -496,7 +563,11 @@ async def stream_training_progress(
current_total_steps,
current_epoch,
)
yield f"data: {progress_payload.model_dump_json()}\n\n"
yield format_sse(
progress_payload.model_dump_json(),
event="progress",
event_id=current_step,
)
last_step = current_step
no_update_count = 0
else:
@ -510,30 +581,46 @@ async def stream_training_progress(
current_total_steps,
current_epoch,
)
yield f"data: {heartbeat_payload.model_dump_json()}\n\n"
yield format_sse(
heartbeat_payload.model_dump_json(),
event="heartbeat",
event_id=current_step,
)
else:
# No steps yet, but training is active
# No steps yet, but training is active (model loading, etc.)
no_update_count += 1
if no_update_count % 5 == 0:
preparing_payload = build_progress(0, 0.0, 0.0, 0)
yield f"data: {preparing_payload.model_dump_json()}\n\n"
yield format_sse(
preparing_payload.model_dump_json(),
event="heartbeat",
event_id=0,
)
# Timeout check
if no_update_count > max_no_updates:
logger.warning("Progress stream timeout - no updates received")
timeout_payload = build_progress(last_step, 0.0, 0.0, 0)
yield f"data: {timeout_payload.model_dump_json()}\n\n"
yield format_sse(
timeout_payload.model_dump_json(),
event="error",
event_id=last_step if last_step >= 0 else 0,
)
break
await asyncio.sleep(1) # Poll every second
except Exception as e:
logger.error(f"Error in progress stream: {e}", exc_info=True)
error_payload = build_progress(0, 0.0, 0.0, 0)
yield f"data: {error_payload.model_dump_json()}\n\n"
yield format_sse(
error_payload.model_dump_json(),
event="error",
event_id=last_step if last_step >= 0 else 0,
)
break
# Send final status
# ── Final "complete" event ───────────────────────────────
final_step = backend.step_history[-1] if backend.step_history else last_step
final_loss = backend.loss_history[-1] if backend.loss_history else 0.0
final_lr = backend.lr_history[-1] if backend.lr_history else 0.0
@ -551,14 +638,19 @@ async def stream_training_progress(
final_total_steps,
final_epoch,
)
yield f"data: {final_payload.model_dump_json()}\n\n"
yield format_sse(
final_payload.model_dump_json(),
event="complete",
event_id=final_step if final_step >= 0 else 0,
)
return StreamingResponse(
event_generator(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
}
)

0
studio/tests/__init__.py Normal file
View file

320
studio/tests/test_sse_progress.py Normal file
View file

@ -0,0 +1,320 @@
"""
Tests for the SSE training progress endpoint and status fallback.
Validates:
- SSE spec compliance: `retry:`, `id:`, `event:` fields
- Named event types: progress, heartbeat, complete, error
- Last-Event-ID reconnection and history replay
- /status metric_history fallback (Option B)
All tests mock the training backend and bypass auth.
"""
import sys
from pathlib import Path
from typing import Optional
from unittest.mock import MagicMock, patch, PropertyMock
import re
import pytest
# ── Path setup ────────────────────────────────────────────────────
# Add backend root so bare `from routes…`, `from models…` etc. resolve.
_backend_root = Path(__file__).resolve().parent.parent / "backend"
if str(_backend_root) not in sys.path:
sys.path.insert(0, str(_backend_root))
from fastapi.testclient import TestClient
from main import app
from auth.authentication import get_current_subject
# ── Fixtures ──────────────────────────────────────────────────────
def _bypass_auth():
"""Dependency override that skips real JWT validation."""
return "test-user"
def _make_mock_backend(
*,
is_active: bool = False,
step_history: list | None = None,
loss_history: list | None = None,
lr_history: list | None = None,
total_steps: int = 100,
epoch: int | None = 1,
job_id: str = "job_test_001",
):
"""Build a lightweight mock that quacks like TrainingBackend."""
backend = MagicMock()
backend.current_job_id = job_id
backend.step_history = step_history or []
backend.loss_history = loss_history or []
backend.lr_history = lr_history or []
backend.is_training_active.return_value = is_active
backend._training_thread = None
# trainer.training_progress / get_training_progress()
tp = MagicMock()
tp.total_steps = total_steps
tp.epoch = epoch
tp.step = step_history[-1] if step_history else 0
tp.loss = loss_history[-1] if loss_history else 0.0
tp.learning_rate = lr_history[-1] if lr_history else 0.0
tp.status_message = "Training..."
tp.error = None
tp.is_completed = not is_active and bool(step_history)
backend.trainer = MagicMock()
backend.trainer.training_progress = tp
backend.trainer.get_training_progress.return_value = tp
return backend
@pytest.fixture()
def client():
"""TestClient with auth bypassed."""
app.dependency_overrides[get_current_subject] = _bypass_auth
yield TestClient(app)
app.dependency_overrides.clear()
# ── SSE Parsing Helpers ───────────────────────────────────────────
def parse_sse_events(raw: str) -> list[dict]:
"""
Parse raw SSE text into a list of event dicts.
Each dict has optional keys: 'id', 'event', 'data', 'retry'.
"""
events: list[dict] = []
current: dict = {}
for line in raw.split("\n"):
if line.startswith("retry:"):
# retry is a standalone directive, not part of a normal event
events.append({"retry": line.split(":", 1)[1].strip()})
continue
if line.startswith("id:"):
current["id"] = line.split(":", 1)[1].strip()
elif line.startswith("event:"):
current["event"] = line.split(":", 1)[1].strip()
elif line.startswith("data:"):
current["data"] = line.split(":", 1)[1].strip()
elif line == "" and current:
events.append(current)
current = {}
if current:
events.append(current)
return events
# =====================================================================
# Option A — /api/train/progress (SSE)
# =====================================================================
class TestSSERetryDirective:
"""The first thing the stream emits must be `retry: 3000`."""
def test_retry_is_first_event(self, client: TestClient):
mock_backend = _make_mock_backend(is_active=False)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/progress")
assert resp.status_code == 200
assert resp.headers["content-type"].startswith("text/event-stream")
events = parse_sse_events(resp.text)
assert len(events) >= 1
assert events[0] == {"retry": "3000"}
class TestSSEEventFields:
"""Every non-retry event must include `id:`, `event:`, and `data:` fields."""
def test_events_have_id_and_event_type(self, client: TestClient):
mock_backend = _make_mock_backend(
is_active=False,
step_history=[1, 2, 3],
loss_history=[2.0, 1.5, 1.0],
lr_history=[1e-4, 1e-4, 1e-4],
total_steps=3,
)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/progress")
events = parse_sse_events(resp.text)
data_events = [e for e in events if "data" in e]
assert len(data_events) >= 1
for evt in data_events:
assert "id" in evt, f"Missing `id:` field in event: {evt}"
assert "event" in evt, f"Missing `event:` field in event: {evt}"
assert "data" in evt
class TestSSENamedEventTypes:
"""Events use the correct named types: progress, complete, heartbeat, error."""
def test_idle_sends_progress_then_complete(self, client: TestClient):
mock_backend = _make_mock_backend(
is_active=False,
step_history=[10],
loss_history=[1.5],
lr_history=[1e-4],
total_steps=10,
)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/progress")
events = parse_sse_events(resp.text)
data_events = [e for e in events if "event" in e and e.get("event") != "retry"]
event_types = [e["event"] for e in data_events]
assert "progress" in event_types
assert "complete" in event_types
def test_no_history_sends_complete(self, client: TestClient):
mock_backend = _make_mock_backend(is_active=False)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/progress")
events = parse_sse_events(resp.text)
data_events = [e for e in events if "event" in e]
assert any(e["event"] == "complete" for e in data_events)
class TestSSELastEventIDResume:
"""When `Last-Event-ID` header is sent, the server replays missed steps."""
def test_replays_steps_after_last_event_id(self, client: TestClient):
mock_backend = _make_mock_backend(
is_active=False,
step_history=[1, 2, 3, 4, 5],
loss_history=[2.5, 2.0, 1.5, 1.2, 1.0],
lr_history=[1e-4, 1e-4, 1e-4, 1e-4, 1e-4],
total_steps=5,
)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get(
"/api/train/progress",
headers={"Last-Event-ID": "2"},
)
events = parse_sse_events(resp.text)
# Filter to progress events (replayed ones)
progress_events = [e for e in events if e.get("event") == "progress"]
# Steps 3, 4, 5 should have been replayed
replayed_ids = [int(e["id"]) for e in progress_events]
assert 3 in replayed_ids
assert 4 in replayed_ids
assert 5 in replayed_ids
# Steps 1, 2 should NOT be replayed
assert 1 not in replayed_ids
assert 2 not in replayed_ids
def test_no_replay_without_header(self, client: TestClient):
"""Without Last-Event-ID, should start fresh (initial progress event)."""
mock_backend = _make_mock_backend(
is_active=False,
step_history=[1, 2, 3],
loss_history=[2.0, 1.5, 1.0],
lr_history=[1e-4, 1e-4, 1e-4],
total_steps=3,
)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/progress")
events = parse_sse_events(resp.text)
progress_events = [e for e in events if e.get("event") == "progress"]
# Should have initial step=0 progress event
assert any(e.get("id") == "0" for e in progress_events)
def test_invalid_last_event_id_treated_as_fresh(self, client: TestClient):
"""Non-integer Last-Event-ID should be ignored gracefully."""
mock_backend = _make_mock_backend(is_active=False)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get(
"/api/train/progress",
headers={"Last-Event-ID": "not-a-number"},
)
assert resp.status_code == 200
events = parse_sse_events(resp.text)
# Should still work — treated as a fresh connection
assert any(e.get("event") == "progress" or e.get("event") == "complete" for e in events)
class TestSSEResponseHeaders:
"""Verify SSE response headers for proxy compatibility."""
def test_headers(self, client: TestClient):
mock_backend = _make_mock_backend(is_active=False)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/progress")
assert resp.headers["content-type"].startswith("text/event-stream")
assert resp.headers.get("cache-control") == "no-cache"
assert resp.headers.get("x-accel-buffering") == "no"
# =====================================================================
# Option B — /api/train/status (metric_history fallback)
# =====================================================================
class TestStatusMetricHistory:
"""The /status endpoint returns metric_history for chart recovery."""
def test_metric_history_populated_when_history_exists(self, client: TestClient):
mock_backend = _make_mock_backend(
is_active=True,
step_history=[1, 2, 3, 4, 5],
loss_history=[2.5, 2.0, 1.5, 1.2, 1.0],
lr_history=[1e-4, 1e-4, 1e-4, 1e-4, 1e-4],
total_steps=10,
)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/status")
assert resp.status_code == 200
body = resp.json()
assert "metric_history" in body
mh = body["metric_history"]
assert mh is not None
assert mh["steps"] == [1, 2, 3, 4, 5]
assert mh["loss"] == [2.5, 2.0, 1.5, 1.2, 1.0]
assert mh["lr"] == [1e-4, 1e-4, 1e-4, 1e-4, 1e-4]
def test_metric_history_null_when_no_history(self, client: TestClient):
mock_backend = _make_mock_backend(is_active=False)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/status")
assert resp.status_code == 200
body = resp.json()
assert body["metric_history"] is None
def test_status_still_returns_phase_and_details(self, client: TestClient):
"""Ensure adding metric_history didn't break existing fields."""
mock_backend = _make_mock_backend(
is_active=True,
step_history=[5],
loss_history=[1.5],
lr_history=[1e-4],
total_steps=100,
)
with patch("routes.training.get_training_backend", return_value=mock_backend):
resp = client.get("/api/train/status")
body = resp.json()
assert body["phase"] == "training"
assert body["is_training_running"] is True
assert body["job_id"] == "job_test_001"
assert "details" in body