打造 ML/AI 系统的内部开发者平台(IDP)——编排机器学习(ML)流水线

49 阅读21分钟

本章将涵盖:

  • 使用 Kubeflow Pipelines 构建用于模型推理的批处理流水线
  • 创建一条完整的批处理推理工作流:从加载数据到运行模型推理

在第 4 章中,我们用 MLflow 建立了可靠的实验追踪,用 Feast 建立了特征管理。但这些工具仍需要人工介入来协调模型训练、特征更新与推理。这正是流水线编排变得关键的地方(见图 5.1)。

image.png

图 5.1 心智地图:我们现在聚焦 Kubeflow Pipelines 编排(A)

本章我们将使用 Kubeflow Pipelines(KFP)来自动化这些工作流,使我们的 ML 系统更可扩展、更可复现。通过一个收入分类的实践示例,我们将看到如何把手工步骤转化为自动化、可复用的流水线组件。本章全部代码在 GitHub:https://github.com/practical-mlops/chapter-5

5.1 Kubeflow Pipelines:任务编排器(Task orchestrator)

大多数 ML 推理流水线都有一个共同结构:我们需要从某处获取数据(对象存储、数据仓库、文件系统),对数据进行预处理,获取或加载模型,然后执行推理。推理结果会写入数据库或上传到某个云存储。流水线需要周期性运行,我们可能还需要提供运行时参数,例如日期/时间、特征表名等。KFP 可以实现以上所有能力。

KFP 帮助我们构建 ML 流水线,并提供一个友好的 Python SDK(kfp 包),我们可以用它来定义流水线。

我们会先定义流水线步骤(也称为 components),再把这些 components 组合起来形成一条 Kubeflow pipeline。每个 component 都在自己的 Kubernetes pod 中执行,确保每一步都在干净、隔离的环境中运行。

作为 ML 工程师,我们经常会遇到这样的问题:如何让 ML 工作流可重复?如何在不同模型之间标准化步骤?如何保证步骤之间的数据流可靠?Kubeflow components 通过提供一种标准化的方式来打包 ML 任务,帮助解决这些挑战。下面我们用实践来看看它怎么工作。

5.1.1 Kubeflow components

在上一章,我们训练了模型并在特征库中注册了必要特征。现在该做模型推理了。我们将创建一条 Kubeflow pipeline 来完成以下事情:

  1. 从 MinIO 读取数据(包含一份 user IDs 列表)。
  2. 为这份 user IDs 列表检索特征。
  3. 检索模型并执行推理。
  4. 把结果写回 MinIO。

这些任务都可以转成 Kubeflow pipeline 的 component。component 是 ML 工作流中可复用、可组合的工作单元。它们封装单个任务或步骤,执行特定动作,例如数据预处理、模型训练、评估、部署等。把工作流拆成 components 后,我们就能构建模块化、可维护的 ML 流水线。在接下来的小节中,我们会构建这些 components,并把它们组合成一条 pipeline。

读取数据组件(READ DATA COMPONENT)

每个 Kubeflow component 都在一个 Kubernetes pod 中执行(第 3 章 3.3.3 节)。而一个 Kubernetes pod 需要一个容器来执行。因此,为了构建 component,我们需要构建一个包含 component 源码的镜像。我们从读取数据组件开始,如清单 5.1 所示。我们先写 read_data.py:它包含如何从 MinIO 读取数据的逻辑。这是一个普通 Python 文件:先用 host、access key、secret key 初始化 MinIO client,然后用 client 拉取 user ID 文件并读成 pandas DataFrame。

清单 5.1 读取数据组件

