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)