AI 平台技术栈设计：Kubeflow、MLflow、KServe 集成运维

"""
以 Kubeflow Pipelines v2 component 的形式执行训练，
同时把实验结果记录进 MLflow 的示例。
"""
from kfp import dsl
from kfp.dsl import Input, Output, Dataset, Model, Metrics

@dsl.component(
    base_image="python:3.11-slim",
    packages_to_install=[
        "mlflow==2.17.2",
        "scikit-learn==1.5.2",
        "pandas==2.2.3",
        "boto3==1.35.0",
    ],
)
def train_model(
    training_data: Input[Dataset],
    model_output: Output[Model],
    metrics_output: Output[Metrics],
    mlflow_tracking_uri: str,
    mlflow_experiment_name: str,
    n_estimators: int = 200,
    max_depth: int = 10,
    learning_rate: float = 0.1,
):
    import mlflow
    import pandas as pd
    from sklearn.ensemble import GradientBoostingClassifier
    from sklearn.model_selection import cross_val_score
    import json
    import os

    # 连接 MLflow
    mlflow.set_tracking_uri(mlflow_tracking_uri)
    mlflow.set_experiment(mlflow_experiment_name)

    # 加载数据
    df = pd.read_parquet(training_data.path)
    X = df.drop(columns=["label"])
    y = df["label"]

    with mlflow.start_run() as run:
        # 记录超参数
        mlflow.log_params({
            "n_estimators": n_estimators,
            "max_depth": max_depth,
            "learning_rate": learning_rate,
            "feature_count": X.shape[1],
            "training_rows": X.shape[0],
        })

        # 训练
        clf = GradientBoostingClassifier(
            n_estimators=n_estimators,
            max_depth=max_depth,
            learning_rate=learning_rate,
            random_state=42,
        )
        cv_scores = cross_val_score(clf, X, y, cv=5, scoring="f1_weighted")
        clf.fit(X, y)

        # 记录指标
        mlflow.log_metrics({
            "cv_f1_mean": float(cv_scores.mean()),
            "cv_f1_std": float(cv_scores.std()),
        })

        # 保存模型(MLflow Model Registry 格式)
        mlflow.sklearn.log_model(
            clf,
            artifact_path="model",
            registered_model_name="recommendation-classifier",
        )

        # 同时记录到 Kubeflow Metrics(用于 UI 展示)
        metrics_output.log_metric("cv_f1_mean", float(cv_scores.mean()))
        metrics_output.log_metric("cv_f1_std", float(cv_scores.std()))
        metrics_output.log_metric("mlflow_run_id", run.info.run_id)

        # 把模型路径作为 output 传递(供下一个 component 使用)
        model_output.uri = f"runs:/{run.info.run_id}/model"
        model_output.metadata["mlflow_run_id"] = run.info.run_id

from kfp import dsl, compiler

@dsl.pipeline(
    name="recommendation-training-pipeline",
    description="数据校验 -> 训练 -> 模型注册 -> 服务部署",
)
def training_pipeline(
    data_path: str = "gs://ml-data/recommendation/2026-03-04/",
    mlflow_tracking_uri: str = "http://mlflow.ml-platform.svc:5000",
    mlflow_experiment: str = "recommendation-v3",
    serving_namespace: str = "model-serving",
):
    # Step 1: 数据校验
    validate_task = validate_data(data_path=data_path)

    # Step 2: 预处理
    preprocess_task = preprocess_data(
        raw_data=validate_task.outputs["validated_data"],
    )

    # Step 3: 训练 + MLflow 追踪
    train_task = train_model(
        training_data=preprocess_task.outputs["processed_data"],
        mlflow_tracking_uri=mlflow_tracking_uri,
        mlflow_experiment_name=mlflow_experiment,
        n_estimators=300,
        max_depth=8,
        learning_rate=0.05,
    )
    train_task.set_cpu_limit("4").set_memory_limit("16Gi")
    train_task.set_accelerator_type("nvidia.com/gpu").set_accelerator_limit(1)

    # Step 4: 模型质量门禁
    gate_task = quality_gate(
        metrics=train_task.outputs["metrics_output"],
        f1_threshold=0.80,
    )

    # Step 5: 部署到 KServe(质量门禁通过时)
    with dsl.Condition(gate_task.outputs["passed"] == "true"):
        deploy_task = deploy_to_kserve(
            model=train_task.outputs["model_output"],
            serving_namespace=serving_namespace,
        )

compiler.Compiler().compile(
    pipeline_func=training_pipeline,
    package_path="recommendation_pipeline.yaml",
)

