diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 8da0dc508..38a2cab38 100644 --- a/studio/backend/routes/training.py +++ b/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 diff --git a/studio/backend/tests/test_training_progress_prep_timeout.py b/studio/backend/tests/test_training_progress_prep_timeout.py new file mode 100644 index 000000000..a7e6d4f83 --- /dev/null +++ b/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"