机器学习实战：预测批处理任务超时风险

自建任务数据
   │
   ├── 训练集 ──► 预处理 Pipeline ──► 模型训练/交叉验证
   ├── 验证集 ──► 模型比较与阈值选择
   └── 测试集 ──► 一次最终评估
                              │
                              ▼
                    pipeline.joblib + metadata.json
                              │
                              ▼
                         独立推理脚本

timeout-risk/
├── data/
├── model_output/
├── prepare_data.py
├── train.py
├── predict.py
└── requirements.txt

pandas
numpy
scikit-learn
joblib

# prepare_data.py
from pathlib import Path

import numpy as np
import pandas as pd

rng = np.random.default_rng(42)
count = 1200
task_types = np.array(["text", "table", "archive", "image"])
data = pd.DataFrame({
    "file_count": rng.integers(1, 100, count),
    "input_mb": rng.gamma(shape=2.2, scale=22.0, size=count),
    "queue_depth": rng.integers(0, 20, count),
    "worker_load": rng.uniform(0.15, 1.0, count),
    "task_type": rng.choice(task_types, count, p=[0.42, 0.25, 0.18, 0.15]),
})
type_risk = data["task_type"].map(
    {"text": -0.5, "table": 0.0, "archive": 0.5, "image": 0.8}
)
logit = (
    -5.4 + 0.018 * data["file_count"] + 0.025 * data["input_mb"]
    + 0.13 * data["queue_depth"] + 2.0 * data["worker_load"] + type_risk
)
probability = 1 / (1 + np.exp(-logit))
data["timeout"] = rng.binomial(1, probability)

# 注入少量缺失，验证流水线处理能力
data.loc[rng.choice(data.index, 20, replace=False), "input_mb"] = np.nan
output = Path("data/batch_timeout.csv")
output.parent.mkdir(parents=True, exist_ok=True)
data.to_csv(output, index=False)
print(data.shape)
print(data["timeout"].value_counts(normalize=True).sort_index())

# train.py
import json
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
import sklearn
from sklearn.compose import ColumnTransformer
from sklearn.dummy import DummyClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    average_precision_score, classification_report, f1_score,
)
from sklearn.model_selection import cross_val_score, train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

NUMERIC = ["file_count", "input_mb", "queue_depth", "worker_load"]
CATEGORICAL = ["task_type"]
FEATURES = NUMERIC + CATEGORICAL
TARGET = "timeout"


def make_preprocess() -> ColumnTransformer:
    return ColumnTransformer([
        ("number", Pipeline([
            ("impute", SimpleImputer(strategy="median")),
            ("scale", StandardScaler()),
        ]), NUMERIC),
        ("category", OneHotEncoder(handle_unknown="ignore"), CATEGORICAL),
    ])


def make_pipeline(estimator) -> Pipeline:
    return Pipeline([("preprocess", make_preprocess()), ("model", estimator)])


data = pd.read_csv("data/batch_timeout.csv")
if set(FEATURES + [TARGET]) - set(data.columns):
    raise ValueError("训练数据缺少必需字段")
X, y = data[FEATURES], data[TARGET]
X_train, X_temp, y_train, y_temp = train_test_split(
    X, y, test_size=0.4, stratify=y, random_state=42,
)
X_valid, X_test, y_valid, y_test = train_test_split(
    X_temp, y_temp, test_size=0.5, stratify=y_temp, random_state=42,
)

baseline = make_pipeline(DummyClassifier(strategy="most_frequent"))
baseline.fit(X_train, y_train)
print("baseline F1:", f1_score(y_valid, baseline.predict(X_valid)))

candidates = {
    "logistic": make_pipeline(LogisticRegression(max_iter=1000, random_state=42)),
    "forest": make_pipeline(RandomForestClassifier(
        n_estimators=180, max_depth=8, min_samples_leaf=4,
        random_state=42, n_jobs=-1,
    )),
}
best_name = ""
best_model = None
best_ap = -1.0
for name, model in candidates.items():
    cv_scores = cross_val_score(model, X_train, y_train, cv=5, scoring="average_precision")
    model.fit(X_train, y_train)
    valid_probability = model.predict_proba(X_valid)[:, 1]
    valid_ap = average_precision_score(y_valid, valid_probability)
    print(name, "cv AP", round(cv_scores.mean(), 3), "valid AP", round(valid_ap, 3))
    if valid_ap > best_ap:
        best_name, best_model, best_ap = name, model, valid_ap

