From 509659ba9736dc920b5607f01df9edf3f28761f7 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Thu, 12 Feb 2026 17:58:48 +0000 Subject: [PATCH] feat: add SSE reconnection resilience with spec-compliant event fields, Last-Event-ID resume, and metric_history fallback in /status --- studio/backend/models/training.py | 5 + studio/backend/routes/training.py | 188 +++++++++++++----- studio/tests/__init__.py | 0 studio/tests/test_sse_progress.py | 320 ++++++++++++++++++++++++++++++ 4 files changed, 465 insertions(+), 48 deletions(-) create mode 100644 studio/tests/__init__.py create mode 100644 studio/tests/test_sse_progress.py diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 2afee72e89..3afcb187e7 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -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): diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 9f6c298a4d..ee36a9b450 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -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", } ) diff --git a/studio/tests/__init__.py b/studio/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/studio/tests/test_sse_progress.py b/studio/tests/test_sse_progress.py new file mode 100644 index 0000000000..d2c007d99e --- /dev/null +++ b/studio/tests/test_sse_progress.py @@ -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