276 lines
10 KiB
Python
Executable file
276 lines
10 KiB
Python
Executable file
from __future__ import annotations
|
|
|
|
import os
|
|
import signal
|
|
from collections.abc import Callable, Iterator
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from threading import Event, RLock, Timer
|
|
from typing import Any
|
|
|
|
from .config import MlflowConfig, TrainingConfig
|
|
|
|
# Restrict PyTorch checkpoint deserialization to Ultralytics' known model classes.
|
|
os.environ["ULTRALYTICS_SAFE_LOAD"] = "1"
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class TrainingEvent:
|
|
kind: str
|
|
message: str
|
|
epoch: int = 0
|
|
total_epochs: int = 0
|
|
|
|
|
|
EventHandler = Callable[[TrainingEvent], None]
|
|
|
|
|
|
@contextmanager
|
|
def mlflow_environment(config: MlflowConfig) -> Iterator[None]:
|
|
"""Temporarily expose the settings expected by Ultralytics' MLflow callback."""
|
|
keys = {
|
|
"MLFLOW_TRACKING_URI": config.tracking_uri.strip(),
|
|
"MLFLOW_EXPERIMENT_NAME": config.experiment_name.strip(),
|
|
"MLFLOW_RUN": config.run_name.strip(),
|
|
"MLFLOW_KEEP_RUN_ACTIVE": "False",
|
|
}
|
|
previous = {key: os.environ.get(key) for key in keys}
|
|
try:
|
|
if config.enabled:
|
|
for key, value in keys.items():
|
|
if value:
|
|
os.environ[key] = value
|
|
else:
|
|
os.environ.pop(key, None)
|
|
yield
|
|
finally:
|
|
for key, value in previous.items():
|
|
if value is None:
|
|
os.environ.pop(key, None)
|
|
else:
|
|
os.environ[key] = value
|
|
|
|
|
|
class TrainingRunner:
|
|
"""Owns a single YOLO training run and exposes cooperative cancellation."""
|
|
|
|
FORCE_STOP_TIMEOUT_SECONDS = 30.0
|
|
|
|
def __init__(self) -> None:
|
|
self._model: Any | None = None
|
|
self._state_lock = RLock()
|
|
self._stop_requested = Event()
|
|
self._force_stop_triggered = Event()
|
|
self._subprocess: Any | None = None
|
|
self._subprocess_ready = False
|
|
self._force_stop_timer: Timer | None = None
|
|
|
|
def prepare_run(self) -> None:
|
|
"""Reset cancellation state before starting a new training run."""
|
|
with self._state_lock:
|
|
timer = self._force_stop_timer
|
|
self._force_stop_timer = None
|
|
self._subprocess_ready = False
|
|
self._stop_requested.clear()
|
|
self._force_stop_triggered.clear()
|
|
if timer is not None:
|
|
timer.cancel()
|
|
|
|
def set_subprocess(self, process: Any, *, ready: bool = False) -> None:
|
|
"""Register the child process without losing an earlier stop request."""
|
|
with self._state_lock:
|
|
self._subprocess = process
|
|
self._subprocess_ready = ready
|
|
stop_requested = self._stop_requested.is_set()
|
|
if stop_requested:
|
|
if ready:
|
|
self._send_cooperative_stop(process)
|
|
self._schedule_force_stop(process)
|
|
|
|
def mark_subprocess_ready(self) -> None:
|
|
"""Mark the child signal handler as ready and deliver any pending stop."""
|
|
with self._state_lock:
|
|
process = self._subprocess
|
|
self._subprocess_ready = process is not None
|
|
stop_requested = self._stop_requested.is_set()
|
|
if process is not None and stop_requested:
|
|
self._send_cooperative_stop(process)
|
|
self._schedule_force_stop(process)
|
|
|
|
def clear_subprocess(self) -> None:
|
|
with self._state_lock:
|
|
self._subprocess = None
|
|
self._subprocess_ready = False
|
|
timer = self._force_stop_timer
|
|
self._force_stop_timer = None
|
|
if timer is not None:
|
|
timer.cancel()
|
|
|
|
def request_stop(self) -> None:
|
|
self._stop_requested.set()
|
|
with self._state_lock:
|
|
process = self._subprocess
|
|
process_ready = self._subprocess_ready
|
|
trainer = getattr(self._model, "trainer", None)
|
|
if process is not None:
|
|
if process_ready:
|
|
self._send_cooperative_stop(process)
|
|
self._schedule_force_stop(process)
|
|
if trainer is not None:
|
|
trainer.stop = True
|
|
|
|
@property
|
|
def stop_requested(self) -> bool:
|
|
return self._stop_requested.is_set()
|
|
|
|
@property
|
|
def force_stop_triggered(self) -> bool:
|
|
return self._force_stop_triggered.is_set()
|
|
|
|
@staticmethod
|
|
def _send_cooperative_stop(process: Any) -> None:
|
|
try:
|
|
process.send_signal(signal.SIGTERM)
|
|
except (AttributeError, OSError, ProcessLookupError):
|
|
try:
|
|
process.terminate()
|
|
except (AttributeError, OSError, ProcessLookupError):
|
|
pass
|
|
|
|
def _schedule_force_stop(self, process: Any) -> None:
|
|
with self._state_lock:
|
|
if self._subprocess is not process or self._force_stop_timer is not None:
|
|
return
|
|
timer = Timer(
|
|
self.FORCE_STOP_TIMEOUT_SECONDS,
|
|
self._force_stop,
|
|
args=(process,),
|
|
)
|
|
timer.daemon = True
|
|
self._force_stop_timer = timer
|
|
timer.start()
|
|
|
|
def _force_stop(self, process: Any) -> None:
|
|
with self._state_lock:
|
|
if self._subprocess is not process:
|
|
return
|
|
try:
|
|
if process.poll() is None:
|
|
process.kill()
|
|
self._force_stop_triggered.set()
|
|
except (AttributeError, OSError, ProcessLookupError):
|
|
pass
|
|
|
|
def train(self, config: TrainingConfig, on_event: EventHandler) -> Path | None:
|
|
config.validate()
|
|
|
|
train_args = config.train_kwargs()
|
|
|
|
if config.split.enabled:
|
|
on_event(TrainingEvent("info", "Разделение датасета на train/val…"))
|
|
try:
|
|
from .dataset_splitter import split_dataset
|
|
train_count, val_count, yaml_path = split_dataset(
|
|
dataset_dir=config.dataset,
|
|
train_ratio=config.split.train_ratio,
|
|
classes_path=config.split.classes_path,
|
|
)
|
|
on_event(
|
|
TrainingEvent(
|
|
"info",
|
|
f"Разделение завершено: train={train_count}, val={val_count}",
|
|
)
|
|
)
|
|
train_args["data"] = yaml_path
|
|
except Exception as exc:
|
|
on_event(TrainingEvent("warning", f"Ошибка разделения датасета: {exc}"))
|
|
raise exc
|
|
|
|
on_event(TrainingEvent("info", "Загрузка Ultralytics и подготовка модели…"))
|
|
from ultralytics import YOLO, settings
|
|
from .ultralytics_trainers import trainer_for_task
|
|
|
|
workspace_root = Path.cwd().resolve()
|
|
try:
|
|
settings.update({
|
|
"runs_dir": str((workspace_root / "runs").resolve()),
|
|
"datasets_dir": str((workspace_root / "datasets").resolve()),
|
|
"weights_dir": str((workspace_root / "models").resolve()),
|
|
})
|
|
except Exception:
|
|
pass
|
|
|
|
previous_mlflow_setting = settings["mlflow"]
|
|
setting_changed = previous_mlflow_setting is not config.mlflow.enabled
|
|
if setting_changed:
|
|
settings.update({"mlflow": config.mlflow.enabled})
|
|
try:
|
|
with mlflow_environment(config.mlflow):
|
|
model = YOLO(config.resolved_model, task=config.task)
|
|
with self._state_lock:
|
|
self._model = model
|
|
|
|
model.add_callback("on_train_start", self._on_train_start(on_event))
|
|
model.add_callback("on_fit_epoch_end", self._on_epoch_end(on_event))
|
|
model.add_callback("on_train_end", self._on_train_end(on_event))
|
|
|
|
try:
|
|
model.train(trainer=trainer_for_task(config.task), **train_args)
|
|
trainer = getattr(model, "trainer", None)
|
|
save_dir = getattr(trainer, "save_dir", None)
|
|
return Path(save_dir) if save_dir else None
|
|
finally:
|
|
with self._state_lock:
|
|
self._model = None
|
|
finally:
|
|
if setting_changed:
|
|
settings.update({"mlflow": previous_mlflow_setting})
|
|
|
|
def _on_train_start(self, on_event: EventHandler) -> Callable[[Any], None]:
|
|
def callback(trainer: Any) -> None:
|
|
total = int(getattr(getattr(trainer, "args", None), "epochs", 0))
|
|
on_event(TrainingEvent("started", "Обучение началось.", 0, total))
|
|
if self._stop_requested.is_set():
|
|
trainer.stop = True
|
|
|
|
return callback
|
|
|
|
def _on_epoch_end(self, on_event: EventHandler) -> Callable[[Any], None]:
|
|
def callback(trainer: Any) -> None:
|
|
validator = getattr(trainer, "validator", None)
|
|
if validator is not None and getattr(validator, "training", True) is False:
|
|
return # final best-checkpoint validation is not an additional train epoch
|
|
epoch = int(getattr(trainer, "epoch", 0)) + 1
|
|
total = int(getattr(getattr(trainer, "args", None), "epochs", 0))
|
|
metrics = getattr(trainer, "metrics", {}) or {}
|
|
summary = self._metrics_summary(metrics)
|
|
message = f"Эпоха {epoch}/{total} завершена"
|
|
if summary:
|
|
message += f" · {summary}"
|
|
on_event(TrainingEvent("epoch", message, epoch, total))
|
|
if self._stop_requested.is_set():
|
|
trainer.stop = True
|
|
|
|
return callback
|
|
|
|
def _on_train_end(self, on_event: EventHandler) -> Callable[[Any], None]:
|
|
def callback(trainer: Any) -> None:
|
|
if self._stop_requested.is_set():
|
|
on_event(TrainingEvent("cancelled", "Обучение остановлено пользователем."))
|
|
else:
|
|
on_event(TrainingEvent("success", "Ultralytics завершил обучение."))
|
|
|
|
return callback
|
|
|
|
@staticmethod
|
|
def _metrics_summary(metrics: dict[str, Any]) -> str:
|
|
result: list[str] = []
|
|
for key, value in metrics.items():
|
|
try:
|
|
result.append(f"{key.split('/')[-1]}={float(value):.4g}")
|
|
if len(result) == 3:
|
|
break
|
|
except (TypeError, ValueError):
|
|
continue
|
|
return " · ".join(result)
|