Kubeflow Pipelines ML 工作流编排实战指南：从 KFP v2 SDK 到生产部署

# KFP v2 与 v1 的主要差异对比
"""
KFP v1:
- 基于 Argo Workflows 执行
- 使用 kfp.dsl.ContainerOp
- 可以用 YAML 定义流水线

KFP v2:
- 自研工作流引擎(也可以选择 Argo)
- 使用 kfp.dsl.component 装饰器
- 引入 IR(Intermediate Representation) YAML
- 原生集成 ML Metadata
- 类型安全的组件接口
"""

# v2 架构分层
ARCHITECTURE_LAYERS = {
    "SDK Layer": "用 Python DSL 定义流水线 (kfp.dsl)",
    "IR Layer": "平台无关的中间表示 (PipelineSpec YAML)",
    "Backend Layer": "流水线执行与管理 (API Server)",
    "Runtime Layer": "容器编排 (K8s Pod)",
    "Metadata Layer": "执行历史与 Artifact 追踪 (MLMD)",
}

from kfp import dsl
from kfp.dsl import Input, Output, Dataset, Model, Metrics

# 轻量级 Python 组件(依赖较少的情况)
@dsl.component(
    base_image="python:3.11-slim",
    packages_to_install=["pandas==2.1.4", "scikit-learn==1.4.0"],
)
def preprocess_data(
    raw_data_path: str,
    test_size: float,
    train_dataset: Output[Dataset],
    test_dataset: Output[Dataset],
    metrics: Output[Metrics],
):
    """数据预处理组件"""
    import pandas as pd
    from sklearn.model_selection import train_test_split

    df = pd.read_csv(raw_data_path)

    # 处理缺失值
    df = df.dropna(subset=["target"])
    df = df.fillna(df.median(numeric_only=True))

    # 拆分训练/测试集
    train_df, test_df = train_test_split(
        df, test_size=test_size, random_state=42, stratify=df["target"]
    )

    # 保存为 Artifact
    train_df.to_csv(train_dataset.path, index=False)
    test_df.to_csv(test_dataset.path, index=False)

    # 记录指标
    metrics.log_metric("total_samples", len(df))
    metrics.log_metric("train_samples", len(train_df))
    metrics.log_metric("test_samples", len(test_df))
    metrics.log_metric("feature_count", len(df.columns) - 1)

# 基于自定义镜像的组件
@dsl.component(
    base_image="gcr.io/my-project/ml-training:v2.1",
)
def train_model(
    train_dataset: Input[Dataset],
    model_type: str,
    hyperparameters: dict,
    trained_model: Output[Model],
    metrics: Output[Metrics],
):
    """模型训练组件"""
    import pandas as pd
    import joblib
    from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
    from sklearn.metrics import accuracy_score, f1_score

    train_df = pd.read_csv(train_dataset.path)
    X_train = train_df.drop("target", axis=1)
    y_train = train_df["target"]

    # 选择模型
    model_map = {
        "random_forest": RandomForestClassifier,
        "gradient_boosting": GradientBoostingClassifier,
    }
    model_cls = model_map[model_type]
    model = model_cls(**hyperparameters)
    model.fit(X_train, y_train)

    # 保存模型
    joblib.dump(model, trained_model.path)

    # 训练指标
    y_pred = model.predict(X_train)
    metrics.log_metric("train_accuracy", accuracy_score(y_train, y_pred))
    metrics.log_metric("train_f1", f1_score(y_train, y_pred, average="weighted"))
    metrics.log_metric("model_type", model_type)

    # 模型元数据
    trained_model.metadata["framework"] = "scikit-learn"
    trained_model.metadata["model_type"] = model_type

from kfp import dsl, compiler

