""" train_selector.py — Train Heuristic Selector Models (DAHS_2) Trains three classifiers (Decision Tree, Random Forest, XGBoost) to predict which of 6 heuristics achieves the best dispatching outcome for a given system state (snapshot-fork labels). NEW in DAHS_2: - Exports models/feature_ranges.json - Exports models/dt_structure.json (for frontend glass-box) - Exports models/feature_names.json Outputs: - models/selector_dt.joblib - models/selector_rf.joblib - models/selector_xgb.joblib - models/feature_ranges.json - models/dt_structure.json - models/feature_names.json - results/plots/feature_importance.png - results/plots/decision_tree.png """ from __future__ import annotations import hashlib import json import logging import time import warnings from pathlib import Path from typing import Any, Dict, List import joblib import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import pandas as pd from sklearn.ensemble import RandomForestClassifier from sklearn.metrics import classification_report from sklearn.model_selection import StratifiedKFold, cross_val_score, train_test_split from sklearn.tree import DecisionTreeClassifier, plot_tree from xgboost import XGBClassifier warnings.filterwarnings("ignore", category=UserWarning) logger = logging.getLogger(__name__) DATA_PATH = Path(__file__).parent.parent / "data" / "raw" / "selector_dataset.csv" MODELS_DIR = Path(__file__).parent.parent / "models" PLOTS_DIR = Path(__file__).parent.parent / "results" / "plots" LABEL_NAMES = ["FIFO", "Priority-EDD", "Critical-Ratio", "ATC", "WSPT", "Slack"] def _extract_dt_structure(dt: DecisionTreeClassifier, feature_names: List[str]) -> Dict[str, Any]: """Extract decision tree node structure for frontend glass-box visualization. Returns a dict with nodes list, each node having: {id, feature, threshold, left, right, class, samples, impurity} """ tree = dt.tree_ nodes = [] def _recurse(node_id: int) -> None: feature_idx = int(tree.feature[node_id]) threshold = float(tree.threshold[node_id]) left_child = int(tree.children_left[node_id]) right_child = int(tree.children_right[node_id]) values = tree.value[node_id][0] dominant = int(np.argmax(values)) samples = int(tree.n_node_samples[node_id]) impurity = float(tree.impurity[node_id]) node: Dict[str, Any] = { "id": node_id, "samples": samples, "impurity": round(impurity, 4), "class": LABEL_NAMES[dominant], "classIdx": dominant, "values": [int(v) for v in values], } if left_child != -1: # not a leaf feat_name = feature_names[feature_idx] if feature_idx < len(feature_names) else f"f{feature_idx}" node["feature"] = feat_name node["featureIdx"] = feature_idx node["threshold"] = round(threshold, 4) node["left"] = left_child node["right"] = right_child _recurse(left_child) _recurse(right_child) nodes.append(node) _recurse(0) return {"nodes": nodes, "featureNames": feature_names, "classNames": LABEL_NAMES} def train_selector_models(data_path: Path = DATA_PATH) -> dict: """Train all three selector classifiers and save artifacts. Returns ------- dict Mapping model_name -> trained sklearn-compatible model. """ MODELS_DIR.mkdir(parents=True, exist_ok=True) PLOTS_DIR.mkdir(parents=True, exist_ok=True) logger.info("Loading selector dataset from %s", data_path) df = pd.read_csv(data_path) feature_cols = [c for c in df.columns if c != "label"] X = df[feature_cols].values.astype(np.float32) # Sanitize: NaN/inf safety (training pipeline bug fix from DAHS_1) X = np.nan_to_num(X, nan=0.0, posinf=999.0, neginf=-999.0) y = df["label"].values.astype(int) logger.info("Dataset shape: X=%s, label distribution: %s", X.shape, dict(zip(*np.unique(y, return_counts=True)))) # Training-run hash binds every artifact in this run together so the # selector loader can detect a stale OOD ranges file or a feature-list # mismatch loudly rather than silently shifting baseline-vs-DAHS results. run_hash = hashlib.sha256( f"{time.time()}|{X.shape}|{','.join(feature_cols)}|{int(y.sum())}".encode() ).hexdigest()[:16] X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=0.20, random_state=42, stratify=y ) # CV seed different from train/test split seed (bug fix) cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=123) from sklearn.utils.class_weight import compute_sample_weight sample_weights_train = compute_sample_weight("balanced", y_train) models = { "dt": DecisionTreeClassifier( max_depth=10, class_weight="balanced", random_state=42, ), "rf": RandomForestClassifier( n_estimators=400, max_depth=14, class_weight="balanced", n_jobs=-1, random_state=42, ), "xgb": XGBClassifier( n_estimators=500, learning_rate=0.03, max_depth=8, num_class=len(LABEL_NAMES), n_jobs=-1, random_state=42, eval_metric="mlogloss", verbosity=0, ), } trained = {} for name, model in models.items(): logger.info("Training %s ...", name.upper()) if name == "xgb": model.fit(X_train, y_train, sample_weight=sample_weights_train) else: model.fit(X_train, y_train) # 5-fold CV accuracy cv_scores = cross_val_score(model, X_train, y_train, cv=cv, scoring="accuracy", n_jobs=-1) logger.info("[%s] CV accuracy: %.4f +/- %.4f", name.upper(), cv_scores.mean(), cv_scores.std()) print(f"[{name.upper()}] 5-Fold CV Accuracy: {cv_scores.mean():.4f} +/- {cv_scores.std():.4f}") y_pred = model.predict(X_test) print(f"\n[{name.upper()}] Classification Report (Test Set):") print(classification_report( y_test, y_pred, labels=list(range(len(LABEL_NAMES))), target_names=LABEL_NAMES, zero_division=0, )) model_path = MODELS_DIR / f"selector_{name}.joblib" # Tag the estimator with the training-run hash so loaders can verify # it matches the on-disk feature_ranges.json / feature_names.json. try: setattr(model, "_dahs_run_hash", run_hash) except Exception: pass joblib.dump(model, model_path) logger.info("Saved model -> %s", model_path) trained[name] = model # ------------------------------------------------------------------ # NEW in DAHS_2: Export interpretability artifacts # ------------------------------------------------------------------ # 1. Feature ranges (for OOD detection in BatchwiseSelector) feature_ranges = {} for i, name in enumerate(feature_cols): feature_ranges[name] = [float(X_train[:, i].min()), float(X_train[:, i].max())] feature_ranges_payload = { "_meta": { "run_hash": run_hash, "n_train": int(X_train.shape[0]), "feature_count": len(feature_cols), }, "ranges": feature_ranges, } with open(MODELS_DIR / "feature_ranges.json", "w") as f: json.dump(feature_ranges_payload, f, indent=2) logger.info("Saved feature_ranges.json -> %s", MODELS_DIR / "feature_ranges.json") # 2. Feature names with descriptions from src.features import FEATURE_DESCRIPTIONS feature_names_data = [ { "name": name, "description": FEATURE_DESCRIPTIONS.get(name, name), "category": ( "disruption" if name in ("disruption_intensity", "queue_imbalance", "job_mix_entropy", "time_pressure_ratio") else "utilization" if "utilization" in name or "bottleneck" in name else "timing" if "due" in name or "tard" in name or "sla" in name else "queue" if "queue" in name or "throughput" in name else "system" ), "index": i, } for i, name in enumerate(feature_cols) ] feature_names_payload = { "_meta": {"run_hash": run_hash}, "features": feature_names_data, } with open(MODELS_DIR / "feature_names.json", "w") as f: json.dump(feature_names_payload, f, indent=2) logger.info("Saved feature_names.json -> %s", MODELS_DIR / "feature_names.json") # 3. Decision tree structure (for frontend glass-box) dt_structure = _extract_dt_structure(trained["dt"], feature_cols) dt_structure["_meta"] = {"run_hash": run_hash} with open(MODELS_DIR / "dt_structure.json", "w") as f: json.dump(dt_structure, f, indent=2) logger.info("Saved dt_structure.json -> %s", MODELS_DIR / "dt_structure.json") # ------------------------------------------------------------------ # Feature importance plot (RF + XGB side-by-side, dark theme) # ------------------------------------------------------------------ rf_importances = trained["rf"].feature_importances_ xgb_importances = trained["xgb"].feature_importances_ fig, axes = plt.subplots(1, 2, figsize=(16, 8)) fig.patch.set_facecolor("#0f1117") for ax, importances, title, color in zip( axes, [rf_importances, xgb_importances], ["Random Forest Feature Importance", "XGBoost Feature Importance"], ["#4fc3f7", "#a5d6a7"], ): ax.set_facecolor("#1a1d27") sorted_idx = np.argsort(importances)[-15:] ax.barh( [feature_cols[i] for i in sorted_idx], importances[sorted_idx], color=color, alpha=0.85, ) ax.set_title(title, color="white", fontsize=13, pad=10) ax.set_xlabel("Importance", color="#aaaaaa") ax.tick_params(colors="#cccccc", labelsize=9) for spine in ax.spines.values(): spine.set_color("#333344") spine.set_linewidth(0.5) fig.suptitle("Heuristic Selector — Feature Importances (DAHS_2)", color="white", fontsize=15, y=1.01) plt.tight_layout() fi_path = PLOTS_DIR / "feature_importance.png" plt.savefig(fi_path, dpi=150, bbox_inches="tight", facecolor=fig.get_facecolor()) plt.close() logger.info("Saved feature importance plot -> %s", fi_path) # ------------------------------------------------------------------ # Decision tree visualization # ------------------------------------------------------------------ fig, ax = plt.subplots(figsize=(24, 10)) fig.patch.set_facecolor("#0f1117") ax.set_facecolor("#0f1117") plot_tree( trained["dt"], feature_names=feature_cols, class_names=LABEL_NAMES, filled=True, max_depth=4, fontsize=7, ax=ax, ) ax.set_title("Decision Tree Classifier (depth≤4 shown)", color="white", fontsize=14) dt_path = PLOTS_DIR / "decision_tree.png" plt.savefig(dt_path, dpi=120, bbox_inches="tight", facecolor=fig.get_facecolor()) plt.close() logger.info("Saved decision tree plot -> %s", dt_path) return trained if __name__ == "__main__": logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") train_selector_models()