feat(ml): add training orchestrator with MLflow and validation
This commit is contained in:
CarterPerez-dev
parent
24eefc5fb4
commit
a26d68e8d1
|
|
@ -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
|
||||
)
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
Loading…
Reference in New Issue