@dsl.component(
    base_image="python:3.11-slim",
    packages_to_install=["kubernetes==31.0.0", "mlflow==2.17.2"],
)
def deploy_to_kserve(
    model: Input[Model],
    serving_namespace: str,
    canary_traffic_percent: int = 20,
):
    """
    把注册在 MLflow 中的模型部署为 KServe InferenceService。
    以 canary 方式为新版本分配一部分流量。
    """
    from kubernetes import client, config
    import json
    import mlflow

    config.load_incluster_config()
    api = client.CustomObjectsApi()

    mlflow_run_id = model.metadata.get("mlflow_run_id")
    model_uri = model.uri  # runs:/<run_id>/model

    # 从 MLflow 中提取 GCS 路径
    mlflow_client = mlflow.tracking.MlflowClient()
    run = mlflow_client.get_run(mlflow_run_id)
    artifact_uri = run.info.artifact_uri  # gs://mlflow-artifacts/<exp>/<run>/artifacts

    storage_uri = f"{artifact_uri}/model"

    inference_service = {
        "apiVersion": "serving.kserve.io/v1beta1",
        "kind": "InferenceService",
        "metadata": {
            "name": "recommendation-classifier",
            "namespace": serving_namespace,
            "labels": {
                "mlflow-run-id": mlflow_run_id,
                "pipeline": "recommendation-training",
            },
            "annotations": {
                "serving.kserve.io/deploymentMode": "Serverless",
                "serving.kserve.io/autoscalerClass": "hpa",
                "serving.kserve.io/metrics": "cpu",
                "serving.kserve.io/targetUtilizationPercentage": "70",
            },
        },
        "spec": {
            "predictor": {
                "canaryTrafficPercent": canary_traffic_percent,
                "minReplicas": 2,
                "maxReplicas": 10,
                "model": {
                    "modelFormat": {"name": "mlflow"},
                    "storageUri": storage_uri,
                    "resources": {
                        "requests": {"cpu": "1", "memory": "2Gi"},
                        "limits": {"cpu": "2", "memory": "4Gi"},
                    },
                },
            },
        },
    }

    try:
        api.patch_namespaced_custom_object(
            group="serving.kserve.io",
            version="v1beta1",
            namespace=serving_namespace,
            plural="inferenceservices",
            name="recommendation-classifier",
            body=inference_service,
        )
        print(f"Updated InferenceService with canary {canary_traffic_percent}%")
    except client.exceptions.ApiException as e:
        if e.status == 404:
            api.create_namespaced_custom_object(
                group="serving.kserve.io",
                version="v1beta1",
                namespace=serving_namespace,
                plural="inferenceservices",
                body=inference_service,
            )
            print("Created new InferenceService")
        else:
            raise

"""
Canary 部署后在一段时间内持续观察指标，
自动判定晋升或回滚的脚本。
以 CronJob 或 Argo Workflow step 的形式运行。
"""
import requests
import time
from dataclasses import dataclass
from typing import Optional

@dataclass
class CanaryConfig:
    service_name: str
    namespace: str
    prometheus_url: str
    observation_minutes: int = 15
    check_interval_seconds: int = 60
    error_rate_threshold: float = 0.02    # 2%
    p99_latency_threshold_ms: float = 500  # 500ms
    min_request_count: int = 100           # 最小观察请求数

def query_prometheus(config: CanaryConfig, query: str) -> Optional[float]:
    """从 Prometheus 查询指标值。"""
    resp = requests.get(
        f"{config.prometheus_url}/api/v1/query",
        params={"query": query},
        timeout=10,
    )
    results = resp.json().get("data", {}).get("result", [])
    if results:
        return float(results[0]["value"][1])
    return None

def evaluate_canary(config: CanaryConfig) -> dict:
    """评估 Canary 版本的 error rate 和 latency。"""
    error_rate_query = (
        f'sum(rate(kserve_request_total{{service_name="{config.service_name}",'
        f'namespace="{config.namespace}",response_code=~"5.."}}[5m])) / '
        f'sum(rate(kserve_request_total{{service_name="{config.service_name}",'
        f'namespace="{config.namespace}"}}[5m]))'
    )
    p99_query = (
        f'histogram_quantile(0.99, sum(rate(kserve_request_duration_seconds_bucket'
        f'{{service_name="{config.service_name}",namespace="{config.namespace}"}}'
        f'[5m])) by (le)) * 1000'
    )
    request_count_query = (
        f'sum(increase(kserve_request_total{{service_name="{config.service_name}",'
        f'namespace="{config.namespace}"}}[{config.observation_minutes}m]))'
    )

    error_rate = query_prometheus(config, error_rate_query) or 0.0
    p99_latency = query_prometheus(config, p99_query) or 0.0
    request_count = query_prometheus(config, request_count_query) or 0

    return {
        "error_rate": error_rate,
        "p99_latency_ms": p99_latency,
        "request_count": request_count,
        "error_rate_ok": error_rate < config.error_rate_threshold,
        "latency_ok": p99_latency < config.p99_latency_threshold_ms,
        "sufficient_traffic": request_count >= config.min_request_count,
    }