valid_probability = best_model.predict_proba(X_valid)[:, 1]
thresholds = np.arange(0.20, 0.81, 0.05)
best_threshold = max(
    thresholds,
    key=lambda value: f1_score(y_valid, valid_probability >= value),
)

test_probability = best_model.predict_proba(X_test)[:, 1]
test_prediction = (test_probability >= best_threshold).astype(int)
print("selected:", best_name, "threshold:", round(float(best_threshold), 2))
print("test AP:", round(average_precision_score(y_test, test_probability), 3))
print(classification_report(y_test, test_prediction, digits=3))

errors = X_test.loc[test_prediction != y_test.to_numpy()].copy()
errors["actual"] = y_test.loc[errors.index]
errors["probability"] = test_probability[test_prediction != y_test.to_numpy()]
print("error samples:\n", errors.head().to_string(index=False))

output = Path("model_output")
output.mkdir(exist_ok=True)
joblib.dump({"model": best_model, "threshold": float(best_threshold)}, output / "pipeline.joblib")
(output / "metadata.json").write_text(json.dumps({
    "model": best_name,
    "features": FEATURES,
    "target": TARGET,
    "sklearn_version": sklearn.__version__,
    "selection_metric": "validation average_precision",
}, ensure_ascii=False, indent=2), encoding="utf-8")

from pathlib import Path

import matplotlib.pyplot as plt
from sklearn.metrics import ConfusionMatrixDisplay, PrecisionRecallDisplay

figure_dir = Path("model_output/figures")
figure_dir.mkdir(parents=True, exist_ok=True)

figure, axis = plt.subplots(figsize=(6, 4))
PrecisionRecallDisplay.from_predictions(y_test, test_probability, ax=axis)
axis.set_title("Test precision-recall curve")
figure.tight_layout()
figure.savefig(figure_dir / "precision_recall.png", dpi=150)
plt.close(figure)

figure, axis = plt.subplots(figsize=(5, 4))
ConfusionMatrixDisplay.from_predictions(y_test, test_prediction, ax=axis)
axis.set_title(f"Test confusion matrix, threshold={best_threshold:.2f}")
figure.tight_layout()
figure.savefig(figure_dir / "confusion_matrix.png", dpi=150)
plt.close(figure)

# predict.py
import joblib
import pandas as pd

artifact = joblib.load("model_output/pipeline.joblib")
sample = pd.DataFrame([{
    "file_count": 55,
    "input_mb": 72.0,
    "queue_depth": 12,
    "worker_load": 0.86,
    "task_type": "archive",
}])
probability = float(artifact["model"].predict_proba(sample)[0, 1])
prediction = int(probability >= artifact["threshold"])
print({"timeout_probability": round(probability, 4), "timeout": prediction})

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt
python prepare_data.py
python train.py
python predict.py

import joblib
import pandas as pd


def test_generated_data_contract() -> None:
    data = pd.read_csv("data/batch_timeout.csv")
    assert len(data) == 1200
    assert set(data["timeout"].unique()) == {0, 1}
    assert (data[["file_count", "queue_depth", "worker_load"]] >= 0).all().all()


def test_saved_pipeline_accepts_raw_fields() -> None:
    artifact = joblib.load("model_output/pipeline.joblib")
    sample = pd.DataFrame([{
        "file_count": 10,
        "input_mb": 20.0,
        "queue_depth": 2,
        "worker_load": 0.3,
        "task_type": "text",
    }])
    probability = artifact["model"].predict_proba(sample)[0, 1]
    assert 0.0 <= probability <= 1.0
    assert 0.0 < artifact["threshold"] < 1.0