fix(studio): surface live step with null loss through the SSE progress stream (#6206)

* fix(studio): surface live step with null loss through the SSE progress stream

The metric histories skip non-finite steps, so during a NaN stretch the
SSE live loop and final complete event replayed the last finite
step/loss pair. Follow the live progress step when it is ahead of the
history tail and report its loss honestly (null until recovery).

Completes the NaN honesty fix for the SSE consumer flagged in review.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Apply live-step handling to inactive streams and clear the UI loss on null for PR #6206

Fresh /progress connections after a finished run took the inactive branch
which still replayed the last finite step and loss pair; apply the same
live-step correction there. On the frontend, applyProgress kept the stale
currentLoss when a payload advanced the step with a null loss; clear it so
the display shows -- until the loss recovers. Widen the runtime state type
to number | null, which the view layer already handles.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

studio/backend/core/training/training.py

@@ -534,9 +534,7 @@ class TrainingBackend:
                 except (TypeError, ValueError):
                     logger.debug("Could not convert loss to float: %s", _raw_loss)
                     _safe_loss = None
-                _loss_is_nonfinite = (
-                    _safe_loss is not None and not math.isfinite(_safe_loss)
-                )
+                _loss_is_nonfinite = _safe_loss is not None and not math.isfinite(_safe_loss)
                 if _loss_is_nonfinite:
                     # Drop the value rather than laundering it back to the last
                     # finite loss; clients see loss=None at this step so the NaN

studio/backend/core/training/worker.py

@@ -1858,7 +1858,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
     # ── 11. Run training ──
     gc.collect()
     mx.synchronize()
-    trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+    trainer.train(resume_from_checkpoint = resume_from_checkpoint)
 
     # ── 12. Save and finalize ──
     if trainer.stop_requested and not _stop_save[0]:

studio/backend/routes/training.py

@@ -665,10 +665,17 @@ async def stream_training_progress(
 
             # If not active, send final state and exit
             if not is_active:
-                if backend.step_history:
-                    final_step = backend.step_history[-1]
+                _live = (getattr(tp, "step", 0) or 0) if tp else 0
+                if backend.step_history or _live > 0:
+                    final_step = backend.step_history[-1] if backend.step_history else 0
                     final_loss = backend.loss_history[-1] if backend.loss_history else None
                     final_lr = backend.lr_history[-1] if backend.lr_history else None
+                    # Histories skip non-finite steps; report the live step with
+                    # loss=None instead of the last finite pair.
+                    if _live > final_step:
+                        final_step = _live
+                        final_loss = getattr(tp, "loss", None)
+                        final_lr = getattr(tp, "learning_rate", final_lr)
                     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(
@@ -697,11 +704,18 @@ async def stream_training_progress(
 
         while backend.is_training_active():
             try:
-                if backend.step_history:
-                    current_step = backend.step_history[-1]
+                tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None)
+                live_step = (getattr(tp_inner, "step", 0) or 0) if tp_inner else 0
+                if backend.step_history or live_step > 0:
+                    current_step = backend.step_history[-1] if backend.step_history else 0
                     current_loss = backend.loss_history[-1] if backend.loss_history else None
                     current_lr = backend.lr_history[-1] if backend.lr_history else None
-                    tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None)
+                    # Histories skip non-finite steps; follow the live progress
+                    # step and report its loss (None until it recovers).
+                    if live_step > current_step:
+                        current_step = live_step
+                        current_loss = getattr(tp_inner, "loss", None)
+                        current_lr = getattr(tp_inner, "learning_rate", current_lr)
                     current_total_steps = (
                         getattr(tp_inner, "total_steps", current_step) if tp_inner else current_step
                     )
@@ -798,6 +812,13 @@ async def stream_training_progress(
         final_loss = backend.loss_history[-1] if backend.loss_history else None
         final_lr = backend.lr_history[-1] if backend.lr_history else None
         final_tp = getattr(getattr(backend, "trainer", None), "training_progress", None)
+        # If the run ended on a non-finite stretch, report the live step with
+        # loss=None instead of rolling back to the last finite pair.
+        _final_live_step = (getattr(final_tp, "step", 0) or 0) if final_tp else 0
+        if _final_live_step > (final_step if final_step is not None else -1):
+            final_step = _final_live_step
+            final_loss = getattr(final_tp, "loss", None)
+            final_lr = getattr(final_tp, "learning_rate", final_lr)
         final_total_steps = getattr(final_tp, "total_steps", final_step) if final_tp else final_step
         final_epoch = getattr(final_tp, "epoch", None) if final_tp else None
         final_payload = build_progress(

studio/backend/tests/test_training_nan_loss_handling.py

@@ -29,7 +29,11 @@ def _make_backend() -> TrainingBackend:
     return TrainingBackend()
 
 
-def _progress_event(step: int, loss: float, lr: float = 1e-4) -> dict:
+def _progress_event(
+    step: int,
+    loss: float,
+    lr: float = 1e-4,
+) -> dict:
     return {
         "type": "progress",
         "step": step,
@@ -43,7 +47,7 @@ def _progress_event(step: int, loss: float, lr: float = 1e-4) -> dict:
 class TestNonfiniteLossSoftHandling:
     def test_finite_loss_updates_progress_normally(self):
         b = _make_backend()
-        b._handle_event(_progress_event(step=1, loss=0.97))
+        b._handle_event(_progress_event(step = 1, loss = 0.97))
         assert b._progress.loss == pytest.approx(0.97)
         assert b._progress.error is None
         assert b._should_stop is False
@@ -51,9 +55,9 @@ class TestNonfiniteLossSoftHandling:
 
     def test_nan_loss_clears_progress_loss(self):
         b = _make_backend()
-        b._handle_event(_progress_event(step=1, loss=0.97))
+        b._handle_event(_progress_event(step = 1, loss = 0.97))
         assert b._progress.loss == pytest.approx(0.97)
-        b._handle_event(_progress_event(step=2, loss=float("nan")))
+        b._handle_event(_progress_event(step = 2, loss = float("nan")))
         # Stale finite loss must NOT leak through
         assert b._progress.loss is None
         # Run is not marked failed
@@ -64,7 +68,7 @@ class TestNonfiniteLossSoftHandling:
 
     def test_inf_loss_clears_progress_loss(self):
         b = _make_backend()
-        b._handle_event(_progress_event(step=1, loss=float("inf")))
+        b._handle_event(_progress_event(step = 1, loss = float("inf")))
         assert b._progress.loss is None
         assert b._progress.error is None
         assert b._should_stop is False
@@ -72,7 +76,7 @@ class TestNonfiniteLossSoftHandling:
 
     def test_negative_inf_loss_clears_progress_loss(self):
         b = _make_backend()
-        b._handle_event(_progress_event(step=1, loss=float("-inf")))
+        b._handle_event(_progress_event(step = 1, loss = float("-inf")))
         assert b._progress.loss is None
         assert b._progress.error is None
         assert b._should_stop is False
@@ -82,12 +86,12 @@ class TestNonfiniteLossSoftHandling:
         """Subsequent NaN events must not re-fire the warning flag setter.
         The flag should already be True after the first NaN."""
         b = _make_backend()
-        b._handle_event(_progress_event(step=1, loss=0.97))
-        b._handle_event(_progress_event(step=2, loss=float("nan")))
+        b._handle_event(_progress_event(step = 1, loss = 0.97))
+        b._handle_event(_progress_event(step = 2, loss = float("nan")))
         assert b._progress._nonfinite_loss_warned is True
         # Further NaN steps don't change anything we care about
-        b._handle_event(_progress_event(step=3, loss=float("nan")))
-        b._handle_event(_progress_event(step=4, loss=float("nan")))
+        b._handle_event(_progress_event(step = 3, loss = float("nan")))
+        b._handle_event(_progress_event(step = 4, loss = float("nan")))
         assert b._progress._nonfinite_loss_warned is True
         assert b._progress.loss is None
         assert b._progress.error is None
@@ -97,10 +101,10 @@ class TestNonfiniteLossSoftHandling:
         """If a NaN step is followed by a finite step, progress.loss must
         reflect the new finite value (not stay stuck at None)."""
         b = _make_backend()
-        b._handle_event(_progress_event(step=1, loss=0.97))
-        b._handle_event(_progress_event(step=2, loss=float("nan")))
+        b._handle_event(_progress_event(step = 1, loss = 0.97))
+        b._handle_event(_progress_event(step = 2, loss = float("nan")))
         assert b._progress.loss is None
-        b._handle_event(_progress_event(step=3, loss=0.85))
+        b._handle_event(_progress_event(step = 3, loss = 0.85))
         assert b._progress.loss == pytest.approx(0.85)
         # Warning flag stays set (we don't reset it on recovery)
         assert b._progress._nonfinite_loss_warned is True

studio/backend/tests/test_training_progress_stream_nan.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""The SSE progress stream must follow the live progress step during
+non-finite-loss stretches (loss reported as null) instead of replaying the
+last finite step/loss pair from the metric histories, which skip NaN steps."""
+
+import asyncio
+import json
+import sys
+import types
+
+import pytest
+
+if "structlog" not in sys.modules:
+
+    class _DummyLogger:
+        def __getattr__(self, _name):
+            return lambda *args, **kwargs: None
+
+    sys.modules["structlog"] = types.SimpleNamespace(
+        BoundLogger = _DummyLogger,
+        get_logger = lambda *args, **kwargs: _DummyLogger(),
+    )
+
+import routes.training as rt
+
+
+class _Progress:
+    def __init__(self):
+        self.step = 5
+        self.total_steps = 10
+        self.loss = None  # cleared by the NaN honesty fix in core training
+        self.learning_rate = 8e-5
+        self.epoch = 0.1
+        self.grad_norm = None
+        self.num_tokens = None
+        self.eval_loss = None
+        self.elapsed_seconds = None
+        self.eta_seconds = None
+
+
+class _FakeBackend:
+    """Finite history stops at step 2; live progress is at step 5 with NaN
+    (loss=None). Active for a few polls, then done."""
+
+    def __init__(self, active_polls = 2):
+        self.current_job_id = "job-1"
+        self.step_history = [1, 2]
+        self.loss_history = [2.0, 1.5]
+        self.lr_history = [1e-4, 9e-5]
+        self.eval_enabled = False
+        self._active_calls = 0
+        self._active_polls = active_polls
+        self.trainer = types.SimpleNamespace(training_progress = _Progress())
+
+    def is_training_active(self):
+        self._active_calls += 1
+        return self._active_calls <= self._active_polls
+
+
+class _FakeRequest:
+    headers = {}
+
+
+def _collect_events(response, timeout = 15):
+    async def _drain():
+        chunks = []
+        async for chunk in response.body_iterator:
+            chunks.append(chunk)
+        return "".join(c.decode() if isinstance(c, bytes) else c for c in chunks)
+
+    return asyncio.run(asyncio.wait_for(_drain(), timeout))
+
+
+def _progress_payloads(raw):
+    payloads = []
+    for block in raw.split("\n\n"):
+        lines = block.strip().splitlines()
+        data = next((l[6:] for l in lines if l.startswith("data: ")), None)
+        if data:
+            payloads.append(json.loads(data))
+    return payloads
+
+
+def test_stream_reports_live_step_with_null_loss_during_nan(monkeypatch):
+    backend = _FakeBackend(active_polls = 2)
+    monkeypatch.setattr(rt, "get_training_backend", lambda: backend)
+
+    response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester"))
+    raw = _collect_events(response)
+    payloads = _progress_payloads(raw)
+    assert payloads, f"no SSE payloads parsed from: {raw!r}"
+
+    live = [p for p in payloads if p.get("step") == 5]
+    assert live, (
+        "stream never advanced to the live progress step during the NaN "
+        f"stretch; steps seen: {[p.get('step') for p in payloads]}"
+    )
+    assert live[0]["loss"] is None
+    # The stale finite pair must not be re-emitted as the latest progress.
+    stale = [p for p in payloads if p.get("step") == 2 and p.get("loss") == 1.5]
+    assert not stale
+
+
+def test_inactive_stream_completes_with_live_step_and_null_loss(monkeypatch):
+    # Fresh connection after the run already ended during a NaN stretch: the
+    # immediate complete event must not replay the stale finite pair either.
+    backend = _FakeBackend(active_polls = 0)
+    monkeypatch.setattr(rt, "get_training_backend", lambda: backend)
+
+    response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester"))
+    payloads = _progress_payloads(_collect_events(response))
+    final = payloads[-1]
+    assert final["step"] == 5
+    assert final["loss"] is None
+
+
+def test_stream_uses_finite_history_when_progress_in_sync(monkeypatch):
+    backend = _FakeBackend(active_polls = 2)
+    # Live progress agrees with the history tail: normal finite behavior.
+    backend.trainer.training_progress.step = 2
+    backend.trainer.training_progress.loss = 1.5
+    monkeypatch.setattr(rt, "get_training_backend", lambda: backend)
+
+    response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester"))
+    payloads = _progress_payloads(_collect_events(response))
+    finite = [p for p in payloads if p.get("step") == 2]
+    assert finite and finite[0]["loss"] == 1.5

studio/frontend/src/features/training/stores/training-runtime-store.ts

@@ -274,7 +274,10 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => (
         jobId: payload.job_id || state.jobId,
         currentStep: step,
         totalSteps: Math.max(payload.total_steps, state.totalSteps),
-        currentLoss: currentLoss ?? state.currentLoss,
+        // A null loss at a new step means the backend reported a non-finite
+        // loss; clear the display instead of keeping the stale value.
+        currentLoss:
+          currentLoss ?? (step > state.currentStep ? null : state.currentLoss),
         currentLearningRate: currentLearningRate ?? state.currentLearningRate,
         progressPercent: payload.progress_percent,
         currentEpoch: payload.epoch ?? state.currentEpoch,

studio/frontend/src/features/training/types/runtime.ts

@@ -90,7 +90,8 @@ export interface TrainingRuntimeState {
   currentStep: number;
   totalSteps: number;
   currentEpoch: number;
-  currentLoss: number;
+  // null while the latest reported loss is non-finite
+  currentLoss: number | null;
   currentLearningRate: number;
   progressPercent: number;
   elapsedSeconds: number | null;