Studio: keep the live progress stream alive during pre-first-step preparation (#6665)

* Studio: don't time out the live progress stream during pre-first-step prep

The live progress SSE counts every 1s poll without a step update toward a
30-minute stall timeout, after which it emits an error event and ends the
stream. But that counter also runs during the pre-first-step phase (model
load + tokenizing the dataset), which is never reset because no step has
happened yet. On a large dataset that prep can take well over 30 minutes, so
the live view is torn down with an error while the run is perfectly healthy
and still preparing -- the run then trains on in the background with the UI
showing nothing, exactly the "no progress for hours" decoupling.

Apply the stall timeout only once the stream has actually seen a live step.
Before the first step the run is preparing and may legitimately emit no step
for a long time; heartbeats still flow so the client stays connected and the
worker's liveness still ends the loop when training finishes. A genuine
post-step stall still times out. Extracted the threshold to a module constant
so it can be tuned/tested.

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

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

* Studio: seed seen_live_step from the resume point on reconnect

Review follow-up: seen_live_step reset to False on every SSE request, so a
client reconnecting past the first step (Last-Event-ID set, or the run already
has step history) only receives heartbeats and never flips it true. A worker
that hangs after step N would then never trip the stall timeout for that
reconnected client. Initialize it from resume_from_step / existing step
history so reconnects keep the post-step timeout behavior, while a genuine
pre-first-step run still stays exempt. Added a reconnect regression test.

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

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

* Tighten prep-phase progress timeout comments

Condense the verbose explanatory comments and docstring on the prep-phase stall
timeout exemption to shorter, clearer forms. Comment/whitespace only; verified no
code changed via AST diff. No behaviour change.

---------

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

studio/backend/routes/training.py

@@ -68,6 +68,11 @@ class TrainingStopRequest(PydanticBaseModel):
 router = APIRouter()
 logger = get_logger(__name__)
 
+# Consecutive 1s polls without a step update that count as a stall. Applied only
+# once stepping: the pre-first-step phase (model load + tokenization) can take far
+# longer, and timing out there made a healthy long-prep run look frozen.
+_PROGRESS_STALL_TIMEOUT_POLLS = 1800  # ~30 min at 1 poll/sec
+
 
 def _validate_local_dataset_paths(paths: list[str], label: str = "Local dataset") -> list[str]:
     """Resolve and validate a list of local dataset paths. Returns validated absolute paths."""
@@ -833,7 +838,13 @@ async def stream_training_progress(
         # ── Live polling loop ────────────────────────────────────
         last_step = resume_from_step if resume_from_step is not None else -1
         no_update_count = 0
-        max_no_updates = 1800  # Timeout after 30 min (large models need compile time)
+        # The stall timeout applies only once the run is stepping (pre-step prep
+        # may legitimately emit no step for a long time). On reconnect to an
+        # already-stepping run, seed from the resume point / history, else a worker
+        # that hangs after step N never times out for a client that reconnects past it.
+        seen_live_step = (resume_from_step is not None and resume_from_step > 0) or bool(
+            backend.step_history
+        )
 
         while backend.is_training_active():
             try:
@@ -871,6 +882,7 @@ async def stream_training_progress(
                         )
                         last_step = current_step
                         no_update_count = 0
+                        seen_live_step = True
                     else:
                         no_update_count += 1
                         # Heartbeat every 10 seconds.
@@ -913,8 +925,9 @@ async def stream_training_progress(
                             event_id = 0,
                         )
 
-                # Timeout check
-                if no_update_count > max_no_updates:
+                # Fires only once stepping: a long pre-first-step prep phase is not
+                # a stall, and ending the stream there made a healthy run look frozen.
+                if seen_live_step and no_update_count > _PROGRESS_STALL_TIMEOUT_POLLS:
                     logger.warning("Progress stream timeout - no updates received")
                     tp_timeout = getattr(
                         getattr(backend, "trainer", None), "training_progress", None

studio/backend/tests/test_training_progress_prep_timeout.py

@@ -0,0 +1,141 @@
+# 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 live progress SSE must not time out during the pre-first-step phase.
+
+A large model load / dataset tokenization can keep a run at step 0 for longer
+than the stall timeout. Treating that as a stall ends the live stream and makes a
+healthy run look frozen, so the timeout must apply only once the run is stepping.
+"""
+
+import asyncio
+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,
+        step = 0,
+        total_steps = 1000,
+    ):
+        self.step = step
+        self.total_steps = total_steps
+        self.loss = None
+        self.learning_rate = None
+        self.epoch = None
+        self.grad_norm = None
+        self.num_tokens = None
+        self.eval_loss = None
+        self.elapsed_seconds = None
+        self.eta_seconds = None
+
+
+class _Backend:
+    def __init__(
+        self,
+        *,
+        active_polls,
+        step_history = None,
+        live_step = 0,
+    ):
+        self.current_job_id = "job-prep"
+        self.step_history = list(step_history or [])
+        self.loss_history = [1.0 for _ in self.step_history]
+        self.lr_history = [1e-4 for _ in self.step_history]
+        self.eval_enabled = False
+        self._active_calls = 0
+        self._active_polls = active_polls
+        self.trainer = types.SimpleNamespace(training_progress = _Progress(step = live_step))
+
+    def is_training_active(self):
+        self._active_calls += 1
+        return self._active_calls <= self._active_polls
+
+
+class _FakeRequest:
+    headers = {}
+
+
+class _ReconnectRequest:
+    # Reconnect carrying the last step the client already received.
+    headers = {"last-event-id": "10"}
+
+
+def _raw(response):
+    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(), 15))
+
+
+@pytest.fixture
+def _fast_short_timeout(monkeypatch):
+    """Make the poll loop instant and the stall timeout tiny."""
+
+    async def _no_sleep(*_a, **_k):
+        return None
+
+    monkeypatch.setattr(rt.asyncio, "sleep", _no_sleep)
+    monkeypatch.setattr(rt, "_PROGRESS_STALL_TIMEOUT_POLLS", 3)
+
+
+def test_prep_phase_does_not_time_out_before_first_step(monkeypatch, _fast_short_timeout):
+    # Step 0 for many polls (far past the timeout), then the run ends. Pre-step
+    # this is preparation, not a stall: no error event may be emitted.
+    backend = _Backend(active_polls = 20, step_history = [], live_step = 0)
+    monkeypatch.setattr(rt, "get_training_backend", lambda: backend)
+
+    raw = _raw(asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")))
+
+    assert (
+        backend._active_calls > rt._PROGRESS_STALL_TIMEOUT_POLLS + 1
+    ), "the loop must have run past the stall threshold for this test to be meaningful"
+    assert "event: heartbeat" in raw, "prep heartbeats should still flow"
+    assert "event: error" not in raw, "a still-preparing run must not be timed out as a stall"
+
+
+def test_stall_after_first_step_still_times_out(monkeypatch, _fast_short_timeout):
+    # Emits a live step (so seen_live_step becomes True) then stays put: a genuine
+    # post-step stall that must still trigger the timeout error.
+    backend = _Backend(active_polls = 100, step_history = [1, 2], live_step = 5)
+    monkeypatch.setattr(rt, "get_training_backend", lambda: backend)
+
+    raw = _raw(asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")))
+
+    assert "event: error" in raw, "a real post-step stall should still time out"
+
+
+def test_reconnect_to_stepped_run_still_times_out(monkeypatch, _fast_short_timeout):
+    # Client reconnects at step 10 (Last-Event-ID) to a run that already stepped
+    # then hangs (only heartbeats): the post-step stall timeout must still fire.
+    # Without seeding seen_live_step from the resume point it resets to False and
+    # never times out for this client.
+    backend = _Backend(active_polls = 100, step_history = [10], live_step = 10)
+    monkeypatch.setattr(rt, "get_training_backend", lambda: backend)
+
+    raw = _raw(
+        asyncio.run(rt.stream_training_progress(_ReconnectRequest(), current_subject = "tester"))
+    )
+
+    assert (
+        "event: error" in raw
+    ), "a reconnect to an already-stepped run that then stalls must still time out"