125 lines
3.7 KiB
Python
125 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
import signal
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
from typing import Any
|
|
|
|
from yolo_webui.config import MlflowConfig, TrainingConfig
|
|
from yolo_webui.trainer import TrainingEvent, TrainingRunner
|
|
|
|
|
|
class FakeProcess:
|
|
def __init__(self) -> None:
|
|
self.signals: list[int] = []
|
|
self.terminate_calls = 0
|
|
self.kill_calls = 0
|
|
|
|
def send_signal(self, signum: int) -> None:
|
|
self.signals.append(signum)
|
|
|
|
def terminate(self) -> None:
|
|
self.terminate_calls += 1
|
|
|
|
def kill(self) -> None:
|
|
self.kill_calls += 1
|
|
|
|
def poll(self) -> None:
|
|
return None
|
|
|
|
|
|
def test_runner_wires_yolo_callbacks_and_returns_output(
|
|
monkeypatch: Any, tmp_path: Path
|
|
) -> None:
|
|
settings_updates: list[dict[str, bool]] = []
|
|
constructed: list[tuple[str, str]] = []
|
|
train_arguments: list[dict[str, Any]] = []
|
|
|
|
class FakeSettings:
|
|
def update(self, values: dict[str, bool]) -> None:
|
|
settings_updates.append(values)
|
|
|
|
class FakeYOLO:
|
|
def __init__(self, model: str, task: str) -> None:
|
|
constructed.append((model, task))
|
|
self.callbacks: dict[str, Any] = {}
|
|
self.trainer = SimpleNamespace(
|
|
args=SimpleNamespace(epochs=2),
|
|
epoch=0,
|
|
metrics={},
|
|
stop=False,
|
|
save_dir=tmp_path / "run",
|
|
)
|
|
|
|
def add_callback(self, name: str, callback: Any) -> None:
|
|
self.callbacks[name] = callback
|
|
|
|
def train(self, **kwargs: Any) -> None:
|
|
train_arguments.append(kwargs)
|
|
self.callbacks["on_train_start"](self.trainer)
|
|
for epoch in range(2):
|
|
self.trainer.epoch = epoch
|
|
self.trainer.metrics = {"metrics/mAP50": 0.5 + epoch / 10}
|
|
self.callbacks["on_train_epoch_end"](self.trainer)
|
|
self.callbacks["on_train_end"](self.trainer)
|
|
|
|
fake_ultralytics = ModuleType("ultralytics")
|
|
fake_ultralytics.YOLO = FakeYOLO # type: ignore[attr-defined]
|
|
fake_ultralytics.settings = FakeSettings() # type: ignore[attr-defined]
|
|
monkeypatch.setitem(sys.modules, "ultralytics", fake_ultralytics)
|
|
|
|
config = TrainingConfig(
|
|
dataset="dataset.yaml",
|
|
model="model.pt",
|
|
task="pose",
|
|
epochs=2,
|
|
mlflow=MlflowConfig(enabled=False),
|
|
)
|
|
events: list[TrainingEvent] = []
|
|
|
|
output = TrainingRunner().train(config, events.append)
|
|
|
|
assert output == tmp_path / "run"
|
|
assert constructed == [("models/model.pt", "pose")]
|
|
assert settings_updates == [{"mlflow": False}]
|
|
assert train_arguments[0]["data"] == "dataset.yaml"
|
|
assert train_arguments[0]["verbose"] is True
|
|
assert [event.kind for event in events] == ["info", "started", "epoch", "epoch", "success"]
|
|
|
|
|
|
def test_prepare_run_clears_previous_stop_request() -> None:
|
|
runner = TrainingRunner()
|
|
runner.request_stop()
|
|
assert runner.stop_requested is True
|
|
|
|
runner.prepare_run()
|
|
|
|
assert runner.stop_requested is False
|
|
|
|
|
|
def test_early_stop_is_delivered_after_subprocess_ready() -> None:
|
|
runner = TrainingRunner()
|
|
process = FakeProcess()
|
|
runner.request_stop()
|
|
|
|
runner.set_subprocess(process, ready=False)
|
|
assert process.signals == []
|
|
|
|
runner.mark_subprocess_ready()
|
|
|
|
assert process.signals == [signal.SIGTERM]
|
|
assert process.terminate_calls == 0
|
|
runner.clear_subprocess()
|
|
|
|
|
|
def test_ready_subprocess_receives_cooperative_signal_not_terminate() -> None:
|
|
runner = TrainingRunner()
|
|
process = FakeProcess()
|
|
runner.set_subprocess(process, ready=True)
|
|
|
|
runner.request_stop()
|
|
|
|
assert process.signals == [signal.SIGTERM]
|
|
assert process.terminate_calls == 0
|
|
runner.clear_subprocess()
|