def get_data(
    minio_host: str,
    access_key: str,
    secret_key: str,
    bucket_name: str,
    file_name: str,
    data_output_path: str,
):

    client = Minio(
        endpoint=minio_host,
        access_key=access_key,
        secret_key=secret_key,
        secure=False,
    )
    local_temp_file = os.path.join("/tmp", file_name)
    print(f"Downloading {file_name} from bucket {bucket_name}...")

    client.fget_object(
        bucket_name,
        file_name,
        local_temp_file,
    )  #1
    print(f"Reading the downloaded file {local_temp_file}...")
    df = pd.read_parquet(local_temp_file)    #2
    print(f"Saving the processed data to {data_output_path}..."
#1 从 MinIO 下载 Parquet 文件
#2 将 Parquet 文件加载为 DataFrame

pandas DataFrame 需要写到某个位置,这样下一个 component 才能读取或使用它。这个位置就是 data_output_path。当我们把 DataFrame 写到 data_output_path,后续流水线组件就能复用该文件。为此,我们先用 mkdir 创建输出路径目录,然后把 DataFrame 写到该目录:

Path(data_output_path).parent.mkdir(parents=True, exist_ok=True)
df.to_parquet(data_output_path)

然后我们在 read_data.py 中定义 main,以便把参数传给 get_data,如清单 5.2 所示。我们用 argparse 来定义参数。

清单 5.2 读取数据组件的 main

def main():
    parser = argparse.ArgumentParser(
        description="Download data from Minio and save as a Dataset."
    )

    parser.add_argument(
        "--minio_host",
        type=str,
        required=True,
        help="Minio host URL",
    )
    parser.add_argument(
        "--access_key", type=str, required=True, help="Minio access key"
    )
    parser.add_argument(
        "--secret_key", type=str, required=True, help="Minio secret key"
    )
    parser.add_argument(
        "--bucket_name", type=str, required=True, help="Minio bucket name"
    )
    parser.add_argument(
        "--file_name", type=str, required=True, help="File name to download"
    )
    parser.add_argument(
        "--data_output_path",
        type=str,
        required=True,
        help="Output path for the Dataset",
    )
    args = parser.parse_args()
    get_data(
        minio_host=args.minio_host,
        access_key=args.access_key,
        secret_key=args.secret_key,
        bucket_name=args.bucket_name,
        file_name=args.file_name,
        data_output_path=args.data_output_path,
    )
if __name__ == "__main__":
    main()

现在我们有了 Python 代码:可以从 MinIO 读数据并把它写到流水线的 artifact store。下一步要把这段代码容器化:构建 Docker 镜像,因此需要写 Dockerfile,如清单 5.3 所示。Dockerfile 以 Python 3.10 基础镜像开始,设置工作目录 /app,安装 build-essential 并清理临时目录。然后复制 requirements.txt 并安装依赖(包含 component 所需的全部依赖)。最后把源码复制进工作目录。

清单 5.3 Component Dockerfile

FROM python:3.10-slim-buster
WORKDIR /app
RUN apt-get update && \
    apt-get install --no-install-recommends -y \
    build-essential \
    && apt-get clean && rm -rf /tmp/* /var/tmp/*

COPY requirements.txt /app/requirements.txt    #1
RUN pip3 install --upgrade pip
RUN pip3 install --no-cache-dir -r requirements.txt
ENV PYTHONPATH "/app"
COPY . /app    #2
#1 把 requirement.txt 复制进镜像
#2 把根目录所有内容复制到工作目录

现在目录结构如下:

├── Dockerfile
└── src
    └── read_data
        └── read_data.py

我们把 read_data.py 放在 src/read_data 下,Dockerfile 在根目录。接下来用 docker build 构建镜像(把 <docker_id> 替换成你的 Docker ID):

docker build . -t <docker_id>/read-minio-data:latest

再用 docker push 推送到镜像仓库:

docker push <docker_id>/read-minio-data:latest

到这里,我们已经写好了 component 代码并完成容器化。下一步是定义 Kubeflow component 的规格文件(spec file)。这类似于定义一个 Kubernetes 对象,只不过对象是 Kubeflow component。和所有 Kubernetes 对象一样,我们用一个 YAML 文件来承载 component 定义。

读取数据组件的 component.yaml 以组件名称与描述开头。然后指定组件的 inputs 与 outputs 以及它们的数据类型。注意在 component.yaml 中,我们把 data_output 放在 outputs 下,并把类型设为 Dataset——这就是 Kubeflow 用来理解 data_output 将承载组件输出的方式。接着我们指定 Docker 镜像名以及读取数据所需执行的命令:这里命令会执行 read_data.py,如清单 5.4 所示。我们还提供执行 component 所需的参数,这些参数值来自预定义 inputs 与 outputs。

清单 5.4 读取组件定义文件

name: Read Data From MinIO
description: Fetches data from MinIO and outputs as a Dataset.
inputs:    #1
- name: minio_host
  type: String
  description: MinIO host URL.
- name: access_key
  type: String
  description: MinIO access key.
- name: secret_key
  type: String
  description: MinIO secret key.
- name: bucket_name
  type: String
  description: Name of the MinIO bucket.
- name: file_name
  type: String
  description: Name of the file to fetch from MinIO.
outputs:    #2
- name: data_output
  type: Dataset
  description: Output dataset artifact.
implementation:
  container:
    image: 'varunmallya/read-minio-data:latest'    #3
    command:
    - python3
    - /app/src/read_data/read_data.py    #4
    - --minio_host
    - {inputValue: minio_host}
    - --access_key
    - {inputValue: access_key}
    - --secret_key
    - {inputValue: secret_key}
    - --bucket_name
    - {inputValue: bucket_name}
    - --file_name
    - {inputValue: file_name}
    - --data_output_path
    - {outputPath: data_output}
#1 列出组件 inputs
#2 列出组件 outputs
#3 组件的 Docker 镜像
#4 执行组件的 Python 命令,并用 inputs/outputs 作为参数

现在我们已经定义了可用于 Kubeflow pipeline 的 component。为了构建这个 component,我们做了三件事:

  1. 用 Python 文件写清 component 的功能
  2. 用 Docker 将其容器化
  3. component.yaml 里写 component 规格

带上 component 规格后的目录结构如下:

├── Dockerfile
├── components
│   └── read_data
│       └── component.yaml
└── src
    └── read_data
        └── read_data.py

接下来我们会以同样方式定义流水线的其他 components:component 代码放在 src 下,component 定义 YAML 放在 components 下。下一个 component 是特征检索组件。

特征检索组件(FEATURE RETRIEVAL COMPONENT)

特征检索组件会从 feature store 取回特征。在第 4 章中,我们创建了一个 registry_local.db 文件,用来存储特征信息。由于我们在本地运行 feast apply 来注册 features 与 entities,当时把这些特征的 file source 指向了 localhost——尽管文件实际在 MinIO 中。因此,在运行下一步之前,我们需要在一个非本地的 registry 里重新注册这些特征,并把 file source 直接指向 MinIO。此时 s3_endpoint_override 会从 http://localhost:9000 变为 https://minio-service.kubeflow.svc.cluster.local:9000

我们必须这么做,因为我们没有对外公开的 MinIO 端点,且 MinIO service 注册为 cluster service,而不是 LoadBalancer 类型(第 3 章 3.3.4 节)。同时我们需要在 Kubernetes 集群内运行 feast apply。为此,我们会运行一个 Kubernetes job——它类似 Kubernetes deployment,但命令执行完就会结束。请按第 5 章代码仓库中的说明来完成特征注册。

完成特征注册后,我们开始编写 retrieve_features.py。我们通过从 MinIO 下载 feature_store.yaml 来初始化 feature store 对象。然后用 get_historical_features 拉取历史特征,该方法需要 entity data frame 和一个 features 列表;这两者都会作为 component inputs 传入,如清单 5.5 所示。entity_df 是 Dataset 类型的输入路径(稍后会在 component.yaml 里定义);在 Python 文件中,entity_df 的数据类型是 string。同样,data_output 是该组件的输出。

清单 5.5 特征检索组件

def get_features(
    minio_host: str,
    access_key: str,
    secret_key: str,
    bucket_name: str,
    file_name: str,
    entity_df: str,
    feature_list: str,
    data_output: str,
):
    store = init_feature_store(
        minio_host, access_key, secret_key, bucket_name, file_name
    )    #1
    print("Feature store initialized")
    feature_list = feature_list.split(",")    #2
    print("Requested features:", feature_list)
    print(entity_df)
    entity_df = pd.read_parquet(entity_df)    #3
    print("Entity DataFrame head:")
    print(entity_df.head())
    feature_df = store.get_historical_features(
        entity_df=entity_df,
        features=feature_list,
    ).to_df()
    print("Retrieved historical features:")
    print(feature_df.head())
    Path(data_output).parent.mkdir(parents=True, exist_ok=True)
    feature_df.to_parquet(data_output)    #4
#1 初始化 feature store
#2 获取 feature list
#3 从输入路径读取 entity DataFrame
#4 将 feature DataFrame 写到输出路径

接着我们把它容器化:构建镜像并推送到镜像仓库。Dockerfile 不需要修改,因为我们只是把新文件(retrieve_features.py)复制进工作目录。但我们需要在 requirements.txt 中加入 Feast 的依赖。镜像可以命名为 retrieve-feast-features

最后一步是写 component.yaml,如清单 5.6 所示。我们为组件命名并写描述,然后列出 inputs 与 outputs。在 container 部分指定镜像名与命令。entity_dfdata_output 的类型都是 Dataset,分别位于 inputs 与 outputs。

清单 5.6 特征检索组件定义文件

name: Retrieve Features From Feast

description: >
  Retrieves features from Feast where the
  feature_store.yaml is stored in MinIO
inputs:
- name: minio_host
  type: String
- name: access_key
  type: String
- name: secret_key
  type: String
- name: bucket_name
  type: String
- name: file_name
  type: String
- name: entity_df    #1
  type: Dataset
- name: feature_list
  type: String
outputs:
- name: data_output    #2
  type: Dataset

implementation:
  container:
    image: 'varunmallya/retrieve-feast-features:latest'
    command:
    - python3
    - /app/src/retrieve_features/retrieve_features.py
    - --minio_host
    - {inputValue: minio_host}
    - --access_key
    - {inputValue: access_key}
    - --secret_key
    - {inputValue: secret_key}
    - --bucket_name
    - {inputValue: bucket_name}
    - --file_name
    - {inputValue: file_name}
    - --entity_df    #3
    - {inputPath: entity_df}
    - --feature_list
    - {inputValue: feature_list}
    - --data_output    #4
    - {outputPath: data_output}
#1 entity_df 作为 Dataset 输入
#2 data_output 作为 Dataset 输出
#3 entity_df 通过 inputPath 传给参数 entity_df
#4 data_output 通过 outputPath 传给参数 data_output

这个组件会返回特征数据,接下来我们就具备了运行预测的条件。为此,我们需要下一个组件——推理(inference)。

推理组件(INFERENCE COMPONENT)

在检索到特征之后,下一步就是用收入分类模型生成预测。但模型在哪里?模型定义在 MLflow model registry 中。这个组件需要从 MLflow 拉取模型,并使用前一个组件得到的特征来生成预测。我们将把这段逻辑写在一个 Python 文件里(run_inference.py),如清单 5.7 所示。我们会初始化 MLflow client,并用它获取模型 URI;然后用该 URI 加载模型。根据模型所用框架(sklearn、XGBoost、TensorFlow 等),选择合适的 MLflow 方法来加载模型。模型加载后,我们调用 model.predict_proba 得到预测,并把预测写入一个新列。最后把整份文件写到输出路径。

清单 5.7 推理组件

def perform_inference(
    minio_host: str,
    access_key: str,
    secret_key: str,
    model_name: str,
    model_type: str,
    model_stage: str,
    mlflow_host: str,
    input_data: str,  
    data_output: str,
):
    os.environ["AWS_ACCESS_KEY_ID"] = access_key
    os.environ["AWS_SECRET_ACCESS_KEY"] = secret_key
    if not minio_host.startswith("http"):
        os.environ["AWS_ENDPOINT_URL"] = "http://" + minio_host
    else:
        os.environ["AWS_ENDPOINT_URL"] = minio_host
    mlflow.set_tracking_uri(mlflow_host)
    mlflow_client = MlflowClient(mlflow_host)    #1
    model_run_id = None

    for model in mlflow_client.search_model_versions(
        f"name='{model_name}'"
    ):  #2
        if model.current_stage == model_stage:
            model_run_id = model.run_id
            break

    if not model_run_id:
        raise ValueError(
            f"No model found in stage {model_stage} for model {model_name}."
        )

    mlflow.artifacts.download_artifacts(
        f"runs:/{model_run_id}/column_list.pkl", dst_path="column_list"
    )
    input_data_df = pd.read_parquet(input_data)    #3
    input_data_df.drop(columns=["user_id", "event_timestamp"], inplace=True)

    with open("column_list/column_list.pkl", "rb") as f:
        col_list = pickle.load(f)
    input_data_df = pd.get_dummies(
        input_data_df, drop_first=True, sparse=False, dtype=float
    )
    input_data_df = input_data_df.reindex(columns=col_list, fill_value=0)

    if model_type == "sklearn":    #4
        model = mlflow.sklearn.load_model(
            model_uri=f"models:/{model_name}/{model_stage}"
        )
    elif model_type == "xgboost":
        model = mlflow.xgboost.load_model(
            model_uri=f"models:/{model_name}/{model_stage}"
        )
    elif model_type == "tensorflow":
        model = mlflow.tensorflow.load_model(
            model_uri=f"models:/{model_name}/{model_stage}"
        )
    else:
        raise NotImplementedError(
        f"Model type '{model_type}' is not supported.")

    predicted_classes = [
        x[1]
        for x in model.predict_proba(
            input_data_df
        )
    ]  #5
    input_data_df["Predicted_Income_Class"] = predicted_classes
    Path(data_output).parent.mkdir(parents=True, exist_ok=True)
    input_data_df.to_parquet(data_output)    #6
#1 初始化 MLflow client
#2 通过模型名搜索 MLflow 模型版本
#3 读取包含 entities 与 features 的推理数据集
#4 根据建模框架加载 MLflow 模型
#5 生成模型预测
#6 把预测写到输出路径

我们将这个 Docker 镜像命名为 run-inference,构建后推送到存放该代码的镜像仓库。接下来就可以定义该组件的 component.yaml,并写上合适的 inputs 与 outputs,如清单 5.8 所示。再次注意输入/输出路径的数据类型:在定义组件 inputs 与 outputs 时,路径会是 Output[Dataset] 的类型——也就是说 input_datadata_output 都是 Dataset 类型。但在 command 中引用这些输入/输出时,必须分别使用 inputPath--input_data, {inputPath: input_data})与 outputPath--data_output, {outputPath: data_output})来传递输入与输出数据。

清单 5.8 推理组件定义文件

name: Model Inference
description: Run model inference after retrieving the model from MLflow
inputs:
- name: minio_host
  type: String
- name: access_key
  type: String
- name: secret_key
  type: String
- name: model_name
  type: STRING
- name: model_type
  type: STRING
- name: model_stage
  type: STRING
- name: mlflow_host
  type: STRING
- name: input_data
  type: Dataset
outputs:
- name: data_output
  type: Dataset
implementation:
  container:
    image: 'varunmallya/run-inference:latest'
    command:
    - python3
    - /app/src/run_inference/run_inference.py
    - --minio_host
    - {inputValue: minio_host}
    - --access_key
    - {inputValue: access_key}
    - --secret_key
    - {inputValue: secret_key}
    - --model_name
    - {inputValue: model_name}
    - --model_type
    - {inputValue: model_type}
    - --model_stage
    - {inputValue: model_stage}
    - --mlflow_host
    - {inputValue: mlflow_host}
    - --input_data
    - {inputPath: input_data}
    - --data_output
    - {outputPath: data_output}

我们的预测最终会通过“写数据组件(write data component)”写回 MinIO。

写数据组件(WRITE DATA COMPONENT)

与读数据组件类似,我们会初始化 MinIO client,并把一个 Parquet 文件写到 MinIO bucket。代码写在 write_data.py 中,如下所示。

清单 5.9 写数据组件

def write_data(
    minio_host: str,
    access_key: str,
    secret_key: str,
    bucket_name: str,
    file_name: str,
    input_data_path: str,
):

    client = Minio(
        endpoint=minio_host,
        access_key=access_key,
        secret_key=secret_key,
        secure=False,
    )
    input_data = pd.read_parquet(
        input_data_path
    ) 
    input_data.to_parquet(file_name, index=False)
    client.fput_object(bucket_name, file_name, file_name)

我们还需要写一个组件定义文件,其中组件输入来自推理组件的输出,如下所示。

清单 5.10 写组件定义文件

name: Write Data To Minio
description: Writes data back to MinIO
inputs:
  - name: minio_host
    type: STRING
    description: The MinIO host address.
  - name: access_key
    type: STRING
    description: Access key for MinIO.
  - name: secret_key
    type: STRING
    description: Secret key for MinIO.
  - name: bucket_name
    type: STRING
    description: Name of the bucket in MinIO.
  - name: file_name
    type: STRING
    description: Name of the file to write to MinIO.
  - name: input_data
    type: Artifact
    description: Path to the input data file.

implementation:
  container:
    image: 'varunmallya/write-minio-data:latest'
    command:
      - python3
      - /app/src/write_data/write_data.py
    args:
      - --minio_host
      - {inputValue: minio_host}
      - --access_key
      - {inputValue: access_key}
      - --secret_key
      - {inputValue: secret_key}
      - --bucket_name
      - {inputValue: bucket_name}
      - --file_name
      - {inputValue: file_name}
      - --input_data_path
      - {inputPath: input_data}

现在我们已经定义了基础推理流水线所需的全部 components。完整目录结构如下:

── Dockerfile
├── components
│   ├── read_data
│   │   └── component.yaml
│   ├── retrieve_features
│   │   └── component.yaml
│   ├── run_inference
│   │   └── component.yaml
│   └── write_data
│       └── component.yaml
├── requirements.txt
└── src
    ├── read_data
    │   └── read_data.py
    ├── retrieve_features
    │   └── retrieve_features.py
    ├── run_inference
    │   └── run_inference.py
    └── write_data
        └── write_data.py

流水线与 MinIO、Feast、MLflow 的交互如图 5.2 所示。由于这些 components 足够通用,可以在多条流水线中复用,从而帮助我们在团队与组织内标准化 ML 工作流。

image.png

图 5.2 Kubeflow pipeline 与 MinIO、Feast、MLflow 的交互。所有中间数据集都存储在 MinIO 中,以便后续步骤取用。模型与特征分别从 MinIO 与 Feast 获取。

在设计 components 时,请记住以下几点:

  • 模块化(Modularity) ——尽量让组件遵循单一职责原则:一个组件只做一件事。这促进代码复用,并让维护与更新单个组件更容易。在我们的例子里,我们把推理任务与写数据任务分开。
  • 组件输入与输出(Component inputs and outputs) ——每个组件的输入与输出都应详细指定,定义组件所需的输入以及会产生的输出数据或 artifacts。我们在每个组件的 component.yaml 中指定了 inputs 与 outputs。
  • 参数化(Parameterization) ——允许组件参数化。参数可以在不改代码的情况下定制组件行为,尤其适用于超参数、文件路径与配置项。在我们的例子里,我们把 mlflow_hostbucket_nameminio_host 设为参数,而不是写死,这样也能在其他环境复用。
  • 文档(Documentation) ——清晰描述每个组件的功能、输入、输出与用法。在文档中加入示意图与使用说明,帮助团队其他成员理解并高效使用组件。

良好设计的组件能在多个方面帮助我们:

  • 可复用性(Reusability) ——组件定义完成后可被不同流水线复用,减少重复代码与重复劳动,从而提升开发效率。团队可以维护一套标准化组件库,供多个项目使用。
  • 协作(Collaboration) ——组件促进数据科学家、ML 工程师与领域专家之间的协作。不同成员可以并行开发不同组件,并在不同 ML 领域分工专长。
  • 测试(Testing) ——组件可以单独测试,使调试更高效。把问题隔离到特定组件,会降低排查整条流水线时的复杂度,更容易定位并修复缺陷。

下一节我们会把所有组件整合为一条 Kubeflow pipeline。

5.1.2 收入分类器流水线(Income classifier pipeline)

Kubeflow pipeline 用于描述完整的 ML 工作流。组件会以有向无环图(DAG)的形式表示。DAG 有有向边,意味着边从一个节点指向另一个节点;并且没有环(cycle),也就是说图中没有闭环。我们可以指定任务执行顺序,并在需要时强制任务之间的依赖关系——也就是只有当某个任务依赖的任务完成后,它才会运行。我们将使用前面定义的 components 来构建一条 Kubeflow pipeline。

和 components 一样,Kubeflow pipeline 也是用一个 Python 文件定义的,我们把它命名为 income-classifier-pipeline.py。我们首先用 kfp 提供的 load_component_from_file 函数来定义所有 components(见清单 5.11),该函数会创建 component 对象。它需要传入 component 文件路径。这样我们就得到四个 component:fetch_data_opretrieve_features_oprun_inference_opwrite_data_op

清单 5.11 收入分类器 Kubeflow pipeline

import kfp
from kfp import dsl, compiler
fetch_data_op = kfp.components.load_component_from_file(
    "components/read_data/component.yaml"
)
retrieve_features_op = kfp.components.load_component_from_file(
    "components/retrieve_features/component.yaml"
)
run_inference_op = kfp.components.load_component_from_file(
    "components/run_inference/component.yaml"
)

write_data_op = kfp.components.load_component_from_file(
    "components/write_data/component.yaml"
)

定义好 components 后,我们定义一个名为 income_classifier_pipeline 的 pipeline 函数。该函数的参数就是 pipeline 的运行时参数(pipeline parameters),它们的值可以在运行时修改。我们用 @dsl.pipeline 装饰器标注该函数是 pipeline 定义:

@dsl.pipeline(
    name="income-classifier-pipeline",

    description=(
      "A Kubeflow pipeline to classify income categories "
      "using KFP v2"
    ),
    )
def income_classifier_pipeline(
    minio_host: str,
    access_key: str,
    secret_key: str,
    entity_df_bucket: str,
    entity_df_filename: str,
    feature_store_bucket_name: str,
    feature_store_config_file_name: str,
    feature_list: str,
    model_name: str,
    model_type: str,
    model_stage: str,
    mlflow_host: str,
    output_bucket: str,
    output_file_name: str,
):

接下来我们用预定义的 components 来定义任务:本质上就是给 component 的 input 字段赋值,而这些 input 值来自 pipeline parameters。比如,读取数据组件需要 MinIO 的连接信息以及 entity DataFrame 的位置,这些都通过 pipeline parameters 传入。四个 components 会分别对应四个 tasks:fetch_data_taskretrieve_features_taskrun_inference_taskwrite_data_task

但组件之间怎么连起来?fetch_data 的输出如何变成 retrieve_features 的输入?kfp 包提供了一个简单方式:使用 task.outputs 获取 task 的输出,然后把它喂给下游输入。例如,retrieve_features_task 有一个输入叫 entity_df,它就是 fetch_data_task 的输出(即 fetch_data_task.outputs["data_output"])。这样 Kubeflow 就能理解:retrieve_features_task 依赖 fetch_data_task 完成。同样,run_inference_task 依赖 retrieve_features_task 完成,其中 input_data 来自 retrieve_features_task 的输出。

fetch_data_task = fetch_data_op(
    minio_host=minio_host,
    access_key=access_key,
    secret_key=secret_key,
    bucket_name=entity_df_bucket,
    file_name=entity_df_filename,
)
retrieve_features_task = retrieve_features_op(
    minio_host=minio_host,
    access_key=access_key,
    secret_key=secret_key,
    bucket_name=feature_store_bucket_name,
    file_name=feature_store_config_file_name,

    entity_df=fetch_data_task.outputs[
        "data_output"
    ],
    feature_list=feature_list,
)
run_inference_task = run_inference_op(
    minio_host=minio_host,
    access_key=access_key,
    secret_key=secret_key,
    model_name=model_name,
    model_type=model_type,
    model_stage=model_stage,
    mlflow_host=mlflow_host,
    input_data=retrieve_features_task.outputs["data_output"],
)
write_data_task = write_data_op(
    minio_host=minio_host,
    access_key=access_key,
    secret_key=secret_key,
    bucket_name=output_bucket,
    file_name=output_file_name,
    input_data=run_inference_task.outputs["data_output"],
)

我们也可以在参数定义时为 pipeline parameters 指定默认值。pipeline 函数写好之后,需要把它编译成一个 YAML 文件,交给 Kubeflow 使用。kfp 包提供了 compiler,会把 pipeline 函数翻译成一个名为 income_classifier_pipeline.yaml 的 YAML 文件:

Compiler().compile(
    income_classifier_pipeline,
    "income_classifier_pipeline.yaml",
)

然后就可以把 income_classifier_pipeline.yaml 上传到 KFP UI(图 5.3)。操作是:登录 Kubeflow,在侧边栏点击 Pipelines,然后点击 Upload pipeline 按钮。

image.png

图 5.3 KFP UI:列出账号下所有 pipelines

接着填写 pipeline 详情,包括 Pipeline Name、Pipeline Description,然后提供已编译的 income_classifier_pipeline.yaml 文件路径。点击 Create 创建 pipeline(图 5.4)。

image.png

图 5.4 KFP UI:输入 pipeline 名称并上传 pipeline YAML 文件来创建 pipeline

创建后会得到 pipeline 可视化图(图 5.5)。要运行该 pipeline,需要按以下步骤:

image.png

图 5.5 上传后由 KFP UI 可视化的 pipeline

  1. 创建一个 experiment。
  2. 创建一个 run。

Kubeflow pipeline experiment 用于管理与跟踪 ML 工作流的运行,便于组织、监控并对比不同的 pipeline runs。Kubeflow run 则是一次 pipeline 执行实例。通常建议把某条 pipeline 的多次 runs 归到同一个 experiment 下,便于跟踪。首先我们点击 Create Experiment 按钮创建 experiment(图 5.6),填写 Experiment name 与(可选)Description。experiment 创建完成后,点击 Create Run 创建一次 pipeline run。在这里我们需要填写 pipeline 的运行时参数,并选择 experiment 名称(图 5.7)。然后点击 Start Run 启动运行。

image.png

图 5.6 在 Kubeflow 中创建 experiment:填写 experiment 名称与描述。

image.png

图 5.7 创建一次 pipeline run:填写所有 pipeline 参数。

点击 Start Run 后,就能在 KFP UI 里看到 run 的执行进度(图 5.8)。

image.png

图 5.8 一次成功的 Kubeflow pipeline run

我们成功构建并运行了收入分类器 pipeline,并且可以在 MinIO 的 inference-datasets bucket 中找到预测结果。KFP 是在 Kubernetes 环境中管理与执行 ML 工作流的优秀工具。它为复杂 ML pipelines 的设计、调度与跟踪提供结构化、组织化的方式,从而提升数据科学与 ML 项目的可复现性、可扩展性与协作效率。通过把工作流定义为代码,KFP 让数据科学家与 ML 工程师能自动化、串联并编排从数据预处理到模型部署的各类任务。团队也能在 experiments 中系统地尝试不同配置、超参数与数据集,并在其中跟踪大量 runs。

KFP 非常适合批处理推理。至于实时推理与生产部署模型,我们会在下一章介绍另一个工具——BentoML。

总结(Summary)

编排(Orchestration)是指在运行多条复杂流水线的同时,确保流水线各步骤具备所需的数据与计算资源的过程。Kubeflow 平台提供了 Kubeflow Pipelines(KFP)解决方案。

组件(component)是 ML 工作流中可复用、可组合的工作单元。这些组件封装单个任务或步骤,执行特定动作,例如数据预处理、模型训练、评估、部署等。

在设计组件时,应当注意尽可能让组件保持独立与通用,从而提升可复用性,并使单个流水线阶段的开发、测试与维护更容易。

更大的工作流可以拆解为更小的组件。该方法有助于用共享组件构建模块化、易维护的 ML 流水线。

KFP 与任务编排器可用于自动化、串联并管理 ML 工作流,确保数据预处理、模型训练与部署任务被高效执行。它提供 SDK,用于构建单个组件、把组件连接成流水线,以及完成流水线的部署与执行。