如何在Vertex AI Pipeline中复用Kubeflow v1的Docker镜像组件？

问题 #

我正尝试将Kubeflow中创建的自定义组件迁移至Vertex AI。在Kubeflow中，我曾将组件创建为Docker容器镜像，并通过以下方式加载到流水线中：

def my_custom_component_op(gcs_dataset_path: str, some_param: str):
    return kfp.dsl.ContainerOp(
        name='My Custom Component Step',
        image='gcr.io/my-project-23r2/my-custom-component:latest',
        arguments=["--gcs_dataset_path", gcs_dataset_path,
                   '--component_param', some_param],
        file_outputs={
            'output': '/app/output.csv',
        }
    )

随后我会在流水线中这样使用：

@kfp.dsl.pipeline(
    name='My custom pipeline',
    description='The custom pipeline'
)
def generic_pipeline(project_id, some_param):
    
    output_component = my_custom_component_op(
        gcs_dataset_path=gcs_dataset_path,
        some_param=some_param
    )

    output_next_op = next_op(gcs_dataset_path=dsl.InputArgumentPath(
        output_component.outputs['output']),
        next_op_param="some other param"
    )

我能否在Vertex AI Pipeline中复用Kubeflow v1的同一组件Docker镜像？该如何操作？希望无需对组件本身做任何修改。

我在网上找到的Vertex AI Pipeline示例均使用@component装饰器，如下所示：

@component(base_image=PYTHON37, packages_to_install=[PANDAS])
def my_component_op(
    gcs_dataset_path: str,
    some_param: str
    dataset: Output[Dataset],
):
   ...perform some op....

但这种方式需要我将Docker中的代码复制粘贴到流水线中，这并非我想要的。是否存在复用Docker镜像并传递参数的方法？我未能找到相关示例。

解决方案 #

完全可以复用Kubeflow v1的Docker镜像在Vertex AI Pipelines中，无需修改组件本身。以下是两种可行的实现方式：

方法1：适配KFP v2的ContainerOp调用 #

Vertex AI Pipelines基于Kubeflow Pipelines (KFP) v2，调整原有ContainerOp代码即可适配：

from kfp import dsl

def my_custom_component_op(gcs_dataset_path: str, some_param: str):
    return dsl.ContainerOp(
        name='My Custom Component Step',
        image='gcr.io/my-project-23r2/my-custom-component:latest',
        arguments=[
            "--gcs_dataset_path", gcs_dataset_path,
            "--component_param", some_param
        ],
        # KFP v2用outputs替代v1的file_outputs，需明确输出路径与名称
        outputs=[
            dsl.OutputPath(path='/app/output.csv', name='output')
        ]
    )

流水线使用方式与v1基本一致，仅需调整输入引用写法：

from kfp import dsl

@dsl.pipeline(
    name='My custom pipeline',
    description='The custom pipeline'
)
def generic_pipeline(project_id: str, some_param: str, gcs_dataset_path: str):
    
    output_component = my_custom_component_op(
        gcs_dataset_path=gcs_dataset_path,
        some_param=some_param
    )

    # KFP v2中直接引用组件输出，无需InputArgumentPath
    output_next_op = next_op(
        gcs_dataset_path=output_component.outputs['output'],
        next_op_param="some other param"
    )

方法2：使用load_component_from_container_image（推荐） #

这是KFP v2官方推荐的加载现有容器镜像的方式，会自动生成符合v2规范的组件函数：

from kfp.components import load_component_from_container_image

# 加载容器镜像生成组件函数
my_custom_component_op = load_component_from_container_image(
    image='gcr.io/my-project-23r2/my-custom-component:latest',
    # 定义组件输入输出结构，需与镜像内参数/输出路径匹配
    input_component_schema={
        "inputs": [
            {"name": "gcs_dataset_path", "type": "String"},
            {"name": "some_param", "type": "String"}
        ],
        "outputs": [
            {"name": "output", "type": "String"}
        ]
    },
    # 映射命令行参数
    arguments=[
        "--gcs_dataset_path", "{{inputs.parameters.gcs_dataset_path}}",
        "--component_param", "{{inputs.parameters.some_param}}"
    ],
    # 映射输出文件路径
    file_outputs={
        "output": "/app/output.csv"
    }
)

流水线使用方式更简洁：

from kfp import dsl

@dsl.pipeline(
    name='My custom pipeline',
    description='The custom pipeline'
)
def generic_pipeline(project_id: str, some_param: str, gcs_dataset_path: str):
    
    output_component = my_custom_component_op(
        gcs_dataset_path=gcs_dataset_path,
        some_param=some_param
    )

    output_next_op = next_op(
        gcs_dataset_path=output_component.outputs['output'],
        next_op_param="some other param"
    )

关键注意事项 #

• 确保Docker镜像已推送到Vertex AI可访问的容器仓库（如GCR、Artifact Registry），且Vertex AI服务账号拥有该镜像的读取权限。
• 两种方法均无需修改原有Docker镜像内的代码，仅需调整流水线侧的调用逻辑。
• 若组件输出为文件（如示例中的output.csv），KFP v2会自动将文件路径以字符串形式传递给下游组件，与v1行为一致。

内容的提问来源于stack exchange，提问作者DarioB