ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #72 from unslothai/fix/sse-progress-timeout

Browse source 
Fix: SSE progress stream timeout during training
This commit is contained in:
Roland Tannous 2026-02-14 09:58:14 +04:00 • committed by GitHub
commit 44d52b4103
2 changed files with 15 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

2
studio/backend/core/training/training.py
View file

@ -35,7 +35,7 @@ class TrainingBackend:
def _on_progress_update(self, progress: TrainingProgress):
"""Callback for progress updates"""
if progress.step > 0 and progress.loss > 0:
if progress.step >= 0 and progress.loss > 0:
self.loss_history.append(progress.loss)
self.lr_history.append(progress.learning_rate)
self.step_history.append(progress.step)

16
studio/backend/routes/training.py
View file

@ -542,7 +542,7 @@ 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 = 300 # Timeout after 5 minutes
max_no_updates = 1800 # Timeout after 30 minutes (large models need time for compilation)
while backend.is_training_active():
try:
@ -596,7 +596,19 @@ async def stream_training_progress(
# 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)
# Pull total_steps and status from trainer so
# the frontend can show "Tokenizing…" etc.
tp_prep = getattr(
getattr(backend, "trainer", None),
"training_progress", None,
)
prep_total = (
getattr(tp_prep, "total_steps", 0)
if tp_prep else 0
)
preparing_payload = build_progress(
0, 0.0, 0.0, prep_total,
)
yield format_sse(
preparing_payload.model_dump_json(),
event="heartbeat",