my-homelab-configs/apps/mlops-platform/training/export_model.py

53 lines
1.8 KiB
Python

import json
from pathlib import Path
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
FEATURES = [
"latency_ms",
"error_rate",
"cpu_utilization",
"memory_utilization",
"queue_depth",
]
OUTPUT_DIR = Path(__file__).resolve().parent.parent / "models"
def export_model(version: str, seed: int, threshold: float) -> None:
rng = np.random.default_rng(seed)
healthy = rng.normal([170, 0.015, 0.38, 0.45, 5], [45, 0.01, 0.12, 0.12, 4], size=(160, 5))
at_risk = rng.normal([420, 0.12, 0.82, 0.78, 38], [130, 0.08, 0.12, 0.13, 18], size=(160, 5))
x = np.vstack([healthy, at_risk])
y = np.array([0] * len(healthy) + [1] * len(at_risk))
pipeline = make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000, random_state=seed))
pipeline.fit(x, y)
scaler = pipeline.named_steps["standardscaler"]
classifier = pipeline.named_steps["logisticregression"]
artifact = {
"version": version,
"trained_with": "scikit-learn LogisticRegression",
"threshold": threshold,
"bias": float(classifier.intercept_[0]),
"features": FEATURES,
"weights": [float(value) for value in classifier.coef_[0]],
"baseline": {
feature: {"mean": float(mean), "stddev": float(stddev)}
for feature, mean, stddev in zip(FEATURES, scaler.mean_, scaler.scale_, strict=True)
},
}
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
with (OUTPUT_DIR / f"model_{version}.json").open("w", encoding="utf-8") as handle:
json.dump(artifact, handle, indent=2)
handle.write("\n")
if __name__ == "__main__":
export_model("v1", seed=42, threshold=0.55)
export_model("v2", seed=84, threshold=0.5)