@dsl.pipeline(
    name="ml-training-pipeline",
    description="End-to-end ML training pipeline with evaluation and deployment",
)
def ml_training_pipeline(
    raw_data_path: str = "gs://my-bucket/data/raw.csv",
    test_size: float = 0.2,
    model_type: str = "random_forest",
    accuracy_threshold: float = 0.85,
):
    # Step 1: 数据预处理
    preprocess_task = preprocess_data(
        raw_data_path=raw_data_path,
        test_size=test_size,
    )
    preprocess_task.set_cpu_limit("2")
    preprocess_task.set_memory_limit("4Gi")

    # Step 2: 模型训练
    train_task = train_model(
        train_dataset=preprocess_task.outputs["train_dataset"],
        model_type=model_type,
        hyperparameters={
            "n_estimators": 200,
            "max_depth": 10,
            "min_samples_split": 5,
        },
    )
    train_task.set_cpu_limit("4")
    train_task.set_memory_limit("8Gi")
    train_task.set_accelerator_type("nvidia.com/gpu")
    train_task.set_accelerator_limit(1)

    # Step 3: 模型评估
    eval_task = evaluate_model(
        test_dataset=preprocess_task.outputs["test_dataset"],
        trained_model=train_task.outputs["trained_model"],
        accuracy_threshold=accuracy_threshold,
    )

    # Step 4: 条件部署(超过准确率阈值时)
    with dsl.Condition(
        eval_task.outputs["deploy_decision"] == "approved",
        name="check-accuracy",
    ):
        deploy_task = deploy_model(
            model=train_task.outputs["trained_model"],
            serving_endpoint="ml-model-serving",
        )

# 编译流水线
compiler.Compiler().compile(
    pipeline_func=ml_training_pipeline,
    package_path="ml_pipeline.yaml",
)

@dsl.pipeline(name="parallel-training-pipeline")
def parallel_training_pipeline(
    raw_data_path: str,
    accuracy_threshold: float = 0.85,
):
    # 数据预处理(公共步骤)
    preprocess_task = preprocess_data(
        raw_data_path=raw_data_path,
        test_size=0.2,
    )

    # 并行训练多个模型
    models = ["random_forest", "gradient_boosting", "xgboost"]
    train_tasks = []

    for model_type in models:
        train_task = train_model(
            train_dataset=preprocess_task.outputs["train_dataset"],
            model_type=model_type,
            hyperparameters={"n_estimators": 200, "max_depth": 10},
        )
        train_task.set_display_name(f"Train {model_type}")
        train_tasks.append(train_task)

    # 选出最优模型
    select_task = select_best_model(
        models=[t.outputs["trained_model"] for t in train_tasks],
        metrics=[t.outputs["metrics"] for t in train_tasks],
    )

    # 部署冠军模型
    with dsl.Condition(
        select_task.outputs["best_accuracy"] >= accuracy_threshold,
        name="accuracy-gate",
    ):
        deploy_model(
            model=select_task.outputs["best_model"],
            serving_endpoint="champion-model",
        )

@dsl.pipeline(name="robust-ml-pipeline")
def robust_ml_pipeline(raw_data_path: str):
    # Exit Handler：流水线完成/失败时发送通知
    notify_task = send_notification(
        pipeline_name="robust-ml-pipeline",
        notification_channel="slack",
    )

    with dsl.ExitHandler(exit_task=notify_task):
        # 主流水线逻辑
        preprocess_task = preprocess_data(
            raw_data_path=raw_data_path,
            test_size=0.2,
        )

        train_task = train_model(
            train_dataset=preprocess_task.outputs["train_dataset"],
            model_type="gradient_boosting",
            hyperparameters={"n_estimators": 300, "max_depth": 12},
        )

        eval_task = evaluate_model(
            test_dataset=preprocess_task.outputs["test_dataset"],
            trained_model=train_task.outputs["trained_model"],
            accuracy_threshold=0.85,
        )

# 设置 Recurring Run(KFP 客户端)
from kfp.client import Client

client = Client(host="https://kubeflow.example.com/pipeline")

