mirror of https://github.com/razor-ai/soup.git
101 lines
2.9 KiB
Python
101 lines
2.9 KiB
Python
"""Tests for SoupTrainerCallback."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from soup_cli.monitoring.callback import SoupTrainerCallback
|
|
|
|
|
|
def _make_state(global_step=10, max_steps=100, epoch=1.0):
|
|
"""Create a mock TrainerState."""
|
|
state = MagicMock()
|
|
state.global_step = global_step
|
|
state.max_steps = max_steps
|
|
state.epoch = epoch
|
|
return state
|
|
|
|
|
|
def _make_args():
|
|
"""Create a mock TrainingArguments."""
|
|
return MagicMock()
|
|
|
|
|
|
def test_on_train_begin_starts_display():
|
|
"""on_train_begin should call display.start with total_steps."""
|
|
display = MagicMock()
|
|
callback = SoupTrainerCallback(display=display)
|
|
state = _make_state(max_steps=500)
|
|
|
|
callback.on_train_begin(_make_args(), state, MagicMock())
|
|
|
|
display.start.assert_called_once_with(total_steps=500)
|
|
|
|
|
|
def test_on_train_end_stops_display():
|
|
"""on_train_end should call display.stop."""
|
|
display = MagicMock()
|
|
callback = SoupTrainerCallback(display=display)
|
|
|
|
callback.on_train_end(_make_args(), _make_state(), MagicMock())
|
|
|
|
display.stop.assert_called_once()
|
|
|
|
|
|
def test_on_log_updates_display():
|
|
"""on_log should call display.update with metrics from logs."""
|
|
display = MagicMock()
|
|
callback = SoupTrainerCallback(display=display)
|
|
|
|
logs = {
|
|
"loss": 1.234,
|
|
"learning_rate": 2e-5,
|
|
"grad_norm": 0.5,
|
|
"train_steps_per_second": 3.0,
|
|
}
|
|
state = _make_state(global_step=42, epoch=1.5)
|
|
|
|
with patch("soup_cli.monitoring.callback.torch", create=True):
|
|
callback.on_log(_make_args(), state, MagicMock(), logs=logs)
|
|
|
|
display.update.assert_called_once()
|
|
call_kwargs = display.update.call_args
|
|
assert call_kwargs[1]["step"] == 42 or call_kwargs[0][0] == 42
|
|
|
|
|
|
def test_on_log_none_logs():
|
|
"""on_log with logs=None should do nothing."""
|
|
display = MagicMock()
|
|
callback = SoupTrainerCallback(display=display)
|
|
|
|
callback.on_log(_make_args(), _make_state(), MagicMock(), logs=None)
|
|
|
|
display.update.assert_not_called()
|
|
|
|
|
|
def test_on_log_with_tracker():
|
|
"""on_log should forward metrics to tracker if provided."""
|
|
display = MagicMock()
|
|
tracker = MagicMock()
|
|
callback = SoupTrainerCallback(display=display, tracker=tracker, run_id="run_123")
|
|
|
|
logs = {"loss": 0.5, "learning_rate": 1e-5}
|
|
state = _make_state(global_step=10, epoch=1.0)
|
|
|
|
callback.on_log(_make_args(), state, MagicMock(), logs=logs)
|
|
|
|
tracker.log_metrics.assert_called_once()
|
|
call_kwargs = tracker.log_metrics.call_args[1]
|
|
assert call_kwargs["run_id"] == "run_123"
|
|
assert call_kwargs["step"] == 10
|
|
assert call_kwargs["loss"] == 0.5
|
|
|
|
|
|
def test_on_log_without_tracker():
|
|
"""on_log without tracker should not crash."""
|
|
display = MagicMock()
|
|
callback = SoupTrainerCallback(display=display, tracker=None, run_id="")
|
|
|
|
logs = {"loss": 0.5}
|
|
callback.on_log(_make_args(), _make_state(), MagicMock(), logs=logs)
|
|
|
|
display.update.assert_called_once()
|