训练与实验跟踪

这一课你将能做到

1. 把训练代码组织成「加载数据 → 按时间切分 → 构造防泄漏特征 → 训练 → 评估」的结构化管道
2. 用 MLflow 自动记录每次训练的超参、指标、模型文件，并在 UI 里对比实验
3. 给训练脚本加 --smoke 模式，为第 7 课的 CI 训练冒烟测试做准备

为什么要学

训练本身不难，难的是三个月后回答：「线上这个模型是哪天、用哪份数据、什么参数训出来的？F1 多少？为什么不用另一次实验？」没有实验跟踪的团队靠聊天记录和 Excel 找答案，这是 MLOps 里最典型的混乱源。

MLflow Tracking 是事实标准的轻量方案：每次训练成为一个 run，参数、指标、模型 artifact 自动入库，一个网页看全所有实验。人话：把「调参日记」从脑子和小本本里搬进数据库，还能一键翻旧账。

先建立图像

Experiment: timeout-prediction（一个业务问题一个实验）
├─ Run #1  params: {max_iter=200, lr=0.1}   metrics: {f1=0.87, auc=0.93}
│          artifacts: model/ (模型文件+签名)   tags: {git_sha=abc123, data=judge_clean@md5}
├─ Run #2  params: {max_iter=400, lr=0.05}  metrics: {f1=0.89, auc=0.94}
└─ Run #3  ...

一次 run = 一份完整档案：代码版本 + 数据版本 + 参数 + 指标 + 模型文件

跟着做

步骤 1：安装并初始化 MLflow

pip install mlflow

步骤 2：写出结构化训练脚本

数据沿用第 3 课产出的 data/processed/judge_clean.parquet。重点看两处纪律：按时间切分（不能随机切，否则用未来数据预测过去）和防泄漏历史特征（shift(1) 把当前行排除）。

# src/train.py
import argparse, subprocess
import pandas as pd
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.metrics import f1_score, roc_auc_score
import mlflow
from mlflow.models import infer_signature

FEATURES = ["code_lines", "lang_java", "rank", "problem_tle_hist"]

def load_data() -> pd.DataFrame:
    df = pd.read_parquet("data/processed/judge_clean.parquet")
    df = df.sort_values("submit_time").reset_index(drop=True)
    df["is_tle"] = (df["verdict"] == "TLE").astype(int)
    df["lang_java"] = (df["lang"] == "Java").astype(int)
    # 防泄漏：每题历史 TLE 率，shift(1) 排除当前提交自己
    df["problem_tle_hist"] = (df.groupby("problem_id")["is_tle"]
                                .transform(lambda s: s.shift(1).expanding().mean())
                                .fillna(0.15))
    return df

def git_sha() -> str:
    try:
        return subprocess.check_output(
            ["git", "rev-parse", "--short", "HEAD"], text=True).strip()
    except Exception:
        return "unknown"

def main(smoke: bool, max_iter: int, learning_rate: float):
    df = load_data()
    if smoke:
        df = df.sample(500, random_state=42)

    split = int(len(df) * 0.8)                    # 前 80% 训练，后 20% 测试
    X_tr, y_tr = df[FEATURES][:split], df["is_tle"][:split]
    X_te, y_te = df[FEATURES][split:], df["is_tle"][split:]

    mlflow.set_experiment("timeout-prediction")
    with mlflow.start_run() as run:
        mlflow.log_params({
            "max_iter": max_iter, "learning_rate": learning_rate,
            "train_rows": len(X_tr), "test_rows": len(X_te),
        })
        mlflow.set_tags({"git_sha": git_sha(), "data": "judge_clean.parquet"})

        model = HistGradientBoostingClassifier(
            max_iter=max_iter, learning_rate=learning_rate, random_state=42)
        model.fit(X_tr, y_tr)

        proba = model.predict_proba(X_te)[:, 1]
        mlflow.log_metrics({
            "f1": f1_score(y_te, proba > 0.5),
            "roc_auc": roc_auc_score(y_te, proba),
        })
        mlflow.sklearn.log_model(
            model, name="model",
            signature=infer_signature(X_tr.head(5)),
            input_example=X_tr.head(5))
        print("run_id =", run.info.run_id)

if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--smoke", action="store_true", help="小样本冒烟训练，供 CI 使用")
    p.add_argument("--max-iter", type=int, default=200)
    p.add_argument("--lr", type=float, default=0.1)
    a = p.parse_args()
    main(a.smoke, a.max_iter, a.lr)