78 lines
2.6 KiB
Python
78 lines
2.6 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Heartbeat throttle for repeated progress logs."""
|
|
|
|
import pytest
|
|
|
|
from loggers.progress import ProgressThrottle
|
|
|
|
|
|
@pytest.fixture(autouse = True)
|
|
def _not_verbose(monkeypatch):
|
|
monkeypatch.delenv("UNSLOTH_STUDIO_VERBOSE", raising = False)
|
|
monkeypatch.setenv("LOG_LEVEL", "INFO")
|
|
|
|
|
|
def test_first_message_logs_then_identical_repeats_throttle():
|
|
t = ProgressThrottle(interval_s = 100) # large window: only the first logs
|
|
k = "load"
|
|
assert t.should_log(k, "Downloading...") is True # first
|
|
assert t.should_log(k, "Downloading...") is False # identical repeat within window
|
|
assert t.should_log(k, "Downloading...") is False
|
|
|
|
|
|
def test_message_change_always_logs():
|
|
t = ProgressThrottle(interval_s = 100)
|
|
k = "load"
|
|
assert t.should_log(k, "Loading model...") is True
|
|
assert t.should_log(k, "Loading model...") is False
|
|
assert t.should_log(k, "Importing Unsloth...") is True # phase change logs
|
|
|
|
|
|
def test_heartbeat_emits_after_interval():
|
|
import time
|
|
|
|
t = ProgressThrottle(interval_s = 0.05)
|
|
k = "load"
|
|
assert t.should_log(k, "same") is True
|
|
assert t.should_log(k, "same") is False
|
|
time.sleep(0.06)
|
|
assert t.should_log(k, "same") is True # interval elapsed -> heartbeat
|
|
|
|
|
|
def test_empty_message_is_pure_time_heartbeat():
|
|
# A step counter changes every tick, so callers pass an empty message and rely
|
|
# only on the interval. Identical empty messages must not all log.
|
|
t = ProgressThrottle(interval_s = 100)
|
|
k = ("training", "job1")
|
|
assert t.should_log(k) is True
|
|
assert t.should_log(k) is False
|
|
assert t.should_log(k) is False
|
|
|
|
|
|
def test_reset_logs_next_immediately():
|
|
t = ProgressThrottle(interval_s = 100)
|
|
k = "load"
|
|
assert t.should_log(k, "x") is True
|
|
assert t.should_log(k, "x") is False
|
|
t.reset(k)
|
|
assert t.should_log(k, "x") is True # fresh after reset
|
|
|
|
|
|
def test_zero_interval_logs_everything():
|
|
t = ProgressThrottle(interval_s = 0)
|
|
assert all(t.should_log("k", "same") for _ in range(5))
|
|
|
|
|
|
def test_verbose_disables_throttle(monkeypatch):
|
|
monkeypatch.setenv("UNSLOTH_STUDIO_VERBOSE", "1")
|
|
t = ProgressThrottle(interval_s = 100)
|
|
assert all(t.should_log("k", "same") for _ in range(5))
|
|
|
|
|
|
def test_distinct_keys_are_independent():
|
|
t = ProgressThrottle(interval_s = 100)
|
|
assert t.should_log("a", "m") is True
|
|
assert t.should_log("b", "m") is True # different key, independent
|
|
assert t.should_log("a", "m") is False
|