容器化的 Python 组件

以下假设您对 Lightweight Python Components 有基本的了解。

容器化的 Python 组件通过放宽轻量级 Python 组件必须是封闭的（即完全自包含）的限制，扩展了轻量级 Python 组件。这意味着容器化的 Python 组件函数可以依赖于函数外部定义的符号、函数外部的导入、相邻 Python 模块中的代码等。为了实现这一点，KFP SDK 提供了一种便捷的方式将您的 Python 代码打包到一个容器中。

作为生产软件的最佳实践，当组件作者指定 packages_to_install 时，应该优先选择容器化的 Python 组件而不是 轻量级 Python 组件，因为 KFP SDK 会在构建组件的镜像时安装这些依赖，而不是在任务运行时。

以下展示了如何通过修改来自轻量级Python组件示例的add组件来使用容器化Python组件：

from kfp import dsl

@dsl.component
def add(a: int, b: int) -> int:
    return a + b

1. 源代码设置

首先创建一个空的 src/ 目录来包含你的源代码：

src/

接下来，添加以下简单模块，src/math_utils.py，并包含一个辅助函数：

# src/math_utils.py
def add_numbers(num1, num2):
    return num1 + num2

最后，将您的组件移动到 src/my_component.py 并修改它以使用辅助函数：

# src/my_component.py
from kfp import dsl

@dsl.component
def add(a: int, b: int) -> int:
    from math_utils import add_numbers
    return add_numbers(a, b)

src 现在看起来像这样：

src/
├── my_component.py
└── math_utils.py

2. 修改 dsl.component 装饰器

在这一步中，您将提供 base_image 和 target_image 参数给您在 src/my_component.py 中的 @dsl.component 装饰器：

@dsl.component(base_image='python:3.11',
               target_image='gcr.io/my-project/my-component:v1')
def add(a: int, b: int) -> int:
    from math_utils import add_numbers
    return add_numbers(a, b)

设置 target_image 同时 (a) 指定了您将在第 3 步中构建的图像的 tag，以及 (b) 指示 KFP 在使用该标签的图像的容器中运行被装饰的 Python 函数。

在一个容器化的 Python 组件中，base_image 指定了 KFP 在构建新的容器镜像时将使用的基础镜像。具体来说，KFP 在构建镜像时使用 base_image 参数作为 Dockerfile 中的 FROM 指令。

前面的例子包含了 base_image 以提高清晰度，但这不是必需的，因为如果省略 base_image，它将默认为 'python:3.11'。

3. 构建组件

现在您的代码在一个独立的目录中，并且您已经指定了目标镜像，您可以方便地使用kfp component build CLI命令构建一个镜像：

kfp component build src/ --component-filepattern my_component.py --no-push-image

如果您有一个 配置的 Docker 来使用私有镜像仓库，您可以将 --no-push-image 标志替换为 --push-image，以在构建后自动推送镜像。

4. 在管道中使用组件

最后，你可以在管道中使用该组件：

# pipeline.py
from kfp import compiler, dsl
from src.my_component import add

@dsl.pipeline
def addition_pipeline(x: int, y: int) -> int:
    task1 = add(a=x, b=y)
    task2 = add(a=task1.output, b=x)
    return task2.output

compiler.Compiler().compile(addition_pipeline, 'pipeline.yaml')

你的目录现在看起来像这样：

pipeline.py
src/
├── my_component.py
└── math_utils.py

由于 add 的 target_image 使用了 Google Cloud Artifact Registry（由 gcr.io URI 指示），此处所示的管道假定您已将图像推送到 Google Cloud Artifact Registry，您正在 Google Cloud Vertex AI Pipelines 上运行管道，并且您已配置 IAM 权限，以便 Vertex AI Pipelines 可以从 Artifact Registry 中拉取图像。