From e7df5a38502867c86b68e30c4f00d3e5670332f3 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 27 Jun 2026 12:33:28 +0000
Subject: [PATCH] Studio: evict training progress-throttle key on error/stop,
 not just complete

---
 studio/backend/core/training/training.py      |  4 ++
 .../test_training_progress_throttle_reset.py  | 55 +++++++++++++++++++
 2 files changed, 59 insertions(+)
 create mode 100644 studio/backend/tests/test_training_progress_throttle_reset.py

diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py
index 84bcc2a2e8..9d43d14718 100644
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@@ -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:
diff --git a/studio/backend/tests/test_training_progress_throttle_reset.py b/studio/backend/tests/test_training_progress_throttle_reset.py
new file mode 100644
index 0000000000..d9f5bb7ad3
--- /dev/null
+++ b/studio/backend/tests/test_training_progress_throttle_reset.py
@@ -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