Trixy7/trixy_core/trainer/service.py

# -*- coding: utf-8 -*-
"""
TrainerService - IService Wrapper fuer das Trainer-Framework.

Verwaltet die Trainer-Registry und fuehrt Trainings als
isolierte Subprocesse aus.
"""

from __future__ import annotations

import multiprocessing
import asyncio
import dataclasses
from typing import TYPE_CHECKING

from trixy_core.service import IService, ServicePriority, ServiceGroup
from trixy_core.trainer.registry import TrainerRegistry
from trixy_core.trainer.enums import TrainingState
from trixy_core.trainer.dataclasses import (
    TrainerInfo,
    TrainingProgress,
    ValidationResult,
    FormSchema,
)
from trixy_core.utils.debug import pinfo, pdebug, perror

if TYPE_CHECKING:
    from trixy_core.application import IApplication
    from trixy_core.trainer.itrainer import ITrainer


def _load_nvidia_libs() -> None:
    """Laedt CUDA-Bibliotheken aus nvidia Pip-Packages (fuer onnxruntime-gpu im Subprocess)."""
    import os
    import sys
    import ctypes
    from pathlib import Path

    # Suche nvidia libs im venv
    for site_dir in sys.path:
        nvidia_base = Path(site_dir) / "nvidia"
        if nvidia_base.is_dir():
            lib_dirs = [str(p) for p in nvidia_base.glob("*/lib") if p.is_dir()]
            if lib_dirs:
                existing = os.environ.get("LD_LIBRARY_PATH", "")
                os.environ["LD_LIBRARY_PATH"] = ":".join(lib_dirs) + (":" + existing if existing else "")
                for lib_dir in lib_dirs:
                    for so in sorted(Path(lib_dir).glob("*.so.*")):
                        try:
                            ctypes.cdll.LoadLibrary(str(so))
                        except OSError:
                            pass
            break


def _training_worker(
    trainer: "ITrainer",
    progress_queue: multiprocessing.Queue,
    pause_flag: multiprocessing.Event,
    stop_flag: multiprocessing.Event,
) -> None:
    """Einstiegspunkt fuer den Trainings-Subprocess."""
    _load_nvidia_libs()

    try:
        trainer.train(progress_queue, pause_flag, stop_flag)
    except Exception as e:
        progress_queue.put(TrainingProgress(
            state=TrainingState.FAILED,
            error=str(e),
        ))


