Run ID in Kubeflow Pipelines on Vertex AI

Viewed 739

I am trying to run Kubeflow Pipelines with the new Vertex AI on GCP.

Previously, in Kubeflow Pipelines, I was able to use the Run ID in my Pipeline by utilizing dsl.RUN_ID_PLACEHOLDER or {{workflow.uid}}. My understanding was that dsl.RUN_ID_PLACEHOLDER would resolve to {{workflow,uid}} at compile time, and then at run time, the {{workflow.uid}} tag would be resolved to the Run's ID. This is at least how it has worked in my experience using Kubeflow Pipelines and the Kubeflow Pipelines UI.

However, when I try to access the Run ID in a similar way in a pipeline that I run in Vertex AI Pipelines, it seems that dsl.RUN_ID_PLACEHOLDER resolves to {{workflow.uid}} but that this never subsequently resolves to the ID of the run.

I created the following Test Pipeline, which tries to get the Run ID using the DSL Placeholder, then uses a lightweight component to print out the value of the run_id parameter of the pipeline. The result of running the pipeline in the UI is that the print_run_id component prints {{workflow.uid}}, where as on Kubeflow Pipelines previously, it would have resolved to the Run ID.

from kfp import dsl
from kfp import components as comp
import logging
from kfp.v2.dsl import (
    component,
    Input,
    Output,
    Dataset,
    Metrics,
)

@component
def print_run_id(run_id:str):
    print(run_id)

RUN_ID = dsl.RUN_ID_PLACEHOLDER

@dsl.pipeline(
    name='end-to-end-pipeline',
    description='End to end XGBoost cover type training pipeline'
)
def end_to_end_pipeline(
    run_id: str = RUN_ID
):
    print_task = print_run_id(run_id=run_id)

Is there a way to access the Run ID using the KFP SDK with Vertex AI Pipelines?

2 Answers

This isn't documented well. Trying to access them during the pipeline build or from within individuals components directly will return the unmodified placeholder string values. But it does work during the actual pipeline run (at least for kfp v2)

Example for kfp v2 based on this link that works as expected.

import kfp.v2.dsl as dsl

@component
def print_op(msg: str, value: str):
    print(msg, value) # <-- Prints the correct value

@component
def incorrect_print_op():
    print(dsl.PIPELINE_JOB_NAME_PLACEHOLDER) # <-- Prints incorrect placeholder value

@dsl.pipeline(name='pipeline-with-placeholders')
def my_pipeline():

    # Correct
    print_op(msg='job name:', value=dsl.PIPELINE_JOB_NAME_PLACEHOLDER)
    print_op(msg='job resource name:', value=dsl.PIPELINE_JOB_RESOURCE_NAME_PLACEHOLDER)
    print_op(msg='job id:', value=dsl.PIPELINE_JOB_ID_PLACEHOLDER)
    print_op(msg='task name:', value=dsl.PIPELINE_TASK_NAME_PLACEHOLDER)
    print_op(msg='task id:', value=dsl.PIPELINE_TASK_ID_PLACEHOLDER)
    
    # Incorrect
    print(dsl.PIPELINE_TASK_ID_PLACEHOLDER) # Prints only placeholder value

Related