def run_canary_judgment(config: CanaryConfig) -> str:
    """
    在 observation_minutes 时间内观察指标，并决定晋升或回滚。
    Returns: "promote" or "rollback"
    """
    checks_passed = 0
    total_checks = config.observation_minutes * 60 // config.check_interval_seconds

    for i in range(total_checks):
        result = evaluate_canary(config)
        print(f"Check {i+1}/{total_checks}: {result}")

        if not result["sufficient_traffic"]:
            print("Insufficient traffic, waiting...")
            time.sleep(config.check_interval_seconds)
            continue

        if result["error_rate_ok"] and result["latency_ok"]:
            checks_passed += 1
        else:
            # 立即回滚条件：error rate 超过阈值的 3 倍
            if result["error_rate"] > config.error_rate_threshold * 3:
                print(f"Immediate rollback: error_rate={result['error_rate']}")
                return "rollback"

        time.sleep(config.check_interval_seconds)

    pass_ratio = checks_passed / max(total_checks, 1)
    decision = "promote" if pass_ratio >= 0.9 else "rollback"
    print(f"Decision: {decision} (pass_ratio={pass_ratio:.2f})")
    return decision

症状：Pipeline 显示成功(green)，但 MLflow UI 里没有对应的 Run
错误日志(Pipeline pod)：
  requests.exceptions.ConnectionError: HTTPConnectionPool(host='mlflow.ml-platform.svc', port=5000):
  Max retries exceeded with url: /api/2.0/mlflow/runs/create

原因：Kubeflow Pipeline pod 的 ServiceAccount 没有被 NetworkPolicy 允许
      访问 MLflow service

解决：
  1. 在 NetworkPolicy 中放行 pipeline runner namespace -> mlflow namespace
  2. 确认 MLflow tracking URI 是否为 cluster-internal DNS
  3. 给 Pipeline component 加上 retry 逻辑

症状：InferenceService 卡在 "FailedCreate" 状态
错误日志(kserve-container)：
  mlflow.exceptions.MlflowException: Could not find an "MLmodel" configuration file
  at "gs://mlflow-artifacts/3/abc123def/artifacts/model"

原因：MLflow log_model 的 artifact_path 与 KServe storageUri 的路径不一致
      (artifact_path="model"，但 storageUri 里缺了 "/model")

解决：
  storageUri 需要按 artifact_uri + "/model" 的形式拼接。
  错误示例：gs://mlflow-artifacts/3/abc123def/artifacts
  正确示例：gs://mlflow-artifacts/3/abc123def/artifacts/model

症状：即便把 canaryTrafficPercent 设为 0，也回不到旧版本
错误日志(kserve-controller)：
  RevisionFailed: Revision "recommendation-classifier-predictor-prev"
  has no ready pods

原因：KServe 把旧 Revision 的 pod scale-to-zero 之后，
      cold start 耗时超过了 readiness probe timeout

解决：
  1. 给服务设置 minReplicas >= 1(重要服务应禁用 scale-to-zero)
  2. 把 readinessProbe timeout 按模型加载时间调大
  3. 回滚时不要靠调整 canary 比例，而要把 storageUri 换回旧版本

# 回滚时应用的 KServe 清单
apiVersion: serving.kserve.io/v1beta1
kind: InferenceService
metadata:
  name: recommendation-classifier
  namespace: model-serving
spec:
  predictor:
    canaryTrafficPercent: 0 # 清除 canary 流量
    minReplicas: 2 # 防止 scale-to-zero
    model:
      modelFormat:
        name: mlflow
      # 替换为此前稳定版本的 model URI
      storageUri: gs://mlflow-artifacts/3/previous-stable-run/artifacts/model
      resources:
        requests:
          cpu: '1'
          memory: '2Gi'
        limits:
          cpu: '2'
          memory: '4Gi'
    containerSpec:
      readinessProbe:
        initialDelaySeconds: 30
        timeoutSeconds: 10
        periodSeconds: 5