diff --git a/PROJECTS/advanced/ai-threat-detection/backend/ml/orchestrator.py b/PROJECTS/advanced/ai-threat-detection/backend/ml/orchestrator.py new file mode 100644 index 00000000..3b423a80 --- /dev/null +++ b/PROJECTS/advanced/ai-threat-detection/backend/ml/orchestrator.py @@ -0,0 +1,253 @@ +""" +©AngelaMos | 2026 +orchestrator.py +""" + +import json +import logging +from dataclasses import dataclass +from pathlib import Path + +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: + logger.exception("Ensemble validation failed") + ensemble = None + passed = False + + for artifact in self._output_dir.iterdir(): + experiment.log_artifact(artifact) + + 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: + """ + 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: + """ + 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: + """ + 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, + rf_result: dict, + if_result: dict, + ) -> 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 + ) diff --git a/PROJECTS/advanced/ai-threat-detection/backend/ml/splitting.py b/PROJECTS/advanced/ai-threat-detection/backend/ml/splitting.py index 89c91a91..d42e21c4 100644 --- a/PROJECTS/advanced/ai-threat-detection/backend/ml/splitting.py +++ b/PROJECTS/advanced/ai-threat-detection/backend/ml/splitting.py @@ -64,8 +64,13 @@ def prepare_training_data( class_counts = np.bincount(y_train) minority_count = class_counts.min() + majority_count = class_counts.max() + current_ratio = minority_count / majority_count - if minority_count >= smote_k + 1: + if ( + minority_count >= smote_k + 1 + and current_ratio < smote_strategy + ): sampler = SMOTE( sampling_strategy=smote_strategy, k_neighbors=smote_k, diff --git a/PROJECTS/advanced/ai-threat-detection/backend/tests/test_orchestrator.py b/PROJECTS/advanced/ai-threat-detection/backend/tests/test_orchestrator.py new file mode 100644 index 00000000..8d2c76d0 --- /dev/null +++ b/PROJECTS/advanced/ai-threat-detection/backend/tests/test_orchestrator.py @@ -0,0 +1,180 @@ +""" +©AngelaMos | 2026 +test_orchestrator.py +""" + +import json +from pathlib import Path + +import numpy as np + +from ml.orchestrator import TrainingOrchestrator, TrainingResult + +N_FEATURES = 35 +EXPECTED_FILES = [ + "ae.onnx", + "rf.onnx", + "if.onnx", + "scaler.json", + "threshold.json", +] + + +def _make_dataset() -> tuple[np.ndarray, np.ndarray]: + """ + Generate a small synthetic dataset for testing + """ + rng = np.random.default_rng(42) + X_normal = rng.standard_normal( + (200, N_FEATURES) + ).astype(np.float32) + X_attack = ( + rng.standard_normal( + (80, N_FEATURES) + ).astype(np.float32) + + 2.0 + ) + X = np.vstack([X_normal, X_attack]) + y = np.array( + [0] * 200 + [1] * 80, dtype=np.int32 + ) + return X, y + + +class TestTrainingOrchestrator: + """ + Test the end-to-end training orchestrator + """ + + def test_produces_all_output_files( + self, tmp_path: Path + ) -> None: + """ + Orchestrator produces all 5 expected output files + """ + X, y = _make_dataset() + orch = TrainingOrchestrator( + output_dir=tmp_path, epochs=3 + ) + orch.run(X, y) + + for filename in EXPECTED_FILES: + assert ( + tmp_path / filename + ).exists(), f"Missing {filename}" + + def test_returns_training_result( + self, tmp_path: Path + ) -> None: + """ + Returns a TrainingResult dataclass + """ + X, y = _make_dataset() + orch = TrainingOrchestrator( + output_dir=tmp_path, epochs=3 + ) + result = orch.run(X, y) + + assert isinstance(result, TrainingResult) + + def test_scaler_json_has_required_keys( + self, tmp_path: Path + ) -> None: + """ + scaler.json contains center, scale, and n_features + """ + X, y = _make_dataset() + orch = TrainingOrchestrator( + output_dir=tmp_path, epochs=3 + ) + orch.run(X, y) + + scaler_data = json.loads( + (tmp_path / "scaler.json").read_text() + ) + assert "center" in scaler_data + assert "scale" in scaler_data + assert "n_features" in scaler_data + + def test_threshold_json_has_float( + self, tmp_path: Path + ) -> None: + """ + threshold.json contains a float threshold value + """ + X, y = _make_dataset() + orch = TrainingOrchestrator( + output_dir=tmp_path, epochs=3 + ) + orch.run(X, y) + + threshold_data = json.loads( + (tmp_path / "threshold.json").read_text() + ) + assert "threshold" in threshold_data + assert isinstance( + threshold_data["threshold"], float + ) + + def test_result_has_per_model_metrics( + self, tmp_path: Path + ) -> None: + """ + TrainingResult includes metrics for each model + """ + X, y = _make_dataset() + orch = TrainingOrchestrator( + output_dir=tmp_path, epochs=3 + ) + result = orch.run(X, y) + + assert "ae_threshold" in result.ae_metrics + assert "f1" in result.rf_metrics + assert "n_samples" in result.if_metrics + + def test_ensemble_metrics_present( + self, tmp_path: Path + ) -> None: + """ + Ensemble validation metrics are populated + """ + X, y = _make_dataset() + orch = TrainingOrchestrator( + output_dir=tmp_path, epochs=3 + ) + result = orch.run(X, y) + + assert result.ensemble_metrics is not None + assert 0.0 <= result.ensemble_metrics.f1 <= 1.0 + + def test_mlflow_run_id_set( + self, tmp_path: Path + ) -> None: + """ + MLflow run ID is captured in the result + """ + X, y = _make_dataset() + mlflow_dir = tmp_path / "mlruns" + mlflow_dir.mkdir() + orch = TrainingOrchestrator( + output_dir=tmp_path / "models", + epochs=3, + ) + result = orch.run(X, y) + + assert result.mlflow_run_id is not None + assert len(result.mlflow_run_id) == 32 + + def test_passed_gates_is_bool( + self, tmp_path: Path + ) -> None: + """ + passed_gates is a boolean value + """ + X, y = _make_dataset() + orch = TrainingOrchestrator( + output_dir=tmp_path, epochs=3 + ) + result = orch.run(X, y) + + assert isinstance(result.passed_gates, bool)