Build your own pipeline components

It's common that when running a component, you want to not only see the link to the component job being launched, but also the link to the underlying cloud resources, such as the Vertex batch prediction jobs or dataflow jobs.

The gcp_resource proto is a special parameter that you can use in your component to enable the Google Cloud console to provide a customized view of the resource's logs and status in the Vertex AI Pipelines console.

Output the gcp_resource parameter

Via container-based component

First, you'll need to define the gcp_resource parameter in your component as shown in the following example component.yaml file:

name: dataflow_python
inputs:
...
outputs:
- {name: gcp_resources, type: String}
implementation:
  container:
    image: gcr.io/ml-pipeline/google-cloud-pipeline-components:latest
    command: [python3, -u, -m, google_cloud_pipeline_components.container.experimental.dataflow.dataflow_launcher]
    args: [
      --project, {inputValue: project},
      --location, {inputValue: location},
      --python_module_path, {inputValue: python_module_path},
      --temp_location, {inputValue: temp_location},
      --requirements_file_path, {inputValue: requirements_file_path},
      --args, {inputValue: args},
      --gcp_resources, {outputPath: gcp_resources},
    ]

Next, inside the container, install the Google Cloud Pipeline Components package:

pip install -upgrade google-cloud-pipeline-components

Then in the Python code, write out the resource as a gcp_resource parameter:

from google_cloud_pipeline_components.experimental.proto.gcp_resources_pb2 import GcpResources
from google.protobuf.json_format import MessageToJson

dataflow_resources = GcpResources()
dr = dataflow_resources.resources.add()
dr.resource_type='DataflowJob'
dr.resource_uri='https://dataflow.googleapis.com/v1b3/projects/[your-project]/locations/us-east1/jobs/[dataflow-job-id]'

with open(gcp_resources, 'w') as f:
    f.write(MessageToJson(dataflow_resources))

Via Python component

Alternatively, you can return the gcp_resources output parameter as you would any string output parameter:

@dsl.component(
    base_image='python:3.9',
    packages_to_install=['google-cloud-pipeline-components==1.0.7'],
)
def launch_dataflow_component(project: str, location:str) -> NamedTuple("Outputs",  [("gcp_resources", str)]):
  # Launch the dataflow job
  dataflow_job_id = [dataflow-id]
  dataflow_resources = GcpResources()
  dr = dataflow_resources.resources.add()
  dr.resource_type='DataflowJob'
  dr.resource_uri=f'https://dataflow.googleapis.com/v1b3/projects/{project}/locations/{location}/jobs/{dataflow_job_id}'
  gcp_resources=MessageToJson(dataflow_resources)
  return gcp_resources

Supported resource_type values

You can set the resource_type to be an arbitrary string, but only the following types have links in the Google Cloud console:

  • BatchPredictionJob
  • BigQueryJob
  • CustomJob
  • DataflowJob
  • HyperparameterTuningJob

Write a component to cancel the underlying resources

When a pipeline job is canceled, the default behavior is for the underlying Google Cloud resources to keep running. They are not canceled automatically. To change this behavior, you should attach a SIGTERM handler to the pipeline job. A good place to do this is just before a polling loop for a job that could run for a long time.

Cancellation has been implemented on several Google Cloud Pipeline Components, including:

  • Batch prediction job
  • BigQuery ML job
  • Custom job
  • Hyperparameter tuning job

For more information, including sample code that shows how to attach a SIGTERM handler, see the following GitHub links:

Consider the following when implementing your SIGTERM handler:

  • Cancellation propagation works only after the component has been running for a few minutes. This is typically due to background startup tasks that need to be processed before the Python signal handlers are called.
  • Some Google Cloud resources might not have cancellation implemented. For example, creating or deleting an Vertex Endpoint or Model could create a long-running operation that accepts a cancellation request through its REST API, but doesn't implement the cancellation operation itself.