# 每天凌晨 2 点执行流水线
client.create_recurring_run(
    experiment_id="daily-training-exp",
    job_name="daily-model-retraining",
    pipeline_id="robust-ml-pipeline-v2",
    cron_expression="0 2 * * *",
    max_concurrency=1,
    params={
        "raw_data_path": "gs://my-bucket/data/daily/latest.csv",
    },
)

# 缓存设置
@dsl.pipeline(name="cached-pipeline")
def cached_pipeline(raw_data_path: str):
    # 启用缓存(默认值：True)
    preprocess_task = preprocess_data(
        raw_data_path=raw_data_path,
        test_size=0.2,
    )
    preprocess_task.set_caching_options(enable_caching=True)

    # 训练阶段禁用缓存(始终用最新数据重新训练)
    train_task = train_model(
        train_dataset=preprocess_task.outputs["train_dataset"],
        model_type="random_forest",
        hyperparameters={"n_estimators": 200},
    )
    train_task.set_caching_options(enable_caching=False)

from kfp.dsl import (
    Input, Output,
    Dataset, Model, Metrics,
    ClassificationMetrics, SlicedClassifications,
    Artifact, HTML, Markdown,
)

@dsl.component(base_image="python:3.11-slim")
def generate_evaluation_report(
    test_dataset: Input[Dataset],
    trained_model: Input[Model],
    classification_metrics: Output[ClassificationMetrics],
    html_report: Output[HTML],
    eval_metrics: Output[Metrics],
):
    """评估报告生成组件"""
    import json

    # ClassificationMetrics：混淆矩阵可视化
    classification_metrics.log_confusion_matrix(
        categories=["negative", "positive"],
        matrix=[[850, 50], [30, 270]],
    )

    # 记录 ROC 曲线
    classification_metrics.log_roc_curve(
        fpr=[0.0, 0.1, 0.2, 0.5, 1.0],
        tpr=[0.0, 0.6, 0.8, 0.95, 1.0],
        threshold=[1.0, 0.8, 0.5, 0.2, 0.0],
    )

    # 生成 HTML 报告
    report_content = "<h1>Model Evaluation Report</h1>"
    report_content += "<p>Accuracy: 0.933</p>"
    report_content += "<p>F1 Score: 0.891</p>"

    with open(html_report.path, "w") as f:
        f.write(report_content)

    # 数值指标
    eval_metrics.log_metric("accuracy", 0.933)
    eval_metrics.log_metric("f1_score", 0.891)
    eval_metrics.log_metric("precision", 0.844)
    eval_metrics.log_metric("recall", 0.900)

from kfp import dsl, compiler
from kfp.dsl import Input, Output, Dataset, Model, Metrics

@dsl.component(
    base_image="python:3.11-slim",
    packages_to_install=["pandas==2.1.4", "great-expectations==0.18.8"],
)
def validate_data(
    raw_data_path: str,
    validated_data: Output[Dataset],
    validation_metrics: Output[Metrics],
) -> str:
    """数据质量校验"""
    import pandas as pd

    df = pd.read_csv(raw_data_path)

    # 基本数据质量检查
    checks = {
        "row_count_check": len(df) > 100,
        "null_ratio_check": df.isnull().mean().max() < 0.3,
        "duplicate_check": df.duplicated().mean() < 0.05,
        "target_balance_check": df["target"].value_counts(normalize=True).min() > 0.1,
    }

    all_passed = all(checks.values())

    for check_name, passed in checks.items():
        validation_metrics.log_metric(check_name, int(passed))

    validation_metrics.log_metric("total_rows", len(df))
    validation_metrics.log_metric("all_checks_passed", int(all_passed))

    if all_passed:
        df.to_csv(validated_data.path, index=False)
        return "passed"
    else:
        failed = [k for k, v in checks.items() if not v]
        raise ValueError(f"Data validation failed: {failed}")


