#Kubeflow Pipeline + MLflow 集成

本指南展示了 Kubeflow Pipelines (KFP) 组件如何使用 MLflow Python client 将参数、指标和模型记录到 Alauda AI 上的 MLflow。认证以及 workspace/RBAC 遵循 将 MLflow Python SDK 与认证和 RBAC 一起使用 —— 每个组件都会使用用户身份 token 进行认证，server 会以该用户身份记录 run。

#适用范围

• Alauda AI 2.5 及更高版本。
• 已安装 Kubeflow Pipelines 和 MLflow Operator，并且 MLflow custom resource 正在运行。
• MLflow workspace 是一个标注了 mlflow-enabled=true 的 namespace。
• 对于 bearer-token 方法，请保持默认的 spec.auth.oauth.skipJwtBearerTokens: true — 请参见 SDK 指南中的 验证 token 方法。无需更改全局认证，cookie 方法也不需要任何额外配置。

#前提条件

• kfp 和 kfp-kubernetes Python SDK（pip install kfp kfp-kubernetes）。
• 可以访问 KFP endpoint（请参见 使用 Kubeflow Pipelines）。
• 一个用于专用非个人账户的 Dex id token（普通平台登录，不是 Kubernetes ServiceAccount），通过 authorization-code flow 在无浏览器环境下签发（请参见 SDK 指南）。将其存储在 Kubernetes Secret 中，并注入到 component 中。
• 该账户可以访问的 MLflow workspace（一个带有 mlflow-enabled=true 的 namespace）。

#组件如何访问 MLflow

pipeline component 运行在集群内部，因此它通过集群内的 Service http://mlflow-tracking-server.kubeflow:5000 访问 MLflow（该 Service 前面由 OAuth proxy 提供入口 — 组件不会直接使用 MLflow container 端口）。它的认证方式与其他 MLflow client 完全一致：

• MLFLOW_TRACKING_TOKEN — 一个 Dex id token；MLflow client 会将其作为 Authorization: Bearer … 发送。
• mlflow.set_workspace(...) — 选择 workspace（X-MLFLOW-WORKSPACE）。

server 会从 token 中读取身份，并将 run 记录到该用户名下。有关 token 的获取方式以及授权的工作原理，请参见 SDK 指南。

#完整示例：使用 MLflow 的训练 pipeline

该 component 使用 MLflow client，并从通过 kfp-kubernetes 注入的 Secret 中读取 MLFLOW_TRACKING_TOKEN。KFP v2 会从各自的 source 打包每个 component，因此 import mlflow 位于函数内部。

from kfp import dsl, compiler
from kfp import kubernetes


@dsl.component(base_image="python:3.11-slim", packages_to_install=["mlflow==3.13.0"])
def train_model(
    workspace: str,
    model_name: str,
    learning_rate: float,
    epochs: int,
    run_id: str,
) -> dict:
    """Simulated training component that logs to MLflow as the calling user."""
    import mlflow   # MLFLOW_TRACKING_TOKEN is injected from a Secret (see the pipeline below)

    mlflow.set_tracking_uri("http://mlflow-tracking-server.kubeflow:5000")  # in-cluster Service, via the OAuth proxy
    mlflow.set_workspace(workspace)
    mlflow.set_experiment("kfp-training-experiment")

    metrics = {}
    with mlflow.start_run(run_name=f"run-{run_id}"):
        mlflow.log_param("model_name", model_name)
        mlflow.log_param("learning_rate", learning_rate)
        mlflow.log_param("epochs", epochs)
        for epoch in range(1, epochs + 1):
            loss = 2.0 * (0.95 ** epoch)
            accuracy = 1.0 - loss
            mlflow.log_metric("loss", loss, step=epoch)
            mlflow.log_metric("accuracy", accuracy, step=epoch)
            metrics = {"final_loss": loss, "final_accuracy": accuracy}

    print("logged run:", mlflow.last_active_run().info.run_id)
    return metrics


@dsl.pipeline(name="mlflow-training-pipeline", description="Train with MLflow tracking")
def training_pipeline(
    workspace: str = "team-a",
    model_name: str = "qwen3-0.6b",
    learning_rate: float = 2e-4,
    epochs: int = 10,
):
    task = train_model(
        workspace=workspace,
        model_name=model_name,
        learning_rate=learning_rate,
        epochs=epochs,
        # PIPELINE_JOB_ID_PLACEHOLDER resolves to the run's job id at runtime;
        # pass it in as an argument (a component cannot reference dsl.* itself).
        run_id=dsl.PIPELINE_JOB_ID_PLACEHOLDER,
    )
    # Inject the Dex id token from a Secret as MLFLOW_TRACKING_TOKEN.
    kubernetes.use_secret_as_env(
        task, secret_name="mlflow-token", secret_key_to_env={"token": "MLFLOW_TRACKING_TOKEN"}
    )


