Skip to content

Integrate with Vertex AI

Comet integrates with Google Vertex AI.

Google Vertex AI lets you build, deploy, and scale ML models faster, with pre-trained and custom tooling within a unified artificial intelligence platform.

Visualize Vertex AI Pipelines

Comet partnered with Vertex AI to allow you to track not just individual training runs but also Vertex Pipelines within the Comet UI.

The state of the pipeline can be visualized using the Vertex AI Pipeline Panel available in the Featured tab.

vertex-ai-pipelines-integration.png

You can recreate the view above by:

  • Grouping Experiments by vertex_run_name.
  • Adding the Vertex AI Pipelines panel available in the Featured tab.
  • Saving the view as a Vertex dashboard.

The Vertex AI Pipelines can be used to either visualize the latest state of the DAG or the static graph -- similar to what is available in the Vertex AI Pipelines UI. In addition, all tasks that have the Comet logo are also tracked in Comet and can be accessed by clicking the task.

Integrate in two steps

  1. Add a Comet logger component to track the state of the pipeline.
  2. Add Comet to each individual tasks that you would like to track as a Comet Experiment.

Log the state of a pipeline

The Vertex AI Pipelines integration relies on a component that runs in parallel to the rest of the pipeline. The comet_logger_component logs the state of the pipeline as a whole and reports it back to the Comet UI, allowing you to track the pipeline progress as a whole. This component can be used like this:

import comet_ml.integration.vertex
import kfp.dsl as dsl


@dsl.pipeline(name="ML training pipeline")
def ml_training_pipeline():
    # Add the Comet logger component
    comet_ml.integration.vertex.comet_logger_component(
        api_key="<>", workspace="<>", project_name="<>"
    )

    # Rest of the training pipeline

The Comet logger component can also be configured using environment variables, in which case you will not need to specify the api_key, workspace or project_name arguments.

Log automatically

The Comet logger component creates an Experiment when the Vertex pipeline is compiled into JSON. Through the Experiment, the Vertex pipeline code and all of the other environment details listed here are logged automatically. In addition the following Vertex-specific information is logged:

Item NameItem typeDescription
vertex_run_nameOtherThe Vertex pipeline name
vertex_run_idOtherThe Vertex unique run id
vertex_task_typeOtherInternal field used to distinguish between task and pipeline Experiments
pipeline_typeOtherInternal field used to distinguish between integrations
vertex-pipelineAssetThe status of the Vertex pipeline

If you want to have more control over what is logged automatically or want to log additional information, you can create your own Experiment and pass it to the Comet logger component, like this:

import comet_ml
import comet_ml.integration.vertex
import kfp.dsl as dsl


@dsl.pipeline(name="ML training pipeline")
def ml_training_pipeline():
    # Your own experiment object
    experiment = comet_ml.Experiment(api_key="<>", workspace="<>", project_name="<>")

    # Log an additional code file
    experiment.log_code("./utils/helper.py")

    # Log additional information
    trigger_reason = os.environ.get("TRIGGER_REASON", "Manual trigger")
    experiment.log_other("Trigger reason", trigger_reason)

    # Add the Comet logger component
    comet_ml.integration.vertex.comet_logger_component(custom_experiment=experiment)

    # Flush all of the information
    experiment.end()

    # Rest of the training pipeline

Log the state of each task

In addition to the component, you will need to initialize the Vertex task logger within each task of the pipeline that you wish to run. This can be done by creating an Experiment within each task and using the initialize_comet_logger function:

@kfp.dsl.v2.component
def my_component() -> None:
    import comet_ml.integration.vertex

    experiment = comet_ml.Experiment(api_key="<>", project_name="<>")
    pipeline_run_name = "{{$.pipeline_job_name}}"
    pipeline_task_name = "{{$.pipeline_task_name}}"
    pipeline_task_id = "{{$.pipeline_task_uuid}}"

    comet_ml.integration.vertex.initialize_comet_logger(
        experiment, pipeline_run_name, pipeline_task_name, pipeline_task_id
    )

    # Rest of the task code

The strings "{{$.pipeline_job_name}}", "{{$.pipeline_task_name}}", "{{$.pipeline_task_uuid}}" are placeholders that are replaced at runtime by Vertex AI; you do not need to update them before running the pipeline.

Log automatically

The Experiment automatically logs the Vertex pipeline code and all of the other environment details listed here. In addition, the following Vertex-specific information is logged:

Item NameItem typeDescription
vertex_run_nameOtherThe Vertex pipeline name
vertex_task_idOtherThe Vertex unique task id
vertex_task_nameOtherThe Vertex task name
vertex_task_typeOtherInternal field used to distinguish between task and pipeline Experiments

Note

There are alternatives to setting the API key programatically. See more here.

Nov. 27, 2022