@dsl.component(
    base_image="python:3.11-slim",
    packages_to_install=["pandas==2.1.4", "scikit-learn==1.4.0"],
)
def feature_engineering(
    validated_data: Input[Dataset],
    feature_config: dict,
    features_dataset: Output[Dataset],
    feature_metrics: Output[Metrics],
):
    """特征工程"""
    import pandas as pd
    import numpy as np
    from sklearn.preprocessing import StandardScaler, LabelEncoder

    df = pd.read_csv(validated_data.path)

    # 数值型特征做标准化
    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
    numeric_cols = [c for c in numeric_cols if c != "target"]

    scaler = StandardScaler()
    df[numeric_cols] = scaler.fit_transform(df[numeric_cols])

    # 类别型特征编码
    cat_cols = df.select_dtypes(include=["object"]).columns.tolist()
    for col in cat_cols:
        le = LabelEncoder()
        df[col] = le.fit_transform(df[col].astype(str))

    df.to_csv(features_dataset.path, index=False)

    feature_metrics.log_metric("numeric_features", len(numeric_cols))
    feature_metrics.log_metric("categorical_features", len(cat_cols))
    feature_metrics.log_metric("total_features", len(df.columns) - 1)


@dsl.component(
    base_image="gcr.io/my-project/ml-serving:v1.0",
)
def deploy_to_kserve(
    model: Input[Model],
    serving_endpoint: str,
    namespace: str,
) -> str:
    """把模型部署到 KServe"""
    import subprocess
    import json
    import yaml

    inference_service = {
        "apiVersion": "serving.kserve.io/v1beta1",
        "kind": "InferenceService",
        "metadata": {
            "name": serving_endpoint,
            "namespace": namespace,
        },
        "spec": {
            "predictor": {
                "model": {
                    "modelFormat": {"name": "sklearn"},
                    "storageUri": model.uri,
                    "resources": {
                        "requests": {"cpu": "1", "memory": "2Gi"},
                        "limits": {"cpu": "2", "memory": "4Gi"},
                    },
                }
            }
        },
    }

    manifest_path = "/tmp/isvc.yaml"
    with open(manifest_path, "w") as f:
        yaml.dump(inference_service, f)

    result = subprocess.run(
        ["kubectl", "apply", "-f", manifest_path],
        capture_output=True, text=True,
    )

    if result.returncode != 0:
        raise RuntimeError(f"Deploy failed: {result.stderr}")

    return f"Deployed to {namespace}/{serving_endpoint}"


@dsl.pipeline(
    name="e2e-ml-pipeline",
    description="从数据校验到模型部署的完整 ML 流水线",
)
def e2e_ml_pipeline(
    raw_data_path: str = "gs://ml-data/raw/dataset.csv",
    model_type: str = "gradient_boosting",
    accuracy_threshold: float = 0.85,
    serving_endpoint: str = "fraud-detector",
    namespace: str = "ml-serving",
):
    # 用于通知的 Exit Handler
    notify = send_notification(
        pipeline_name="e2e-ml-pipeline",
        notification_channel="slack",
    )

    with dsl.ExitHandler(exit_task=notify):
        # 1. 数据校验
        validate_task = validate_data(raw_data_path=raw_data_path)

        # 2. 特征工程
        feature_task = feature_engineering(
            validated_data=validate_task.outputs["validated_data"],
            feature_config={"scaling": "standard", "encoding": "label"},
        )

        # 3. 数据拆分
        split_task = preprocess_data(
            raw_data_path=feature_task.outputs["features_dataset"].uri,
            test_size=0.2,
        )

        # 4. 模型训练
        train_task = train_model(
            train_dataset=split_task.outputs["train_dataset"],
            model_type=model_type,
            hyperparameters={"n_estimators": 300, "max_depth": 12},
        )
        train_task.set_cpu_limit("4")
        train_task.set_memory_limit("16Gi")

        # 5. 模型评估
        eval_task = evaluate_model(
            test_dataset=split_task.outputs["test_dataset"],
            trained_model=train_task.outputs["trained_model"],
            accuracy_threshold=accuracy_threshold,
        )

        # 6. 条件部署
        with dsl.Condition(
            eval_task.outputs["deploy_decision"] == "approved",
            name="deploy-gate",
        ):
            deploy_to_kserve(
                model=train_task.outputs["trained_model"],
                serving_endpoint=serving_endpoint,
                namespace=namespace,
            )

