From 3733e0b2741e666ebdd677424c2741c348dc9cfc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 07:50:13 -0700 Subject: [PATCH] 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 | 4 +- studio/backend/core/training/worker.py | 2 +- studio/backend/routes/training.py | 31 ++++- .../tests/test_training_nan_loss_handling.py | 30 ++-- .../test_training_progress_stream_nan.py | 129 ++++++++++++++++++ .../training/stores/training-runtime-store.ts | 5 +- .../src/features/training/types/runtime.ts | 3 +- 7 files changed, 180 insertions(+), 24 deletions(-) create mode 100644 studio/backend/tests/test_training_progress_stream_nan.py diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 5bd9bc5338..2764101d93 100644 --- a/studio/backend/core/training/training.py +++ b/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 diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 427687451c..18b25cb4fe 100644 --- a/studio/backend/core/training/worker.py +++ b/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]: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 96403e21a5..d7687ffdee 100644 --- a/studio/backend/routes/training.py +++ b/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( diff --git a/studio/backend/tests/test_training_nan_loss_handling.py b/studio/backend/tests/test_training_nan_loss_handling.py index 0ab444f9bb..a2dc78bee2 100644 --- a/studio/backend/tests/test_training_nan_loss_handling.py +++ b/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 diff --git a/studio/backend/tests/test_training_progress_stream_nan.py b/studio/backend/tests/test_training_progress_stream_nan.py new file mode 100644 index 0000000000..899527a04d --- /dev/null +++ b/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 diff --git a/studio/frontend/src/features/training/stores/training-runtime-store.ts b/studio/frontend/src/features/training/stores/training-runtime-store.ts index f1966f03fb..97fbd32d57 100644 --- a/studio/frontend/src/features/training/stores/training-runtime-store.ts +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -274,7 +274,10 @@ export const useTrainingRuntimeStore = create()((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, diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts index d30c3f75c2..5fd8286d22 100644 --- a/studio/frontend/src/features/training/types/runtime.ts +++ b/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;