Studio: evict training progress-throttle key on error/stop, not just complete

This commit is contained in:
Daniel Han 2026-06-27 12:33:28 +00:00
commit e7df5a3850
2 changed files with 59 additions and 0 deletions

View file

@ -1054,6 +1054,10 @@ class TrainingBackend:
elif etype == "error":
self._progress.is_training = False
self._progress.error = event.get("error", "Unknown error")
# Evict the throttle key on error/stop too (not just on complete),
# so a re-run reusing the job_id logs its first heartbeat at once
# and the entry doesn't linger for the process lifetime.
progress_throttle.reset(("training", self.current_job_id))
logger.error("Training error: %s", event.get("error"))
stack = event.get("stack", "")
if stack:

View file

@ -0,0 +1,55 @@
# 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 training progress-log throttle key must be evicted on terminal events.
A long run logs a throttled "Training progress" heartbeat keyed by
("training", job_id). That key has to be released on BOTH complete and error/stop
(not just complete), so a re-run reusing the job_id logs its first heartbeat
immediately and the entry doesn't linger for the process lifetime.
"""
import os
import sys
import pytest
_backend = os.path.join(os.path.dirname(__file__), "..")
if _backend not in sys.path:
sys.path.insert(0, _backend)
from core.training.training import TrainingBackend
from loggers.progress import progress_throttle
@pytest.fixture(autouse = True)
def _not_verbose(monkeypatch):
# Verbose makes should_log() always True, which would mask the throttle.
monkeypatch.delenv("UNSLOTH_STUDIO_VERBOSE", raising = False)
monkeypatch.setenv("LOG_LEVEL", "INFO")
def _seed_throttled(job_id: str):
# First call logs, the immediate repeat is throttled within the 10s window.
key = ("training", job_id)
progress_throttle.reset(key)
assert progress_throttle.should_log(key) is True
assert progress_throttle.should_log(key) is False
return key
def test_error_event_resets_progress_throttle():
b = TrainingBackend()
b.current_job_id = "job-err"
key = _seed_throttled("job-err")
b._handle_event({"type": "error", "error": "boom"})
# Evicted -> the next run's first heartbeat logs at once.
assert progress_throttle.should_log(key) is True
def test_complete_event_resets_progress_throttle():
b = TrainingBackend()
b.current_job_id = "job-done"
key = _seed_throttled("job-done")
b._handle_event({"type": "complete"})
assert progress_throttle.should_log(key) is True