如何在Vertex AI Pipeline组件间传递对象URI至预构建组件

核心要点 #

要解决自定义组件输出无法传递给预构建组件的问题，关键要做好两点：正确定义自定义组件的输出类型，以及按照预构建组件的参数要求传递数据。

1. 规范自定义组件的输出定义 #

以save_gcs_file_paths_to_csv组件为例，需要明确输出是GCS上的CSV文件URI，可通过两种方式定义输出：

方式一：使用Output[Artifact]自动管理URI

这种方式下，KFP会自动分配GCS存储路径，组件只需将生成的CSV文件上传到指定URI即可，同时返回该URI供后续组件使用：

from kfp.v2 import dsl
from kfp.v2.dsl import component, Output, Artifact
import csv
from google.cloud import storage

@component(base_image='python:3.9', packages_to_install=['google-cloud-storage'])
def save_gcs_file_paths_to_csv(
    image_paths: list[str],
    output_file: Output[Artifact]
) -> str:
    # 生成本地CSV文件
    with open('/tmp/file_paths.csv', 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['GCS_PATH'])
        for path in image_paths:
            writer.writerow([path])
    
    # 上传到KFP分配的GCS URI
    client = storage.Client()
    bucket_name, blob_path = output_file.uri.replace('gs://', '').split('/', 1)
    bucket = client.bucket(bucket_name)
    blob = bucket.blob(blob_path)
    blob.upload_from_filename('/tmp/file_paths.csv')
    
    return output_file.uri

方式二：直接返回字符串类型的GCS URI

如果需要手动指定GCS路径，可直接将输出定义为字符串类型：

@component(base_image='python:3.9', packages_to_install=['google-cloud-storage'])
def save_gcs_file_paths_to_csv(
    image_paths: list[str],
    gcs_output_path: str
) -> str:
    with open('/tmp/file_paths.csv', 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['GCS_PATH'])
        for path in image_paths:
            writer.writerow([path])
    
    client = storage.Client()
    bucket_name, blob_path = gcs_output_path.replace('gs://', '').split('/', 1)
    bucket = client.bucket(bucket_name)
    blob = bucket.blob(blob_path)
    blob.upload_from_filename('/tmp/file_paths.csv')
    
    return gcs_output_path

2. 按照预构建组件的参数要求传递数据 #

ImageDatasetCreateOp的gcs_source参数要求传入字符串列表（即使只有一个CSV文件），因此需要将自定义组件的输出包装成列表传递：

from google_cloud_pipeline_components.v1.dataset import ImageDatasetCreateOp

@dsl.pipeline(
    name='image-dataset-pipeline',
    pipeline_root='gs://your-pipeline-root-bucket'
)
def pipeline(project_id: str, region: str):
    # 调用自定义组件生成CSV文件
    save_csv_task = save_gcs_file_paths_to_csv(
        image_paths=['gs://your-bucket/image1.jpg', 'gs://your-bucket/image2.jpg']
        # 如果是方式二，需传入gcs_output_path='gs://your-bucket/output/file_paths.csv'
    )
    
    # 调用预构建组件，将CSV URI包装为列表传入gcs_source
    create_dataset = ImageDatasetCreateOp(
        project=project_id,
        location=region,
        display_name='custom-image-dataset',
        gcs_source=[save_csv_task.output]
    )

额外注意事项 #

• 确保自定义组件使用的服务账号拥有目标GCS桶的读写权限，以及Vertex AI的数据集创建权限。
• 如果自定义组件输出的URI格式不正确（比如不是gs://开头），预构建组件会报错，需在自定义组件中验证输出路径的有效性。

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