如何修改KFP输出Dataset artifact的URI后缀为.csv？

问题场景 #

当前使用KubeFlow的Output[Dataset]输出CSV格式数据时，生成的URI为gs://<bucket-name>/<id>/<pipeline-name>+<id>/<script-name>+<id>/dataset，结尾无.csv后缀。手动在dataset.uri后追加.csv会导致后续组件无法直接通过train_data.path引用数据，必须手动拼接后缀，希望利用KFP原生能力解决该问题。

解决方案 #

方案1：修改Output[Dataset]对象的URI和Path属性 #

直接在组件函数中修改Dataset对象的uri和path属性，让KFP后续引用时直接获取带后缀的路径：

def extract_bq_to_dataset(
    bq_client_project_id: str,
    source_project_id: str,
    dataset_id: str,
    table_name: str,
    dataset: Output[Dataset],
    dataset_location: str = "EU",
    extract_job_config: dict = None,
    ):
    from google.cloud import bigquery
    from google.cloud.exceptions import GoogleCloudError

    # 自定义带.csv后缀的URI和Path
    csv_uri = dataset.uri + ".csv"
    csv_path = dataset.path + ".csv"
    # 更新Dataset对象的属性
    dataset.uri = csv_uri
    dataset.path = csv_path
    # 可选：自定义文件名（如training_data.csv）
    # dataset.name = "training_data.csv"

    full_table_id = f"{source_project_id}.{dataset_id}.{table_name}"
    table = bigquery.table.Table(table_ref=full_table_id)

    if extract_job_config is None:
        extract_job_config = {}
    job_config = bigquery.job.ExtractJobConfig(**extract_job_config)

    client = bigquery.Client(project=bq_client_project_id, location=dataset_location)
    extract_job = client.extract_table(
        source=table,
        destination_uris=csv_uri,
        job_config=job_config,
        location=dataset_location,
    )

    try:
        result = extract_job.result()
    except GoogleCloudError as e:
        raise e

修改后，后续组件引用该Dataset时，train_data.path和train_data.uri都会直接返回带.csv后缀的路径，无需手动拼接。

方案2：组件定义时指定输出文件名 #

使用KFP的@dsl.component装饰器，通过OutputPath的file_name参数直接指定输出文件名：

@dsl.component(
    base_image="python:3.9",
    packages_to_install=["google-cloud-bigquery"],
    outputs={
        "dataset": dsl.OutputPath(type=Dataset, file_name="training_data.csv")
    }
)
def extract_bq_to_dataset(
    bq_client_project_id: str,
    source_project_id: str,
    dataset_id: str,
    table_name: str,
    dataset_location: str = "EU",
    extract_job_config: dict = None,
    ):
    from google.cloud import bigquery
    from google.cloud.exceptions import GoogleCloudError

    # 获取KFP生成的带指定文件名的路径
    dataset_path = dsl.get_output_path("dataset")
    csv_uri = f"gs://{dataset_path.lstrip('/')}"

    full_table_id = f"{source_project_id}.{dataset_id}.{table_name}"
    table = bigquery.table.Table(table_ref=full_table_id)

    if extract_job_config is None:
        extract_job_config = {}
    job_config = bigquery.job.ExtractJobConfig(**extract_job_config)

    client = bigquery.Client(project=bq_client_project_id, location=dataset_location)
    extract_job = client.extract_table(
        source=table,
        destination_uris=csv_uri,
        job_config=job_config,
        location=dataset_location,
    )

    try:
        result = extract_job.result()
    except GoogleCloudError as e:
        raise e

这种方式下，KFP会自动生成包含指定文件名（如training_data.csv）的URI和路径，后续组件直接引用输出即可。

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