# 编译并提交
compiler.Compiler().compile(
    pipeline_func=e2e_ml_pipeline,
    package_path="e2e_ml_pipeline.yaml",
)

@dsl.pipeline(name="resource-managed-pipeline")
def resource_managed_pipeline():
    train_task = train_model(
        train_dataset=preprocess_task.outputs["train_dataset"],
        model_type="xgboost",
        hyperparameters={"n_estimators": 500},
    )

    # 资源限制
    train_task.set_cpu_limit("8")
    train_task.set_memory_limit("32Gi")
    train_task.set_accelerator_type("nvidia.com/gpu")
    train_task.set_accelerator_limit(2)

    # 节点选择器(在 GPU 节点上执行)
    train_task.add_node_selector_constraint(
        label_name="cloud.google.com/gke-accelerator",
        value="nvidia-tesla-v100",
    )

    # Toleration 设置
    train_task.set_gpu_limit(2).add_toleration(
        key="nvidia.com/gpu",
        operator="Exists",
        effect="NoSchedule",
    )

    # 挂载 PVC(大容量数据)
    train_task.add_pvolumes({
        "/mnt/data": dsl.PipelineVolume(
            pvc="ml-data-pvc",
            volume_name="data-volume",
        ),
    })

    # 超时设置(以秒为单位)
    train_task.set_timeout(3600)  # 1小时

    # 重试设置
    train_task.set_retry(
        num_retries=3,
        policy="Always",
        backoff_duration="30s",
        backoff_factor=2.0,
        backoff_max_duration="600s",
    )

# 检查 Pod 状态
kubectl get pods -n kubeflow -l pipeline/runid=run-abc123
kubectl describe pod train-model-xxxxx -n kubeflow

# 在事件中确认 OOMKilled
# Last State: Terminated
#   Reason: OOMKilled
#   Exit Code: 137

# 1. 提高内存限制
train_task.set_memory_limit("64Gi")

# 2. 把组件改成按数据块(chunk)处理
@dsl.component(base_image="python:3.11-slim")
def train_with_chunks(
    train_dataset: Input[Dataset],
    chunk_size: int,
    trained_model: Output[Model],
):
    import pandas as pd
    from sklearn.linear_model import SGDClassifier

    model = SGDClassifier(loss="log_loss")
    chunks = pd.read_csv(train_dataset.path, chunksize=chunk_size)

    for chunk in chunks:
        X = chunk.drop("target", axis=1)
        y = chunk["target"]
        model.partial_fit(X, y, classes=[0, 1])

    import joblib
    joblib.dump(model, trained_model.path)

from kfp.client import Client

client = Client(host="https://kubeflow.example.com/pipeline")

# 1. 停用既有 Recurring Run
client.disable_recurring_run(recurring_run_id="run-xxx")

# 2. 上传新版本流水线
pipeline_version = client.upload_pipeline_version(
    pipeline_package_path="ml_pipeline_v3.yaml",
    pipeline_version_name="v3.0",
    pipeline_id="ml-training-pipeline",
)

# 3. 创建新的 Recurring Run
client.create_recurring_run(
    experiment_id="daily-training-exp",
    job_name="daily-model-retraining-v3",
    version_id=pipeline_version.pipeline_version_id,
    cron_expression="0 2 * * *",
    max_concurrency=1,
)

# 检查 MLMD 服务状态
kubectl get pods -n kubeflow -l app=metadata-grpc-server
kubectl logs metadata-grpc-server-xxxxx -n kubeflow

# 检查 MySQL/PostgreSQL 连接
kubectl exec -it metadata-grpc-server-xxxxx -n kubeflow -- \
    mysql -h metadata-db -u root -p -e "SHOW DATABASES;"

# 重启 MLMD 服务
kubectl rollout restart deployment metadata-grpc-server -n kubeflow