Cybersecurity-Projects/PROJECTS/advanced/ai-threat-detection/backend/ml/orchestrator.py

"""
©AngelaMos | 2026
orchestrator.py
"""

import json
import logging
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

from ml.experiment import VigilExperiment
from ml.export_onnx import (
    export_autoencoder,
    export_isolation_forest,
    export_random_forest,
)
from ml.splitting import prepare_training_data
from ml.train_autoencoder import train_autoencoder
from ml.train_classifiers import (
    train_isolation_forest,
    train_random_forest,
)
from ml.validation import ValidationResult, validate_ensemble

logger = logging.getLogger(__name__)

N_FEATURES = 35

AE_FILENAME = "ae.onnx"
RF_FILENAME = "rf.onnx"
IF_FILENAME = "if.onnx"
SCALER_FILENAME = "scaler.json"
THRESHOLD_FILENAME = "threshold.json"


@dataclass
class TrainingResult:
    """
    Aggregated results from a full training pipeline run
    """

    ae_metrics: dict[str, float]
    rf_metrics: dict[str, float]
    if_metrics: dict[str, float]
    ensemble_metrics: ValidationResult | None
    passed_gates: bool
    output_dir: Path
    mlflow_run_id: str | None


class TrainingOrchestrator:
    """
    End-to-end training pipeline that splits data, trains all 3 models,
    exports to ONNX, validates the ensemble, and logs to MLflow
    """

    def __init__(
        self,
        output_dir: Path,
        experiment_name: str = "angelusvigil-training",
        epochs: int = 100,
        batch_size: int = 256,
    ) -> None:
        self._output_dir = output_dir
        self._experiment_name = experiment_name
        self._epochs = epochs
        self._batch_size = batch_size

    def run(
        self,
        X: np.ndarray,
        y: np.ndarray,
    ) -> TrainingResult:
        """
        Execute the full training pipeline
        """
        self._output_dir.mkdir(parents=True, exist_ok=True)

        split = prepare_training_data(X, y)

        logger.info(
            "Split: train=%d val=%d test=%d normal_train=%d",
            len(split.X_train),
            len(split.X_val),
            len(split.X_test),
            len(split.X_normal_train),
        )

        with VigilExperiment(self._experiment_name) as experiment:
            experiment.log_params({
                "epochs": self._epochs,
                "batch_size": self._batch_size,
                "n_samples": len(X),
                "n_attack": int(np.sum(y == 1)),
                "n_normal": int(np.sum(y == 0)),
                "n_features": X.shape[1],
            })

            ae_result = self._train_ae(split.X_normal_train)
            ae_metrics = {
                "ae_threshold": ae_result["threshold"],
                "ae_final_train_loss": ae_result["history"]["train_loss"][-1],
                "ae_final_val_loss": ae_result["history"]["val_loss"][-1],
            }

            rf_result = self._train_rf(split.X_train, split.y_train)
            rf_metrics = rf_result["metrics"]

            if_result = self._train_if(split.X_normal_train)
            if_metrics = if_result["metrics"]

            self._export_models(ae_result, rf_result, if_result)

            experiment.log_metrics(ae_metrics)
            experiment.log_metrics({
                f"rf_{k}": v
                for k, v in rf_metrics.items()
            })

            try:
                ensemble = validate_ensemble(
                    self._output_dir,
                    split.X_test,
                    split.y_test,
                )
                experiment.log_metrics({
                    "ensemble_precision": ensemble.precision,
                    "ensemble_recall": ensemble.recall,
                    "ensemble_f1": ensemble.f1,
                    "ensemble_pr_auc": ensemble.pr_auc,
                    "ensemble_roc_auc": ensemble.roc_auc,
                })
                passed = ensemble.passed_gates
            except Exception as exc:
                logger.exception("Ensemble validation failed")
                print(
                    f"  WARNING: validation raised"
                    f" {type(exc).__name__}: {exc}",
                    file=sys.stderr,
                )
                ensemble = None
                passed = False

            for name in (
                AE_FILENAME,
                RF_FILENAME,
                IF_FILENAME,
                SCALER_FILENAME,
                THRESHOLD_FILENAME,
            ):
                experiment.log_artifact(self._output_dir / name)

            run_id = experiment.run_id

        logger.info(
            "Training complete: passed_gates=%s run_id=%s",
            passed,
            run_id,
        )

        return TrainingResult(
            ae_metrics=ae_metrics,
            rf_metrics=rf_metrics,
            if_metrics=if_metrics,
            ensemble_metrics=ensemble,
            passed_gates=passed,
            output_dir=self._output_dir,
            mlflow_run_id=run_id,
        )

    def _train_ae(self, X_normal: np.ndarray) -> dict[str, Any]:
        """
        Train the autoencoder on normal-only data
        """
        logger.info(
            "Training autoencoder (%d epochs, %d samples)",
            self._epochs,
            len(X_normal),
        )
        return train_autoencoder(
            X_normal,
            epochs=self._epochs,
            batch_size=self._batch_size,
        )

    def _train_rf(self, X: np.ndarray, y: np.ndarray) -> dict[str, Any]:
        """
        Train the random forest classifier
        """
        logger.info(
            "Training random forest (%d samples)",
            len(X),
        )
        return train_random_forest(X, y)

    def _train_if(self, X_normal: np.ndarray) -> dict[str, Any]:
        """
        Train the isolation forest on normal-only data
        """
        logger.info(
            "Training isolation forest (%d samples)",
            len(X_normal),
        )
        return train_isolation_forest(X_normal)

    def _export_models(
        self,
        ae_result: dict[str, Any],
        rf_result: dict[str, Any],
        if_result: dict[str, Any],
    ) -> None:
        """
        Export all 3 models to ONNX and save scaler/threshold
        """
        export_autoencoder(
            ae_result["model"],
            self._output_dir / AE_FILENAME,
        )
        ae_result["scaler"].save_json(self._output_dir / SCALER_FILENAME)
        threshold_data = {"threshold": ae_result["threshold"]}
        (self._output_dir / THRESHOLD_FILENAME).write_text(
            json.dumps(threshold_data, indent=2))

        export_random_forest(
            rf_result["model"],
            N_FEATURES,
            self._output_dir / RF_FILENAME,
        )

        export_isolation_forest(
            if_result["model"],
            N_FEATURES,
            self._output_dir / IF_FILENAME,
        )

        logger.info("Exported models to %s", self._output_dir)