作者：PySuper
日期：2025-11-01
标签：MLOps、训练流水线、Argo Workflows、MLflow
一、为什么需要训练流水线
在机器学习项目的早期阶段，很多团队的训练流程是这样的：手动下载数据 → 手动清洗 → 手动训练 → 手动评估 → 手动部署。这种方式在小规模实验时看起来没问题，但随着项目发展，问题会接踵而至：
流程不可复现：同样的代码，这次跑出来效果好，下次可能因为某个环节的微小变化导致效果下降
资源利用率低下：数据准备时 GPU 空闲，训练时人工盯着，产出检查时流程又卡住
实验管理混乱：跑了 100 个实验后，很难说清楚哪个配置对应哪个结果
错误难以追踪：训练失败时，很难快速定位是数据问题、配置问题还是代码问题
一个设计良好的训练流水线，本质上是将数据工程、模型工程和MLOps三者有机结合，形成一条可观测、可复现、自动化的端到端链路。
二、训练流水线的全链路架构
┌─────────────────────────────────────────────────────────────────────────────────┐
│ 模型训练全链路架构图 │
├─────────────────────────────────────────────────────────────────────────────────┤
│ │
│ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │
│ │ 数据源 │ │ 数据仓库 │ │ 特征存储 │ │
│ │ (OSS/S3) │────▶│ (Raw Data) │────▶│ (Feature DB) │ │
│ └──────────────┘ └──────────────┘ └──────────────┘ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ 数据准备阶段 │ │
│ │ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ │
│ │ │ 数据清洗 │ │ 数据标注 │ │ 质量检查 │ │ 数据分割 │ │ │
│ │ └────┬────┘ └────┬────┘ └────┬────┘ └────┬────┘ │ │
│ │ └───────────┴───────────┴───────────┘ │ │
│ │ │ │ │
│ │ ▼ │ │
│ │ ┌─────────────┐ │ │
│ │ │ 训练数据集 │ │ │
│ │ │ (Train/Val) │ │ │
│ │ └─────────────┘ │ │
│ └───────────────────────────┬───────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ 训练配置管理 │ │
│ │ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │ │
│ │ │ 超参数配置 │ │ 实验追踪 │ │ 模型注册 │ │ │
│ │ │ (YAML/MLflow)│ │ (MLflow) │ │ (Model Zoo)│ │ │
│ │ └─────────────┘ └─────────────┘ └─────────────┘ │ │
│ └───────────────────────────┬───────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ 训练执行阶段 │ │
│ │ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ │
│ │ │分布式训练│ │断点续训 │ │早停策略 │ │资源调度 │ │ │
│ │ └────┬────┘ └────┬────┘ └────┬────┘ └────┬────┘ │ │
│ │ └───────────┴───────────┴───────────┘ │ │
│ │ │ │ │
│ │ ▼ │ │
│ │ ┌─────────────┐ │ │
│ │ │ Checkpoint │ │ │
│ │ │ 存档管理 │ │ │
│ │ └─────────────┘ │ │
│ └───────────────────────────┬───────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ 产出检查阶段 │ │
│ │ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ │
│ │ │Loss曲线 │ │自动评测 │ │质量门禁 │ │模型对比 │ │ │
│ │ └────┬────┘ └────┬────┘ └────┬────┘ └────┬────┘ │ │
│ │ └───────────┴───────────┴───────────┘ │ │
│ │ │ │ │
│ │ ▼ │ │
│ │ ┌─────────────┐ │ │
│ │ │ 发布决策 │ │ │
│ │ │ (Proceed/Reject)│ │ │
│ │ └─────────────┘ │ │
│ └─────────────────────────────────────────────────────────────┘ │
│ │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ 流水线编排层 │ │
│ │ Argo Workflows / Apache Airflow │ │
│ └─────────────────────────────────────────────────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────────────────┘整个流水线可以分为四个核心阶段：数据准备、训练配置管理、训练执行和产出检查。每个阶段都有明确的输入输出，通过流水线编排层统一调度。
三、数据准备阶段
数据准备是训练流水线的起点，也是最容易出问题的环节。很多团队在数据清洗上花的时间比训练本身还多，这是正常的——"垃圾进，垃圾出"（Garbage In, Garbage Out） 是机器学习的铁律。
3.1 数据清洗
数据清洗的核心任务是处理缺失值、异常值、重复数据和格式不一致问题。下面是一个通用的数据清洗脚本：
#!/usr/bin/env python3
"""
数据清洗模块
功能：处理原始数据中的各种质量问题
"""
import pandas as pd
import numpy as np
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass
from pathlib import Path
import logging
import json
# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
@dataclass
class CleaningConfig:
"""数据清洗配置"""
drop_duplicates: bool = True # 是否删除重复数据
handle_missing_strategy: str = "mean" # 缺失值处理策略：mean/median/drop
outlier_threshold: float = 3.0 # 异常值 Z-score 阈值
min_sample_length: int = 10 # 最小样本长度（文本场景）
max_sample_length: int = 4096 # 最大样本长度（文本场景）
required_columns: List[str] = None # 必需列
text_columns: List[str] = None # 文本列（用于长度检查）
def __post_init__(self):
if self.required_columns is None:
self.required_columns = ["input", "target"]
if self.text_columns is None:
self.text_columns = ["input", "target"]
class DataCleaner:
"""数据清洗器"""
def __init__(self, config: CleaningConfig):
self.config = config
self.stats = {
"total_samples": 0,
"duplicates_removed": 0,
"missing_values_filled": 0,
"outliers_removed": 0,
"length_filtered": 0,
"final_samples": 0
}
def clean(self, df: pd.DataFrame) -> pd.DataFrame:
"""
执行完整的数据清洗流程
Args:
df: 原始 DataFrame
Returns:
清洗后的 DataFrame
"""
self.stats["total_samples"] = len(df)
original_len = len(df)
# 1. 检查必需列
df = self._validate_columns(df)
# 2. 删除重复数据
if self.config.drop_duplicates:
df = self._remove_duplicates(df)
# 3. 处理缺失值
df = self._handle_missing_values(df)
# 4. 处理异常值
df = self._remove_outliers(df)
# 5. 长度过滤（文本场景）
df = self._filter_by_length(df)
self.stats["final_samples"] = len(df)
# 记录清洗统计
logger.info(f"数据清洗完成: {original_len} -> {len(df)} "
f"(保留率: {len(df)/original_len*100:.2f}%)")
logger.info(f"清洗统计: {json.dumps(self.stats, indent=2)}")
return df
def _validate_columns(self, df: pd.DataFrame) -> pd.DataFrame:
"""验证必需列是否存在"""
missing_cols = set(self.config.required_columns) - set(df.columns)
if missing_cols:
raise ValueError(f"缺少必需列: {missing_cols}")
return df
def _remove_duplicates(self, df: pd.DataFrame) -> pd.DataFrame:
"""删除重复数据"""
before = len(df)
df = df.drop_duplicates()
self.stats["duplicates_removed"] = before - len(df)
if self.stats["duplicates_removed"] > 0:
logger.info(f"删除重复样本: {self.stats['duplicates_removed']}")
return df
def _handle_missing_values(self, df: pd.DataFrame) -> pd.DataFrame:
"""处理缺失值"""
for col in df.columns:
if df[col].isna().sum() > 0:
missing_count = df[col].isna().sum()
if self.config.handle_missing_strategy == "drop":
df = df.dropna(subset=[col])
elif self.config.handle_missing_strategy == "mean":
if df[col].dtype in [np.float64, np.float32, np.int64]:
df[col] = df[col].fillna(df[col].mean())
else:
df[col] = df[col].fillna("")
elif self.config.handle_missing_strategy == "median":
if df[col].dtype in [np.float64, np.float32, np.int64]:
df[col] = df[col].fillna(df[col].median())
else:
df[col] = df[col].fillna("")
self.stats["missing_values_filled"] += missing_count
logger.info(f"列 {col} 填充缺失值: {missing_count}")
return df
def _remove_outliers(self, df: pd.DataFrame) -> pd.DataFrame:
"""使用 Z-score 移除异常值（仅对数值列有效）"""
numeric_cols = df.select_dtypes(include=[np.number]).columns
for col in numeric_cols:
if col in df.columns and df[col].std() > 0:
z_scores = np.abs((df[col] - df[col].mean()) / df[col].std())
outliers_mask = z_scores > self.config.outlier_threshold
outliers_count = outliers_mask.sum()
if outliers_count > 0:
df = df[~outliers_mask]
self.stats["outliers_removed"] += outliers_count
logger.info(f"列 {col} 移除异常值: {outliers_count}")
return df
def _filter_by_length(self, df: pd.DataFrame) -> pd.DataFrame:
"""按长度过滤样本（文本场景）"""
if not self.config.text_columns:
return df
for col in self.config.text_columns:
if col in df.columns:
# 计算文本长度（按字符或词）
df[f"{col}_length"] = df[col].astype(str).apply(len)
# 应用长度过滤
before = len(df)
df = df[
(df[f"{col}_length"] >= self.config.min_sample_length) &
(df[f"{col}_length"] <= self.config.max_sample_length)
]
filtered = before - len(df)
if filtered > 0:
self.stats["length_filtered"] += filtered
logger.info(f"列 {col} 长度过滤: {filtered}")
# 删除临时长度列
df = df.drop(columns=[f"{col}_length"])
return df
def get_stats(self) -> Dict:
"""获取清洗统计信息"""
return self.stats
def main():
"""主函数：演示数据清洗流程"""
# 创建示例数据
np.random.seed(42)
n_samples = 1000
data = {
"input": [f"这是第{i}条输入文本" for i in range(n_samples)],
"target": [f"这是第{i}条目标文本" for i in range(n_samples)],
"score": np.random.randn(n_samples) * 10 + 50,
"category": np.random.choice(["A", "B", "C"], n_samples)
}
# 注入一些问题数据
df = pd.DataFrame(data)
df.loc[10:15, "score"] = np.nan # 缺失值
df.loc[20:25, "score"] = 1000 # 异常值
df = pd.concat([df, df.iloc[:5]], ignore_index=True) # 重复数据
df.loc[30:35, "input"] = "短" # 长度不足
# 执行清洗
config = CleaningConfig(
drop_duplicates=True,
handle_missing_strategy="mean",
outlier_threshold=3.0,
min_sample_length=5,
max_sample_length=100
)
cleaner = DataCleaner(config)
cleaned_df = cleaner.clean(df)
print(f"\n最终数据形状: {cleaned_df.shape}")
print(f"清洗统计: {cleaner.get_stats()}")
if __name__ == "__main__":
main()3.2 数据集分割配置
数据集分割需要考虑训练集、验证集和测试集的比例，以及如何保证分布一致性：
# 数据集分割配置
# dataset_split_config.yaml
dataset_split:
# 分割比例配置
ratios:
train: 0.8 # 训练集 80%
val: 0.1 # 验证集 10%
test: 0.1 # 测试集 10%
# 分割策略
strategy: stratified # stratified（分层）/ random（随机）/ temporal（时序）
# 分层分割配置
stratification:
enabled: true
stratify_columns:
- category # 按类别分层
- length_bucket # 按长度区间分层
# 时序分割配置
temporal:
enabled: false
time_column: created_at
train_cutoff: "2025-01-01"
val_cutoff: "2025-06-01"
# 随机种子（确保可复现）
seed: 42
# 最小样本数要求
min_samples:
train: 1000
val: 100
test: 100
# 输出路径
output:
base_dir: ./data/splits
train_file: train.parquet
val_file: val.parquet
test_file: test.parquet
stats_file: split_statistics.json#!/usr/bin/env python3
"""
数据集分割模块
功能：按配置进行训练/验证/测试集分割，并生成统计报告
"""
import pandas as pd
import numpy as np
from pathlib import Path
from typing import Dict, List, Tuple, Optional
import json
import yaml
from sklearn.model_selection import train_test_split
class DatasetSplitter:
"""数据集分割器"""
def __init__(self, config_path: str):
with open(config_path, 'r', encoding='utf-8') as f:
self.config = yaml.safe_load(f)['dataset_split']
self.seed = self.config['seed']
np.random.seed(self.seed)
def split(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
"""
执行数据集分割
Args:
df: 完整的 DataFrame
Returns:
(train_df, val_df, test_df)
"""
strategy = self.config['strategy']
if strategy == 'stratified':
return self._stratified_split(df)
elif strategy == 'temporal':
return self._temporal_split(df)
else:
return self._random_split(df)
def _stratified_split(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
"""分层分割（保持类别分布一致）"""
stratify_cols = self.config['stratification']['stratify_columns']
# 创建分层键（组合多个分类维度）
df['_stratify_key'] = df[stratify_cols].astype(str).agg('_'.join, axis=1)
# 计算分割后的最小样本数
min_samples = self.config['min_samples']
# 先分割出测试集
train_val_df, test_df = train_test_split(
df,
test_size=self.config['ratios']['test'],
stratify=df['_stratify_key'],
random_state=self.seed
)
# 再从 train_val 中分割出验证集
val_ratio = self.config['ratios']['val'] / (
self.config['ratios']['train'] + self.config['ratios']['val']
)
train_df, val_df = train_test_split(
train_val_df,
test_size=val_ratio,
stratify=train_val_df['_stratify_key'],
random_state=self.seed
)
# 删除临时列
for d in [train_df, val_df, test_df]:
d.drop('_stratify_key', axis=1, inplace=True)
# 检查最小样本数
self._validate_min_samples(train_df, val_df, test_df, min_samples)
return train_df, val_df, test_df
def _random_split(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
"""随机分割"""
min_samples = self.config['min_samples']
# 依次分割
train_df, temp_df = train_test_split(
df,
test_size=self.config['ratios']['val'] + self.config['ratios']['test'],
random_state=self.seed
)
val_ratio = self.config['ratios']['val'] / (
self.config['ratios']['val'] + self.config['ratios']['test']
)
val_df, test_df = train_test_split(
temp_df,
test_size=1 - val_ratio,
random_state=self.seed
)
self._validate_min_samples(train_df, val_df, test_df, min_samples)
return train_df, val_df, test_df
def _temporal_split(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
"""时序分割（用于时间序列数据）"""
temporal_config = self.config['temporal']
time_col = temporal_config['time_column']
# 确保时间列是 datetime 类型
df[time_col] = pd.to_datetime(df[time_col])
# 按时间分割
train_cutoff = pd.to_datetime(temporal_config['train_cutoff'])
val_cutoff = pd.to_datetime(temporal_config['val_cutoff'])
train_df = df[df[time_col] < train_cutoff]
val_df = df[(df[time_col] >= train_cutoff) & (df[time_col] < val_cutoff)]
test_df = df[df[time_col] >= val_cutoff]
return train_df, val_df, test_df
def _validate_min_samples(self, train_df: pd.DataFrame, val_df: pd.DataFrame,
test_df: pd.DataFrame, min_samples: Dict):
"""验证分割后的最小样本数"""
for name, split_df, min_count in [
('train', train_df, min_samples['train']),
('val', val_df, min_samples['val']),
('test', test_df, min_samples['test'])
]:
if len(split_df) < min_count:
raise ValueError(
f"{name} 集样本数 {len(split_df)} 少于最小要求 {min_count}"
)
def save_splits(self, train_df: pd.DataFrame, val_df: pd.DataFrame,
test_df: pd.DataFrame) -> Dict:
"""
保存分割后的数据集，并生成统计报告
Returns:
统计信息字典
"""
output_config = self.config['output']
base_dir = Path(output_config['base_dir'])
base_dir.mkdir(parents=True, exist_ok=True)
# 保存各个数据集
train_df.to_parquet(base_dir / output_config['train_file'], index=False)
val_df.to_parquet(base_dir / output_config['val_file'], index=False)
test_df.to_parquet(base_dir / output_config['test_file'], index=False)
# 生成统计报告
stats = self._generate_statistics(train_df, val_df, test_df)
with open(base_dir / output_config['stats_file'], 'w', encoding='utf-8') as f:
json.dump(stats, f, indent=2, ensure_ascii=False)
return stats
def _generate_statistics(self, train_df: pd.DataFrame,
val_df: pd.DataFrame,
test_df: pd.DataFrame) -> Dict:
"""生成数据集统计报告"""
def get_split_stats(df: pd.DataFrame, name: str) -> Dict:
stats = {
'name': name,
'num_samples': len(df),
'num_columns': len(df.columns),
'columns': list(df.columns),
'dtypes': {col: str(dtype) for col, dtype in df.dtypes.items()},
'missing_values': df.isna().sum().to_dict()
}
# 数值列统计
numeric_cols = df.select_dtypes(include=[np.number]).columns
if len(numeric_cols) > 0:
stats['numeric_summary'] = df[numeric_cols].describe().to_dict()
return stats
return {
'config': self.config,
'splits': {
'train': get_split_stats(train_df, 'train'),
'val': get_split_stats(val_df, 'val'),
'test': get_split_stats(test_df, 'test')
}
}四、训练配置管理与实验追踪
4.1 超参数配置管理
超参数配置是实验可复现性的关键。我们使用 YAML 格式管理配置，支持参数继承和环境覆盖：
# 训练配置 - base_config.yaml
# 基础配置模板
model:
name: "gpt2"
hidden_size: 768
num_layers: 12
num_attention_heads: 12
intermediate_size: 3072
dropout: 0.1
max_position_embeddings: 1024
training:
# 优化器配置
optimizer:
type: "adamw"
learning_rate: 5e-5
weight_decay: 0.01
beta1: 0.9
beta2: 0.999
epsilon: 1e-8
# 学习率调度
lr_scheduler:
type: "cosine"
warmup_steps: 500
warmup_ratio: 0.0
min_lr: 1e-5
# 训练参数
batch_size: 32
gradient_accumulation_steps: 1
max_grad_norm: 1.0
num_train_epochs: 3
logging_steps: 10
save_steps: 500
eval_steps: 500
save_total_limit: 3
# 混合精度
fp16: true
bf16: false
# 分布式训练
deepspeed:
enabled: false
config_path: "configs/deepspeed_config.json"
data:
train_file: "./data/splits/train.parquet"
val_file: "./data/splits/val.parquet"
max_seq_length: 512
preprocessing_num_workers: 4
output:
output_dir: "./outputs/experiment_{{timestamp}}"
logging_dir: "./logs/experiment_{{timestamp}}"
report_to: ["mlflow", "tensorboard"]
random_seed: 42# 实验配置覆盖示例 - experiment_lora.yaml
# 基于 base_config.yaml 的实验配置
# 继承基础配置
_base_config: "base_config.yaml"
model:
name: "gpt2"
# LoRA 配置
lora:
enabled: true
r: 8
lora_alpha: 16
lora_dropout: 0.05
target_modules:
- "q_attention"
- "v_attention"
- "output投影"
bias: "none"
task_type: "CAUSAL_LM"
training:
batch_size: 16 # 减小 batch size（LoRA 显存更省）
learning_rate: 3e-4 # LoRA 通常用更高的学习率
num_train_epochs: 5
warmup_steps: 100
output:
output_dir: "./outputs/lora_experiment"4.2 MLflow 实验追踪集成
#!/usr/bin/env python3
"""
MLflow 实验追踪模块
功能：统一管理训练实验的参数、日志、指标和产物
"""
import mlflow
import mlflow.pytorch
from mlflow.tracking import MlflowClient
from typing import Dict, Any, Optional, List
from pathlib import Path
import json
import logging
import torch
import os
logger = logging.getLogger(__name__)
class MLflowTracker:
"""MLflow 实验追踪器封装"""
def __init__(
self,
tracking_uri: str = "http://localhost:5000",
experiment_name: str = "default",
artifact_root: str = "./mlruns"
):
"""
初始化 MLflow 追踪器
Args:
tracking_uri: MLflow Server 地址
experiment_name: 实验名称
artifact_root: 产物存储根目录
"""
self.tracking_uri = tracking_uri
self.experiment_name = experiment_name
self.artifact_root = Path(artifact_root)
# 设置追踪服务器
mlflow.set_tracking_uri(tracking_uri)
# 创建或获取实验
try:
mlflow.create_experiment(experiment_name, artifact_location=str(self.artifact_root))
except mlflow.exceptions MlflowException:
pass # 实验已存在
mlflow.set_experiment(experiment_name)
logger.info(f"MLflow Tracker 初始化完成: {tracking_uri}, 实验: {experiment_name}")
def start_run(self, run_name: Optional[str] = None, tags: Optional[Dict] = None):
"""
开始一个新的实验 run
Args:
run_name: Run 名称
tags: Run 标签
"""
self.run = mlflow.start_run(
run_name=run_name,
tags=tags,
log_system_metrics=True # 记录系统指标
)
logger.info(f"开始 Run: {self.run.info.run_id}, 名称: {run_name}")
return self.run
def end_run(self, status: str = "FINISHED"):
"""结束当前 Run"""
mlflow.end_run(status=status)
logger.info(f"Run 结束: {self.run.info.run_id}, 状态: {status}")
def log_params(self, params: Dict[str, Any]):
"""
记录超参数
Args:
params: 参数字典
"""
# MLflow 有参数数量限制（100个），需要处理嵌套结构
flat_params = self._flatten_dict(params)
mlflow.log_params(flat_params)
logger.debug(f"记录参数: {len(flat_params)} 个")
def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
"""
记录指标
Args:
metrics: 指标字典
step: 当前步数
"""
mlflow.log_metrics(metrics, step=step)
logger.debug(f"记录指标 (step={step}): {metrics}")
def log_model(
self,
model: torch.nn.Module,
artifact_path: str = "model",
metrics: Optional[Dict] = None
):
"""
记录模型
Args:
model: PyTorch 模型
artifact_path: 产物路径
metrics: 模型相关指标
"""
# 记录模型
mlflow.pytorch.log_model(
model,
artifact_path,
registered_model_name=self.experiment_name
)
# 记录模型元信息
model_info = {
"num_parameters": sum(p.numel() for p in model.parameters()),
"num_trainable_params": sum(
p.numel() for p in model.parameters() if p.requires_grad
)
}
if metrics:
model_info.update(metrics)
mlflow.log_dict(model_info, artifact_file="model_info.json")
logger.info(f"模型已记录: {artifact_path}, 参数: {model_info['num_parameters']}")
def log_artifact(self, local_path: str, artifact_path: Optional[str] = None):
"""
记录产物文件
Args:
local_path: 本地文件路径
artifact_path: 产物中的路径
"""
mlflow.log_artifact(local_path, artifact_path)
logger.debug(f"记录产物: {local_path} -> {artifact_path}")
def log_figure(self, figure, artifact_path: str):
"""
记录 matplotlib/seaborn 图表
Args:
figure: 图表对象
artifact_path: 保存路径
"""
mlflow.log_figure(figure, artifact_path)
def get_best_run(self, metric: str = "val_loss", mode: str = "min") -> Optional[Dict]:
"""
获取最佳 Run
Args:
metric: 用于比较的指标
mode: 'min' 或 'max'
Returns:
最佳 Run 信息
"""
client = MlflowClient()
# 获取当前实验的所有 runs
runs = client.search_runs(
experiment_ids=[mlflow.get_experiment_by_name(self.experiment_name).experiment_id],
filter_string=f"",
order_by=[f"metrics.{metric} {'ASC' if mode == 'min' else 'DESC'}"]
)
if not runs:
return None
best_run = runs[0]
return {
"run_id": best_run.info.run_id,
"run_name": best_run.info.run_name,
"metrics": best_run.data.metrics,
"params": best_run.data.params
}
def _flatten_dict(self, d: Dict, parent_key: str = "", sep: str = ".") -> Dict:
"""将嵌套字典展平"""
items = []
for k, v in d.items():
new_key = f"{parent_key}{sep}{k}" if parent_key else k
if isinstance(v, dict):
items.extend(self._flatten_dict(v, new_key, sep=sep).items())
elif isinstance(v, list):
# 将列表转为字符串（MLflow 参数不支持列表）
items.append((new_key, json.dumps(v)))
else:
items.append((new_key, v))
return dict(items)
# ============ 训练脚本集成示例 ============
def train_with_mlflow_tracking():
"""演示如何使用 MLflowTracker 进行训练追踪"""
# 配置
config = {
"model": {"name": "gpt2", "hidden_size": 768, "num_layers": 12},
"training": {"batch_size": 32, "learning_rate": 5e-5, "num_epochs": 3}
}
# 初始化追踪器
tracker = MLflowTracker(
tracking_uri="http://localhost:5000",
experiment_name="gpt2-finetuning",
artifact_root="./mlruns"
)
# 开始实验
tracker.start_run(
run_name=f"run_{config['model']['name']}_{config['training']['batch_size']}",
tags={"task": "text-generation", "dataset": "custom"}
)
try:
# 记录配置参数
tracker.log_params(config)
# 模拟训练循环
for epoch in range(config["training"]["num_epochs"]):
train_loss = 0.95 ** epoch # 模拟 loss 下降
val_loss = train_loss * 1.1
metrics = {
"train_loss": train_loss,
"val_loss": val_loss,
"train_accuracy": 0.5 + 0.4 * (1 - train_loss),
"val_accuracy": 0.45 + 0.35 * (1 - val_loss),
"learning_rate": config["training"]["learning_rate"],
"epoch": epoch
}
tracker.log_metrics(metrics, step=epoch)
# 定期保存 checkpoint
if epoch % 2 == 0:
# 实际场景中这里会是真实的模型
checkpoint_path = f"./checkpoints/epoch_{epoch}.pt"
# torch.save(model.state_dict(), checkpoint_path)
tracker.log_artifact(checkpoint_path, "checkpoints")
# 记录最终模型
# tracker.log_model(model, metrics={"final_val_loss": val_loss})
logger.info("训练完成!")
except Exception as e:
logger.error(f"训练失败: {e}")
tracker.end_run(status="FAILED")
raise
finally:
tracker.end_run(status="FINISHED")
# 获取最佳 Run
best_run = tracker.get_best_run(metric="val_loss", mode="min")
logger.info(f"最佳 Run: {best_run['run_id']}")
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
train_with_mlflow_tracking()五、训练执行阶段
5.1 分布式训练配置
#!/usr/bin/env python3
"""
分布式训练启动器
支持多卡、多机分布式训练
"""
import os
import argparse
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from typing import Optional, List
import logging
logger = logging.getLogger(__name__)
class DistributedTrainer:
"""分布式训练器"""
def __init__(
self,
backend: str = "nccl",
init_method: str = "env://"
):
"""
初始化分布式训练环境
Args:
backend: 通信后端 ('nccl' for GPU, 'gloo' for CPU)
init_method: 初始化方法
"""
self.backend = backend
self.init_method = init_method
self.world_size = int(os.environ.get("WORLD_SIZE", 1))
self.rank = int(os.environ.get("RANK", 0))
self.local_rank = int(os.environ.get("LOCAL_RANK", 0))
self.is_distributed = self.world_size > 1
self.is_main_process = self.rank == 0
def setup(self):
"""设置分布式环境"""
if not self.is_distributed:
logger.info("单卡训练模式")
return
# 初始化进程组
dist.init_process_group(
backend=self.backend,
init_method=self.init_method,
world_size=self.world_size,
rank=self.rank
)
# 设置当前设备的 CUDA 设备
if torch.cuda.is_available():
torch.cuda.set_device(self.local_rank)
logger.info(
f"分布式训练初始化完成: rank={self.rank}/{self.world_size}, "
f"local_rank={self.local_rank}, backend={self.backend}"
)
def cleanup(self):
"""清理分布式环境"""
if self.is_distributed:
dist.destroy_process_group()
logger.info("分布式进程组已销毁")
def prepare_model(self, model: torch.nn.Module) -> torch.nn.Module:
"""
准备模型用于分布式训练
Args:
model: 原始模型
Returns:
包装后的模型
"""
if not self.is_distributed:
return model
# 同步 BatchNorm（推荐用于分布式训练）
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
# 使用 DDP 包装模型
model = DDP(
model,
device_ids=[self.local_rank] if torch.cuda.is_available() else None,
output_device=self.local_rank if torch.cuda.is_available() else None,
find_unused_parameters=True # 用于有分支结构的模型
)
logger.info(f"模型已使用 DDP 包装 (rank={self.rank})")
return model
def barrier(self):
"""进程同步屏障"""
if self.is_distributed:
dist.barrier()
def reduce_metric(self, metric: float, op=dist.ReduceOp.SUM) -> float:
"""
跨进程归约指标
Args:
metric: 本进程的指标值
op: 归约操作
Returns:
归约后的值
"""
if not self.is_distributed:
return metric
tensor = torch.tensor(metric, device=self._get_device())
dist.all_reduce(tensor, op=op)
if op == dist.ReduceOp.SUM:
return tensor.item() / self.world_size
return tensor.item()
def _get_device(self) -> torch.device:
"""获取当前设备"""
if torch.cuda.is_available():
return torch.device(f"cuda:{self.local_rank}")
return torch.device("cpu")
def launch_distributed_training(
script_path: str,
num_gpus: int = 1,
num_nodes: int = 1,
master_addr: str = "127.0.0.1",
master_port: int = 29500
):
"""
使用 torchrun 启动分布式训练
Args:
script_path: 训练脚本路径
num_gpus: 每节点 GPU 数
num_nodes: 节点数
master_addr: 主节点地址
master_port: 主节点端口
"""
cmd = [
"torchrun",
f"--nnodes={num_nodes}",
f"--nproc_per_node={num_gpus}",
f"--master_addr={master_addr}",
f"--master_port={master_port}",
script_path
]
os.system(" ".join(cmd))
if __name__ == "__main__":
# 分布式训练示例
logging.basicConfig(level=logging.INFO)
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=0)
args = parser.parse_args()
trainer = DistributedTrainer()
trainer.setup()
# 后续训练代码...
trainer.cleanup()5.2 断点续训与 Checkpoint 管理
#!/usr/bin/env python3
"""
Checkpoint 管理模块
支持自动保存、加载、断点续训
"""
import torch
import os
import shutil
import json
import logging
from pathlib import Path
from typing import Dict, Optional, Any, List
from dataclasses import dataclass, field
from datetime import datetime
import glob
logger = logging.getLogger(__name__)
@dataclass
class CheckpointConfig:
"""Checkpoint 配置"""
save_dir: str = "./checkpoints"
save_steps: int = 500
save_total_limit: int = 3 # 最多保留的 checkpoint 数量
save_best_only: bool = False # 仅保存最佳模型
save_optimizer: bool = True # 保存优化器状态
save_scheduler: bool = True # 保存学习率调度器状态
resume_from_checkpoint: Optional[str] = None
class CheckpointManager:
"""Checkpoint 管理器"""
def __init__(self, config: CheckpointConfig):
self.config = config
self.checkpoint_dir = Path(config.save_dir)
self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
self.best_metric = float('inf') if config.save_best_only else None
self.best_metric_mode = "min" # "min" 或 "max"
# 元信息文件
self.meta_file = self.checkpoint_dir / "checkpoint_meta.json"
self.load_meta()
def save_checkpoint(
self,
model: torch.nn.Module,
optimizer: Optional[torch.optim.Optimizer] = None,
scheduler: Optional[Any] = None,
step: int = 0,
metrics: Optional[Dict[str, float]] = None,
epoch: int = 0,
**kwargs
) -> Optional[str]:
"""
保存 checkpoint
Args:
model: 模型
optimizer: 优化器
scheduler: 学习率调度器
step: 当前步数
metrics: 当前指标
epoch: 当前 epoch
Returns:
保存路径，如果未保存则返回 None
"""
# 检查是否应该保存
should_save = self._should_save_checkpoint(step, metrics)
if not should_save:
return None
# 构建 checkpoint 路径
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
checkpoint_name = f"checkpoint-step-{step}-{timestamp}"
checkpoint_path = self.checkpoint_dir / checkpoint_name
checkpoint_path.mkdir(exist_ok=True)
# 收集 checkpoint 数据
checkpoint = {
"step": step,
"epoch": epoch,
"metrics": metrics or {},
"timestamp": timestamp,
"model_state_dict": model.state_dict(),
}
if optimizer and self.config.save_optimizer:
checkpoint["optimizer_state_dict"] = optimizer.state_dict()
if scheduler and self.config.save_scheduler:
checkpoint["scheduler_state_dict"] = scheduler.state_dict()
# 添加额外数据
checkpoint.update(kwargs)
# 保存
torch.save(checkpoint, checkpoint_path / "model.pt")
# 保存元信息
self._save_meta(step, epoch, checkpoint_name, metrics)
# 保存为 latest
latest_path = self.checkpoint_dir / "latest.pt"
shutil.copy(checkpoint_path / "model.pt", latest_path)
# 清理旧 checkpoint
self._cleanup_old_checkpoints()
logger.info(f"Checkpoint 已保存: {checkpoint_path}")
# 如果是最佳模型，保存一份 best.pt
if metrics and self._is_best_metric(metrics):
best_path = self.checkpoint_dir / "best.pt"
shutil.copy(checkpoint_path / "model.pt", best_path)
logger.info(f"新的最佳模型: {metrics}")
return str(checkpoint_path)
def load_checkpoint(
self,
checkpoint_path: str,
model: torch.nn.Module,
optimizer: Optional[torch.optim.Optimizer] = None,
scheduler: Optional[Any] = None,
device: str = "cuda"
) -> Dict[str, Any]:
"""
加载 checkpoint
Args:
checkpoint_path: checkpoint 路径
model: 模型（会被加载状态）
optimizer: 优化器（可选）
scheduler: 学习率调度器（可选）
device: 加载设备
Returns:
checkpoint 元信息
"""
if not os.path.exists(checkpoint_path):
raise FileNotFoundError(f"Checkpoint 不存在: {checkpoint_path}")
checkpoint_file = Path(checkpoint_path) / "model.pt"
if not checkpoint_file.exists():
checkpoint_file = Path(checkpoint_path) # 可能是直接指定 .pt 文件
logger.info(f"加载 Checkpoint: {checkpoint_file}")
checkpoint = torch.load(checkpoint_file, map_location=device)
# 加载模型
model.load_state_dict(checkpoint["model_state_dict"])
# 加载优化器
if optimizer and "optimizer_state_dict" in checkpoint:
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
# 加载调度器
if scheduler and "scheduler_state_dict" in checkpoint:
scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
meta = {
"step": checkpoint.get("step", 0),
"epoch": checkpoint.get("epoch", 0),
"metrics": checkpoint.get("metrics", {}),
}
logger.info(f"Checkpoint 加载完成: step={meta['step']}, epoch={meta['epoch']}")
return meta
def find_latest_checkpoint(self) -> Optional[str]:
"""查找最新的 checkpoint"""
latest = self.checkpoint_dir / "latest.pt"
if latest.exists():
return str(latest)
# 查找所有 checkpoint
checkpoints = sorted(
self.checkpoint_dir.glob("checkpoint-*/model.pt"),
key=lambda x: os.path.getmtime(x),
reverse=True
)
return str(checkpoints[0]) if checkpoints else None
def find_best_checkpoint(self) -> Optional[str]:
"""查找最佳的 checkpoint"""
best = self.checkpoint_dir / "best.pt"
if best.exists():
return str(best)
return None
def _should_save_checkpoint(self, step: int, metrics: Optional[Dict]) -> bool:
"""判断是否应该保存 checkpoint"""
# 如果是最佳模型，总是要保存
if metrics and self._is_best_metric(metrics):
return True
# 按步数保存
if step % self.config.save_steps == 0:
return True
return False
def _is_best_metric(self, metrics: Dict[str, float]) -> bool:
"""判断是否是最佳指标"""
if not self.config.save_best_only:
return False
# 使用第一个有值的指标
for key, value in metrics.items():
if value is not None:
if self.best_metric is None:
self.best_metric = value
return True
if self.best_metric_mode == "min":
is_best = value < self.best_metric
else:
is_best = value > self.best_metric
if is_best:
self.best_metric = value
return True
return False
return False
def _cleanup_old_checkpoints(self):
"""清理旧的 checkpoint"""
if self.config.save_total_limit <= 0:
return
# 查找所有 checkpoint
checkpoints = sorted(
self.checkpoint_dir.glob("checkpoint-*/model.pt"),
key=lambda x: os.path.getmtime(x),
reverse=True
)
# 删除超出限制的 checkpoint
for checkpoint in checkpoints[self.config.save_total_limit:]:
checkpoint_dir = checkpoint.parent
shutil.rmtree(checkpoint_dir)
logger.info(f"清理旧 checkpoint: {checkpoint_dir}")
def _save_meta(self, step: int, epoch: int, checkpoint_name: str,
metrics: Optional[Dict]):
"""保存元信息"""
meta = {
"step": step,
"epoch": epoch,
"checkpoint_name": checkpoint_name,
"metrics": metrics,
"timestamp": datetime.now().isoformat()
}
with open(self.checkpoint_dir / f"{checkpoint_name}" / "meta.json", 'w') as f:
json.dump(meta, f, indent=2)
# 更新全局 meta
all_meta = self.load_meta()
all_meta[checkpoint_name] = meta
all_meta["latest"] = checkpoint_name
if metrics:
# 更新 best
best_key = "best_val_loss" if "val_loss" in metrics else list(metrics.keys())[0]
if best_key in metrics:
all_meta["best"] = {
"checkpoint_name": checkpoint_name,
"metric_key": best_key,
"metric_value": metrics[best_key]
}
with open(self.meta_file, 'w') as f:
json.dump(all_meta, f, indent=2)
def load_meta(self) -> Dict:
"""加载元信息"""
if self.meta_file.exists():
with open(self.meta_file, 'r') as f:
return json.load(f)
return {}六、产出检查阶段
6.1 训练监控与早停
#!/usr/bin/env python3
"""
训练监控与早停模块
监控 Loss 曲线、指标变化，执行早停和质量门禁
"""
import numpy as np
from typing import Dict, List, Optional, Callable
from dataclasses import dataclass, field
import json
import logging
from pathlib import Path
import matplotlib.pyplot as plt
logger = logging.getLogger(__name__)
@dataclass
class EarlyStoppingConfig:
"""早停配置"""
patience: int = 5 # 容忍多少个 epoch 没有改善
min_delta: float = 0.001 # 最小改善量
mode: str = "min" # "min" 或 "max"
baseline: Optional[float] = None # 基准线
restore_best_weights: bool = True # 恢复最佳权重
class EarlyStopping:
"""早停机制"""
def __init__(self, config: EarlyStoppingConfig):
self.config = config
self.best_score = float('inf') if config.mode == "min" else float('-inf')
self.best_epoch = 0
self.counter = 0
self.should_stop = False
self.best_weights = None
def __call__(self, score: float, epoch: int, model: Optional[any] = None) -> bool:
"""
检查是否应该早停
Args:
score: 当前指标值
epoch: 当前 epoch
model: 模型（用于保存最佳权重）
Returns:
True 表示应该停止训练
"""
# 基准线检查
if self.config.baseline is not None:
if epoch == 0 and score > self.config.baseline:
self.should_stop = True
logger.warning(f"初始指标 {score} 超过基准线 {self.config.baseline}")
return True
# 判断是否改善
if self.config.mode == "min":
improved = score < (self.best_score - self.config.min_delta)
else:
improved = score > (self.best_score + self.config.min_delta)
if improved:
self.best_score = score
self.best_epoch = epoch
self.counter = 0
# 保存最佳权重
if model is not None and self.config.restore_best_weights:
self.best_weights = model.state_dict().copy()
logger.info(f"新的最佳 score: {score:.6f} (epoch={epoch})")
else:
self.counter += 1
logger.info(f"指标未改善: counter={self.counter}/{self.config.patience}")
# 判断是否应该停止
if self.counter >= self.config.patience:
self.should_stop = True
logger.info(f"早停触发: 连续 {self.counter} 个 epoch 没有改善")
# 恢复最佳权重
if model is not None and self.config.restore_best_weights and self.best_weights:
model.load_state_dict(self.best_weights)
logger.info("已恢复最佳权重")
return True
return False
def reset(self):
"""重置早停状态"""
self.best_score = float('inf') if self.config.mode == "min" else float('-inf')
self.best_epoch = 0
self.counter = 0
self.should_stop = False
@dataclass
class TrainingMonitor:
"""训练监控器"""
log_dir: str = "./logs"
history: Dict[str, List[float]] = field(default_factory=dict)
steps: List[int] = field(default_factory=list)
def log(self, metrics: Dict[str, float], step: int):
"""记录指标"""
self.steps.append(step)
for key, value in metrics.items():
if key not in self.history:
self.history[key] = []
self.history[key].append(value)
def get_history(self, key: str) -> List[float]:
"""获取指标历史"""
return self.history.get(key, [])
def plot_metrics(self, save_path: Optional[str] = None):
"""绘制指标曲线"""
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# Loss 曲线
ax1 = axes[0]
if "train_loss" in self.history:
ax1.plot(self.steps, self.history["train_loss"], label="Train Loss")
if "val_loss" in self.history:
ax1.plot(self.steps, self.history["val_loss"], label="Val Loss")
ax1.set_xlabel("Step")
ax1.set_ylabel("Loss")
ax1.set_title("Loss 曲线")
ax1.legend()
ax1.grid(True, alpha=0.3)
# 准确率曲线
ax2 = axes[1]
if "train_accuracy" in self.history:
ax2.plot(self.steps, self.history["train_accuracy"], label="Train Acc")
if "val_accuracy" in self.history:
ax2.plot(self.steps, self.history["val_accuracy"], label="Val Acc")
ax2.set_xlabel("Step")
ax2.set_ylabel("Accuracy")
ax2.set_title("准确率曲线")
ax2.legend()
ax2.grid(True, alpha=0.3)
plt.tight_layout()
if save_path:
plt.savefig(save_path, dpi=150)
logger.info(f"指标曲线已保存: {save_path}")
else:
plt.show()
def save_history(self, path: str):
"""保存训练历史"""
Path(path).parent.mkdir(parents=True, exist_ok=True)
with open(path, 'w') as f:
json.dump({
"steps": self.steps,
"history": self.history
}, f, indent=2)
logger.info(f"训练历史已保存: {path}")七、流水线编排：基于 Argo Workflows
7.1 Argo Workflows YAML 配置
Argo Workflows 是 Kubernetes 原生的流水线编排工具，非常适合机器学习流水线的编排：
# training_pipeline.yaml
# 模型训练流水线 Argo Workflows 配置
apiVersion: argoproj.io/v1alpha1
kind: Workflow
metadata:
generateName: ml-training-pipeline-
namespace: ml-platform
spec:
# 入口模板
entrypoint: training-pipeline
# 模板定义
templates:
- name: training-pipeline
dag:
tasks:
# ===== 数据准备阶段 =====
- name: data-download
template: data-download
continue-on:
failed: false
- name: data-clean
template: data-clean
depends: data-download
arguments:
parameters:
- name: data-path
value: "{{tasks.data-download.outputs.parameters.data-path}}"
- name: data-split
template: data-split
depends: data-clean
arguments:
parameters:
- name: cleaned-data-path
value: "{{tasks.data-clean.outputs.parameters.cleaned-data-path}}"
# ===== 训练阶段 =====
- name: training
template: training
depends: data-split
arguments:
parameters:
- name: train-data-path
value: "{{tasks.data-split.outputs.parameters.train-data-path}}"
- name: val-data-path
value: "{{tasks.data-split.outputs.parameters.val-data-path}}"
- name: config-path
value: "{{inputs.parameters.config-path}}"
# ===== 评估阶段 =====
- name: evaluation
template: evaluation
depends: training
arguments:
parameters:
- name: model-path
value: "{{tasks.training.outputs.parameters.model-path}}"
- name: test-data-path
value: "{{tasks.data-split.outputs.parameters.test-data-path}}"
# ===== 质量门禁 =====
- name: quality-gate
template: quality-gate
depends: evaluation
arguments:
parameters:
- name: eval-metrics
value: "{{tasks.evaluation.outputs.parameters.eval-metrics}}"
# ===== 模型注册（仅通过质量门禁后执行）=====
- name: model-register
template: model-register
depends: quality-gate
arguments:
parameters:
- name: model-path
value: "{{tasks.training.outputs.parameters.model-path}}"
- name: model-version
value: "{{tasks.training.outputs.parameters.model-version}}"
# ============ 模板定义 ============
# 数据下载模板
- name: data-download
container:
image: python:3.10-slim
command: [python]
args:
- /scripts/download_data.py
- --source-url
- "{{inputs.parameters.source-url}}"
- --output-path
- /data/raw
env:
- name: AWS_ACCESS_KEY_ID
valueFrom:
secretKeyRef:
name: ml-secrets
key: aws-access-key
- name: AWS_SECRET_ACCESS_KEY
valueFrom:
secretKeyRef:
name: ml-secrets
key: aws-secret-key
volumeMounts:
- name: data-volume
mountPath: /data
- name: scripts
mountPath: /scripts
inputs:
parameters:
- name: source-url
outputs:
parameters:
- name: data-path
valueFrom:
path: /data/raw
volumes:
- name: data-volume
persistentVolumeClaim:
claimName: ml-data-pvc
- name: scripts
configMap:
name: ml-scripts
mountPath: /scripts
# 数据清洗模板
- name: data-clean
container:
image: ml-training:latest
command: [python]
args:
- /scripts/clean_data.py
- --input-path
- "{{inputs.parameters.data-path}}"
- --output-path
- /data/cleaned
- --config
- /configs/cleaning_config.yaml
resources:
requests:
memory: 4Gi
cpu: 2
limits:
memory: 8Gi
cpu: 4
volumeMounts:
- name: data-volume
mountPath: /data
- name: scripts
mountPath: /scripts
- name: configs
mountPath: /configs
inputs:
parameters:
- name: data-path
outputs:
parameters:
- name: cleaned-data-path
valueFrom:
path: /data/cleaned
# 数据分割模板
- name: data-split
container:
image: ml-training:latest
command: [python]
args:
- /scripts/split_data.py
- --input-path
- "{{inputs.parameters.cleaned-data-path}}"
- --output-dir
- /data/splits
- --config
- /configs/split_config.yaml
volumeMounts:
- name: data-volume
mountPath: /data
- name: scripts
mountPath: /scripts
- name: configs
mountPath: /configs
inputs:
parameters:
- name: cleaned-data-path
outputs:
parameters:
- name: train-data-path
valueFrom:
path: /data/splits/train
- name: val-data-path
valueFrom:
path: /data/splits/val
- name: test-data-path
valueFrom:
path: /data/splits/test
# 训练模板
- name: training
container:
image: ml-training:latest
command: [python]
args:
- /scripts/train.py
- --config
- "{{inputs.parameters.config-path}}"
- --train-data
- "{{inputs.parameters.train-data-path}}"
- --val-data
- "{{inputs.parameters.val-data-path}}"
- --output-dir
- /outputs/model
env:
- name: MLFLOW_TRACKING_URI
value: "http://mlflow-server:5000"
- name: WANDB_API_KEY
valueFrom:
secretKeyRef:
name: ml-secrets
key: wandb-api-key
resources:
requests:
memory: 16Gi
cpu: 8
nvidia.com/gpu: "1"
limits:
memory: 32Gi
cpu: 16
nvidia.com/gpu: "1"
volumeMounts:
- name: data-volume
mountPath: /data
- name: model-outputs
mountPath: /outputs
inputs:
parameters:
- name: train-data-path
- name: val-data-path
- name: config-path
outputs:
parameters:
- name: model-path
valueFrom:
path: /outputs/model
- name: model-version
valueFrom:
path: /outputs/model/version.txt
volumes:
- name: data-volume
persistentVolumeClaim:
claimName: ml-data-pvc
- name: model-outputs
persistentVolumeClaim:
claimName: ml-models-pvc
# 评估模板
- name: evaluation
container:
image: ml-training:latest
command: [python]
args:
- /scripts/evaluate.py
- --model-path
- "{{inputs.parameters.model-path}}"
- --test-data
- "{{inputs.parameters.test-data-path}}"
- --output-metrics
- /metrics/eval_results.json
volumeMounts:
- name: model-volume
mountPath: /models
- name: metrics-volume
mountPath: /metrics
inputs:
parameters:
- name: model-path
- name: test-data-path
outputs:
parameters:
- name: eval-metrics
valueFrom:
path: /metrics/eval_results.json
# 质量门禁模板
- name: quality-gate
container:
image: python:3.10-slim
command: [python, /scripts/quality_gate.py]
args:
- --metrics
- "{{inputs.parameters.eval-metrics}}"
- --thresholds
- /configs/quality_thresholds.yaml
inputs:
parameters:
- name: eval-metrics
script:
image: python:3.10-slim
command: [python]
source: |
import json
import sys
# 读取评估指标
metrics = json.loads("{{inputs.parameters.eval-metrics}}")
# 质量门禁阈值
thresholds = {
"min_accuracy": 0.85,
"min_f1": 0.80,
"max_latency_ms": 100
}
# 检查是否通过门禁
passed = True
failures = []
if metrics.get("accuracy", 0) < thresholds["min_accuracy"]:
passed = False
failures.append(f"准确率 {metrics['accuracy']} < {thresholds['min_accuracy']}")
if metrics.get("f1", 0) < thresholds["min_f1"]:
passed = False
failures.append(f"F1 {metrics['f1']} < {thresholds['min_f1']}")
if metrics.get("latency_ms", float('inf')) > thresholds["max_latency_ms"]:
passed = False
failures.append(f"延迟 {metrics['latency_ms']} > {thresholds['max_latency_ms']}")
if not passed:
print("质量门禁未通过:")
for f in failures:
print(f" - {f}")
sys.exit(1)
print("质量门禁通过 ✓")
# 模型注册模板
- name: model-register
container:
image: ml-training:latest
command: [python]
args:
- /scripts/register_model.py
- --model-path
- "{{inputs.parameters.model-path}}"
- --version
- "{{inputs.parameters.model-version}}"
- --metrics
- "{{inputs.parameters.eval-metrics}}"
env:
- name: MLFLOW_TRACKING_URI
value: "http://mlflow-server:5000"
volumeMounts:
- name: model-volume
mountPath: /models
# 默认输入参数
arguments:
parameters:
- name: source-url
value: "s3://ml-bucket/raw-data/dataset_v1.tar.gz"
- name: config-path
value: "s3://ml-bucket/configs/training_config.yaml"7.2 质量门禁阈值配置
# quality_thresholds.yaml
# 质量门禁阈值配置
quality_gate:
# 准确性指标
accuracy:
min_value: 0.85
comparison: "greater_than"
weight: 1.0
# F1 分数
f1_score:
min_value: 0.80
comparison: "greater_than"
weight: 1.0
# 精确率
precision:
min_value: 0.75
comparison: "greater_than"
weight: 0.5
# 召回率
recall:
min_value: 0.75
comparison: "greater_than"
weight: 0.5
# 推理延迟（毫秒）
latency:
max_value: 100
comparison: "less_than"
weight: 1.0
# 模型大小（MB）
model_size:
max_value: 500
comparison: "less_than"
weight: 0.3
# 失败重试配置
retry:
max_attempts: 2
backoff_seconds: 60
# 自动回退目标版本
fallback:
enabled: true
min_improvement: 0.02 # 新版本至少要比旧版本好 2% 才能替换八、总结
本文介绍了模型训练流水线的完整设计，涵盖从数据准备到产出检查的全链路：
核心要点回顾：
数据准备是整个流水线的基础，数据质量决定了模型上限。需要建立完善的数据清洗、分割和质量检查流程。
配置管理是实验可复现的关键。使用 YAML 配置 + 参数继承的方式，可以有效管理大量实验。
MLflow 提供了完整的实验追踪能力，从参数、日志到产物都能统一管理，是事实上的 ML 实验追踪标准。
分布式训练需要考虑通信后端选择、资源调度、故障恢复等多个维度。NCCL + DDP 是当前的主流方案。
Checkpoint 管理是长时间训练的生命线。早停机制配合定期存档，可以在效果和效率之间取得平衡。
Argo Workflows 提供了 Kubernetes 原生的流水线编排能力，可以优雅地串联各个阶段，并通过 DAG 依赖管理执行顺序。
质量门禁是自动化流水线的守门人。通过预定义阈值自动判断模型是否达标，避免不合格模型进入下游。
在实际项目中，建议先从单机版流水线开始验证流程，再逐步扩展到分布式训练和Kubernetes 编排。一步到位往往意味着一步到位的失败。
相关技术栈：Argo Workflows, MLflow, PyTorch DDP, Kubernetes, YAML
参考资料：MLOps Best Practices
评论区