Studio: evict training progress-throttle key on error/stop, not just complete
This commit is contained in:
parent
75e0d2573e
commit
e7df5a3850
2 changed files with 59 additions and 0 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
Loading…
Add table
Add a link
Reference in a new issue