mirror of https://github.com/razor-ai/soup.git
219 lines
7.6 KiB
Python
219 lines
7.6 KiB
Python
"""Tests for loss watchdog: auto-stop on loss spikes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from soup_cli.config.schema import TrainingConfig
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Config validation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWatchdogConfig:
|
|
"""Tests for loss_watchdog config fields."""
|
|
|
|
def test_loss_watchdog_default_false(self):
|
|
"""loss_watchdog defaults to False."""
|
|
cfg = TrainingConfig()
|
|
assert cfg.loss_watchdog is False
|
|
|
|
def test_loss_watchdog_threshold_default(self):
|
|
"""loss_watchdog_threshold defaults to 3.0."""
|
|
cfg = TrainingConfig()
|
|
assert cfg.loss_watchdog_threshold == 3.0
|
|
|
|
def test_loss_watchdog_patience_default(self):
|
|
"""loss_watchdog_patience defaults to 5."""
|
|
cfg = TrainingConfig()
|
|
assert cfg.loss_watchdog_patience == 5
|
|
|
|
def test_loss_watchdog_threshold_positive(self):
|
|
"""loss_watchdog_threshold must be positive."""
|
|
with pytest.raises(Exception):
|
|
TrainingConfig(loss_watchdog_threshold=0.0)
|
|
|
|
def test_loss_watchdog_threshold_negative_rejected(self):
|
|
"""loss_watchdog_threshold rejects negative values."""
|
|
with pytest.raises(Exception):
|
|
TrainingConfig(loss_watchdog_threshold=-1.0)
|
|
|
|
def test_loss_watchdog_patience_positive(self):
|
|
"""loss_watchdog_patience must be >= 1."""
|
|
with pytest.raises(Exception):
|
|
TrainingConfig(loss_watchdog_patience=0)
|
|
|
|
def test_loss_watchdog_enable(self):
|
|
"""loss_watchdog can be enabled."""
|
|
cfg = TrainingConfig(loss_watchdog=True)
|
|
assert cfg.loss_watchdog is True
|
|
|
|
def test_loss_watchdog_custom_values(self):
|
|
"""Custom threshold and patience values are accepted."""
|
|
cfg = TrainingConfig(
|
|
loss_watchdog=True,
|
|
loss_watchdog_threshold=5.0,
|
|
loss_watchdog_patience=10,
|
|
)
|
|
assert cfg.loss_watchdog_threshold == 5.0
|
|
assert cfg.loss_watchdog_patience == 10
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Watchdog logic in callback
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWatchdogCallback:
|
|
"""Tests for loss watchdog logic in SoupTrainerCallback."""
|
|
|
|
def _make_callback(self, threshold: float = 3.0, patience: int = 5):
|
|
"""Create a callback with watchdog enabled."""
|
|
from soup_cli.monitoring.callback import SoupTrainerCallback
|
|
from soup_cli.monitoring.display import TrainingDisplay
|
|
|
|
display = MagicMock(spec=TrainingDisplay)
|
|
cb = SoupTrainerCallback(
|
|
display=display,
|
|
loss_watchdog=True,
|
|
loss_watchdog_threshold=threshold,
|
|
loss_watchdog_patience=patience,
|
|
)
|
|
return cb
|
|
|
|
def test_watchdog_no_stop_normal_loss(self):
|
|
"""Normal loss values do not trigger watchdog."""
|
|
cb = self._make_callback(threshold=3.0, patience=3)
|
|
|
|
state = MagicMock()
|
|
state.global_step = 1
|
|
state.epoch = 1.0
|
|
args = MagicMock()
|
|
control = MagicMock()
|
|
control.should_training_stop = False
|
|
|
|
# Simulate 5 normal loss values
|
|
for step in range(5):
|
|
state.global_step = step + 1
|
|
cb.on_log(args, state, control, logs={"loss": 1.5})
|
|
|
|
assert control.should_training_stop is False
|
|
|
|
def test_watchdog_triggers_after_patience(self):
|
|
"""Watchdog stops training after patience consecutive high losses."""
|
|
cb = self._make_callback(threshold=3.0, patience=3)
|
|
|
|
state = MagicMock()
|
|
args = MagicMock()
|
|
control = MagicMock()
|
|
control.should_training_stop = False
|
|
|
|
# Simulate patience+1 high loss values
|
|
for step in range(4):
|
|
state.global_step = step + 1
|
|
state.epoch = 1.0
|
|
cb.on_log(args, state, control, logs={"loss": 5.0})
|
|
|
|
assert control.should_training_stop is True
|
|
|
|
def test_watchdog_resets_on_good_loss(self):
|
|
"""Counter resets when a normal loss appears."""
|
|
cb = self._make_callback(threshold=3.0, patience=5)
|
|
|
|
state = MagicMock()
|
|
args = MagicMock()
|
|
control = MagicMock()
|
|
control.should_training_stop = False
|
|
|
|
# 3 high, then 1 normal, then 3 high — should NOT trigger (patience=5)
|
|
for loss_val in [5.0, 5.0, 5.0, 1.0, 5.0, 5.0, 5.0]:
|
|
state.global_step = 1
|
|
state.epoch = 1.0
|
|
cb.on_log(args, state, control, logs={"loss": loss_val})
|
|
|
|
assert control.should_training_stop is False
|
|
|
|
def test_watchdog_disabled_by_default(self):
|
|
"""Watchdog does not interfere when disabled."""
|
|
from soup_cli.monitoring.callback import SoupTrainerCallback
|
|
from soup_cli.monitoring.display import TrainingDisplay
|
|
|
|
display = MagicMock(spec=TrainingDisplay)
|
|
cb = SoupTrainerCallback(display=display)
|
|
|
|
state = MagicMock()
|
|
args = MagicMock()
|
|
control = MagicMock()
|
|
control.should_training_stop = False
|
|
|
|
# Very high loss — should NOT trigger since watchdog is off
|
|
for step in range(20):
|
|
state.global_step = step
|
|
state.epoch = 1.0
|
|
cb.on_log(args, state, control, logs={"loss": 100.0})
|
|
|
|
assert control.should_training_stop is False
|
|
|
|
def test_watchdog_exact_patience_boundary(self):
|
|
"""Exactly patience high-loss steps triggers on the next one."""
|
|
cb = self._make_callback(threshold=3.0, patience=3)
|
|
|
|
state = MagicMock()
|
|
args = MagicMock()
|
|
control = MagicMock()
|
|
control.should_training_stop = False
|
|
|
|
# Exactly 3 high losses (patience=3) — should trigger on 4th
|
|
for step in range(3):
|
|
state.global_step = step + 1
|
|
state.epoch = 1.0
|
|
cb.on_log(args, state, control, logs={"loss": 5.0})
|
|
|
|
# After patience steps of high loss, next one triggers
|
|
assert control.should_training_stop is True
|
|
|
|
def test_watchdog_no_loss_in_logs(self):
|
|
"""Watchdog handles logs without loss key gracefully."""
|
|
cb = self._make_callback(threshold=3.0, patience=3)
|
|
|
|
state = MagicMock()
|
|
args = MagicMock()
|
|
control = MagicMock()
|
|
control.should_training_stop = False
|
|
|
|
# Logs without loss key
|
|
cb.on_log(args, state, control, logs={"learning_rate": 1e-5})
|
|
assert control.should_training_stop is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Sweep integration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWatchdogSweep:
|
|
"""Tests for watchdog fields in sweep param support."""
|
|
|
|
def test_watchdog_threshold_in_sweep(self):
|
|
"""loss_watchdog_threshold is a valid sweep param."""
|
|
from soup_cli.commands.sweep import _parse_sweep_params
|
|
|
|
params = _parse_sweep_params(
|
|
["training.loss_watchdog_threshold=2.0,3.0,5.0"]
|
|
)
|
|
assert "training.loss_watchdog_threshold" in params
|
|
assert len(params["training.loss_watchdog_threshold"]) == 3
|
|
|
|
def test_watchdog_patience_in_sweep(self):
|
|
"""loss_watchdog_patience is a valid sweep param."""
|
|
from soup_cli.commands.sweep import _parse_sweep_params
|
|
|
|
params = _parse_sweep_params(
|
|
["training.loss_watchdog_patience=3,5,10"]
|
|
)
|
|
assert "training.loss_watchdog_patience" in params
|
|
assert len(params["training.loss_watchdog_patience"]) == 3
|