compiler.Compiler().compile(training_pipeline, "pipeline.yaml")

使用 Dex id token 创建 mlflow-token Secret。按照 SDK 指南中的 authorization-code flow，在无浏览器环境下签发 ID_TOKEN — 请参见 从命令行获取 token：

# ID_TOKEN: mint it with the curl/Python flow in the SDK guide (browser-free, current grants)
kubectl -n <pipeline-namespace> create secret generic mlflow-token --from-literal=token="$ID_TOKEN"

#上传并运行

#通过 KFP UI

1. 进入 Kubeflow Dashboard → Pipelines → Upload Pipeline 并选择 pipeline.yaml。
2. 点击 Create Run 并填写参数（workspace、model name、epochs）。
3. run 启动后，在 Alauda AI → Tools → MLFlow 下检查 MLflow UI — run owner 是 token 的用户。

#通过 KFP SDK

from kfp.client import Client

client = Client(host="<MY-KFP-ENDPOINT>")
run = client.create_run_from_pipeline_package(
    "pipeline.yaml",
    arguments=dict(workspace="team-a", model_name="qwen3-0.6b", epochs=10),
)
print(f"Run ID: {run.run_id}")

#在 Trainer v2 pipeline 中使用 MLflow

如果你使用 Kubeflow Trainer v2 进行 fine-tune，那么 framework 的 MLflow 集成（例如 LLaMA-Factory 中的 report_to: mlflow）采用相同的认证方式。Trainer v2 使用 apiVersion: trainer.kubeflow.org/v1alpha1、kind: TrainJob，以及 spec.runtimeRef + spec.trainer 的结构。将其指向集群内的 Service，并从 Secret 中注入 id token：

apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
metadata:
  name: mlflow-finetune
spec:
  runtimeRef:
    name: torch-distributed        # a TrainingRuntime / ClusterTrainingRuntime
  trainer:
    image: alaudadockerhub/fine_tune_with_llamafactory:v0.1.1
    env:
      - name: MLFLOW_TRACKING_URI
        value: "http://mlflow-tracking-server.kubeflow:5000"
      - name: MLFLOW_EXPERIMENT_NAME
        value: "trainer-v2-finetune"
      - name: MLFLOW_TRACKING_TOKEN
        valueFrom:
          secretKeyRef:
            name: mlflow-token       # a Secret holding a Dex id token
            key: token

完整的 Trainer v2 + MLflow 示例请参见 使用 Workbench 进行 LLM fine-tuning。

如果你需要一个更完整的方案，将此集成与 MLflow Model Registry、TrustyAI LMEvalJob 以及每日运行的 KFP Recurring Run 结合起来，请参见 基于 MLflow 和 TrustyAI 的每日 Fine-Tuning Pipeline。

#最佳实践

#在 MLflow 中使用 pipeline job ID

KFP v2 提供了 dsl.PIPELINE_JOB_ID_PLACEHOLDER（v1 中的 dsl.RUN_ID_PLACEHOLDER 已移除）。它是 pipeline 级别的 placeholder，因此请作为参数传入 component — component 不能在自身内部直接引用 dsl.*。将接收到的字符串用于 run name，以便使每次 pipeline 执行的 runs 保持唯一。

#将凭据保存在 Secret 中并刷新 token

切勿将 token 或账户凭据硬编码到 pipeline.yaml 中 — 编译后的 pipeline 会被存储和共享。请从 Secret 中注入它们，并在 id token 过期前刷新它（或在 component 内部签发）。

#在 run 内记录指标

每个 metric 都属于一个 mlflow.start_run() block。如果 component 有多个逻辑阶段，应为每个阶段分别打开一个 run，而不是在 run context 之外记录。

#面向生产的 artifact 存储

通过代理的 S3-compatible storage，pipeline component 可以在无需获得 S3 凭据的情况下，通过 tracking server 上传。请在记录 artifacts 或注册模型之前配置 artifacts.s3（参见 MLflow installation → 高可用性和存储）。

#故障排查