class TrainerService(IService):
    """
    Verwaltet Trainer-Discovery und Training-Ausfuehrung.

    Trainings laufen als multiprocessing.Process, Progress
    wird ueber eine Queue gepollt.
    """

    PRIORITY = ServicePriority.OPTIONAL
    GROUP = ServiceGroup.UTILITY
    DEPENDENCIES: list[str] = []

    def __init__(self, application: "IApplication") -> None:
        super().__init__(application)
        self._registry = TrainerRegistry()

        # Aktives Training
        self._active_trainer_id: str = ""
        self._process: multiprocessing.Process | None = None
        self._progress_queue: multiprocessing.Queue | None = None
        self._pause_flag: multiprocessing.Event | None = None
        self._stop_flag: multiprocessing.Event | None = None
        self._current_progress = TrainingProgress()
        self._poll_task: asyncio.Task | None = None

    @property
    def registry(self) -> TrainerRegistry:
        """Trainer-Registry."""
        return self._registry

    @property
    def active_trainer_id(self) -> str:
        """ID des aktuell trainierenden Trainers."""
        return self._active_trainer_id

    @property
    def current_progress(self) -> TrainingProgress:
        """Aktueller Trainings-Fortschritt."""
        return self._current_progress

    async def start(self) -> None:
        """Startet den Service und entdeckt Trainer."""
        core_count = self._registry.discover_core()
        plugin_count = self._registry.discover_plugins()
        pinfo(f"TrainerService: {core_count} Core-Trainer, {plugin_count} Plugin-Trainer entdeckt")

    async def stop(self) -> None:
        """Stoppt den Service und laufende Trainings."""
        await self.stop_training()
        if self._poll_task:
            self._poll_task.cancel()
            self._poll_task = None

    def list_trainers(self) -> list[TrainerInfo]:
        """Liefert alle verfuegbaren Trainer."""
        return self._registry.list_info()

    def get_trainer(self, trainer_id: str) -> "ITrainer | None":
        """Liefert einen Trainer nach ID."""
        return self._registry.get(trainer_id)

    def get_trainer_detail(self, trainer_id: str) -> dict:
        """Liefert detaillierte Trainer-Informationen als dict."""
        trainer = self._registry.get(trainer_id)
        if not trainer:
            return {"error": f"Trainer '{trainer_id}' nicht gefunden"}

        info = trainer.get_info()
        settings_schema = trainer.get_settings_schema()
        dataset_schema = trainer.get_dataset_schema()
        optional_schema = trainer.get_optional_schema()

        return {
            "trainer_id": info.trainer_id,
            "name": info.name,
            "description": info.description,
            "version": info.version,
            "author": info.author,
            "model_type": info.model_type,
            "trainer_type": info.trainer_type.value,
            "settings_schema": dataclasses.asdict(settings_schema),
            "dataset_schema": dataclasses.asdict(dataset_schema),
            "optional_schema": dataclasses.asdict(optional_schema) if optional_schema else None,
            "settings": trainer.get_settings(),
            "dataset_info": trainer.get_dataset_info(),
        }

    def write_settings(self, trainer_id: str, settings: dict) -> tuple[bool, str]:
        """Setzt Einstellungen eines Trainers."""
        trainer = self._registry.get(trainer_id)
        if not trainer:
            return False, f"Trainer '{trainer_id}' nicht gefunden"

        try:
            trainer.set_settings(settings)
            return True, ""
        except Exception as e:
            return False, str(e)

    def validate_trainer(self, trainer_id: str) -> ValidationResult:
        """Validiert einen Trainer."""
        trainer = self._registry.get(trainer_id)
        if not trainer:
            return ValidationResult(errors=[f"Trainer '{trainer_id}' nicht gefunden"])

        try:
            return trainer.validate()
        except Exception as e:
            return ValidationResult(errors=[f"Validierungsfehler: {e}"])

    async def start_training(self, trainer_id: str) -> tuple[bool, str]:
        """Startet ein Training als Subprocess."""
        if self._process and self._process.is_alive():
            return False, "Ein Training laeuft bereits"

        trainer = self._registry.get(trainer_id)
        if not trainer:
            return False, f"Trainer '{trainer_id}' nicht gefunden"

        # Subprocess-Kommunikation vorbereiten (spawn statt fork — CUDA erfordert 'spawn')
        ctx = multiprocessing.get_context("spawn")
        self._progress_queue = ctx.Queue()
        self._pause_flag = ctx.Event()
        self._stop_flag = ctx.Event()
        self._active_trainer_id = trainer_id
        self._current_progress = TrainingProgress(state=TrainingState.STARTING)
        self._process = ctx.Process(
            target=_training_worker,
            args=(trainer, self._progress_queue, self._pause_flag, self._stop_flag),
            daemon=True,
        )
        self._process.start()

        # Polling-Task starten
        self._poll_task = asyncio.create_task(self._poll_progress())

        # ActivityTracker: Training-Start protokollieren
        try:
            tracker = self._application.services.get_service("ActivityTracker")
            if tracker:
                model_name = trainer.get_info().name if hasattr(trainer, "get_info") else trainer_id
                tracker.log_training_start(trainer_id, model_name)
        except Exception:
            pass

        pinfo(f"Training gestartet: {trainer_id} (PID: {self._process.pid})")
        return True, ""

    async def pause_training(self) -> tuple[bool, str]:
        """Pausiert das laufende Training."""
        if not self._process or not self._process.is_alive():
            return False, "Kein aktives Training"

        if self._pause_flag:
            self._pause_flag.set()
            self._current_progress.state = TrainingState.PAUSED
        return True, ""

    async def resume_training(self) -> tuple[bool, str]:
        """Setzt ein pausiertes Training fort."""
        if not self._process or not self._process.is_alive():
            return False, "Kein aktives Training"

        if self._pause_flag:
            self._pause_flag.clear()
            self._current_progress.state = TrainingState.TRAINING
        return True, ""

    async def stop_training(self) -> tuple[bool, str]:
        """Stoppt das laufende Training."""
        if not self._process:
            return False, "Kein aktives Training"

        if self._stop_flag:
            self._stop_flag.set()

        # Warten bis Prozess endet (max 10 Sek.)
        if self._process.is_alive():
            self._process.join(timeout=10)
            if self._process.is_alive():
                self._process.terminate()
                self._process.join(timeout=5)

        self._current_progress.state = TrainingState.COMPLETED
        self._cleanup_training()
        return True, ""

    def execute_action(self, trainer_id: str, action: str, params: dict) -> dict:
        """Fuehrt eine Trainer-Aktion aus."""
        trainer = self._registry.get(trainer_id)
        if not trainer:
            return {"success": False, "message": f"Trainer '{trainer_id}' nicht gefunden"}

        try:
            return trainer.execute_action(action, params)
        except Exception as e:
            return {"success": False, "message": f"Aktionsfehler: {e}"}

    def get_progress(self) -> TrainingProgress:
        """Liefert den aktuellen Fortschritt."""
        return self._current_progress

    async def _poll_progress(self) -> None:
        """Pollt die Progress-Queue des Subprocesses."""
        while True:
            try:
                await asyncio.sleep(0.5)

                if not self._progress_queue:
                    break

                # Alle verfuegbaren Updates lesen
                while not self._progress_queue.empty():
                    try:
                        progress = self._progress_queue.get_nowait()
                        if isinstance(progress, TrainingProgress):
                            self._current_progress = progress
                    except Exception:
                        break

                # Pruefen ob Prozess noch laeuft
                if self._process and not self._process.is_alive():
                    exit_code = self._process.exitcode
                    state = self._current_progress.state
                    if state in (TrainingState.TRAINING, TrainingState.STARTING):
                        if exit_code != 0:
                            self._current_progress.state = TrainingState.FAILED
                            self._current_progress.error = (
                                self._current_progress.error
                                or f"Subprocess beendet mit Exit-Code {exit_code}"
                            )
                        else:
                            self._current_progress.state = TrainingState.COMPLETED

                    # ActivityTracker: Training-Abschluss protokollieren
                    try:
                        tracker = self._application.services.get_service("ActivityTracker")
                        if tracker:
                            is_success = self._current_progress.state == TrainingState.COMPLETED
                            tracker.log_training_complete(
                                self._active_trainer_id,
                                success=is_success,
                                error=self._current_progress.error or "",
                            )
                    except Exception:
                        pass

                    self._cleanup_training()
                    break

            except asyncio.CancelledError:
                break
            except Exception as e:
                perror(f"Progress-Polling Fehler: {e}")
                break

    def _cleanup_training(self) -> None:
        """Raeumt nach einem Training auf."""
        self._process = None
        self._progress_queue = None
        self._pause_flag = None
        self._stop_flag = None
        self._active_trainer_id = ""