Google Cloud VertexAI Orchestrator

Orchestrating your pipelines to run on Vertex AI.

Vertex AI Pipelines is a serverless ML workflow tool running on the Google Cloud Platform. It is an easy way to quickly run your code in a production-ready, repeatable cloud orchestrator that requires minimal setup without provisioning and paying for standby compute.

This component is only meant to be used within the context of a remote ZenML deployment scenario. Usage with a local ZenML deployment may lead to unexpected behavior!

When to use it

You should use the Vertex orchestrator if:

• you're already using GCP.

• you're looking for a proven production-grade orchestrator.

• you're looking for a UI in which you can track your pipeline runs.

• you're looking for a managed solution for running your pipelines.

• you're looking for a serverless solution for running your pipelines.

How to deploy it

In order to use a Vertex AI orchestrator, you need to first deploy ZenML to the cloud. It would be recommended to deploy ZenML in the same Google Cloud project as where the Vertex infrastructure is deployed, but it is not necessary to do so. You must ensure that you are connected to the remote ZenML server before using this stack component.

The only other thing necessary to use the ZenML Vertex orchestrator is enabling Vertex-relevant APIs on the Google Cloud project.

In order to quickly enable APIs, and create other resources necessary for using this integration, you can also consider using the Vertex AI stack recipe, which helps you set up the infrastructure with one click.

How to use it

The Vertex Orchestrator (and GCP integration in general) currently only works for Python versions <3.11. The ZenML team is aware of this dependency clash/issue and is working on a fix. For now, please use Python <3.11 together with the GCP integration.

To use the Vertex orchestrator, we need:

• The ZenML gcp integration installed. If you haven't done so, run

  zenml integration install gcp

• Docker installed and running.

• A remote artifact store as part of your stack.

• A remote container registry as part of your stack.

• GCP credentials with proper permissions

• The GCP project ID and location in which you want to run your Vertex AI pipelines.

GCP credentials and permissions

This part is without doubt the most involved part of using the Vertex orchestrator. In order to run pipelines on Vertex AI, you need to have a GCP user account and/or one or more GCP service accounts set up with proper permissions, depending on whether you wish to practice the principle of least privilege and distribute permissions across multiple service accounts.

You also have three different options to provide credentials to the orchestrator:

• use the gcloud CLI to authenticate locally with GCP

• configure the orchestrator to use a service account key file to authenticate with GCP by setting the service_account_path parameter in the orchestrator configuration.

• (recommended) configure a GCP Service Connector with GCP credentials and then link the Vertex AI Orchestrator stack component to the Service Connector.

This section explains the different components and GCP resources involved in running a Vertex AI pipeline and what permissions they need, then provides instructions for three different configuration use-cases:

1. use the local gcloud CLI configured with your GCP user account, including the ability to schedule pipelines

2. use a GCP Service Connector and a single service account with all permissions, including the ability to schedule pipelines

3. use a GCP Service Connector and multiple service accounts for different permissions, including the ability to schedule pipelines

Vertex AI pipeline components

To understand what accounts you need to provision and why, let's look at the different components of the Vertex orchestrator:

1. the ZenML client environment is the environment where you run the ZenML code responsible for building the pipeline Docker image and submitting the pipeline to Vertex AI, among other things. This is usually your local machine or some other environment used to automate running pipelines, like a CI/CD job. This environment needs to be able to authenticate with GCP and needs to have the necessary permissions to create a job in Vertex Pipelines, (e.g. the Vertex AI User role). If you are planning to run pipelines on a schedule, the ZenML client environment also needs additional permissions:

   • the Storage Object Creator Role to be able to write the pipeline JSON file to the artifact store directly (NOTE: not needed if the Artifact Store is configured with credentials or is linked to Service Connector)

2. the Vertex AI pipeline environment is the GCP environment in which the pipeline steps themselves are running in GCP. The Vertex AI pipeline runs in the context of a GCP service account which we'll call here the workload service account. The workload service account can be explicitly configured in the orchestrator configuration via the workload_service_account parameter. If it is omitted, the orchestrator will use the Compute Engine default service account for the GCP project in which the pipeline is running. This service account needs to have the following permissions:

   • permissions to run a Vertex AI pipeline, (e.g. the Vertex AI Service Agent role).

As you can see, there can be dedicated service accounts involved in running a Vertex AI pipeline. That's two service accounts if you also use a service account to authenticate to GCP in the ZenML client environment. However, you can keep it simple and use the same service account everywhere.

Configuration use-case: local gcloud CLI with user account

This configuration use-case assumes you have configured the gcloud CLI to authenticate locally with your GCP account (i.e. by running gcloud auth login). It also assumes the following:

• your GCP account has permissions to create a job in Vertex Pipelines, (e.g. the Vertex AI User role).

• the Compute Engine default service account for the GCP project in which the pipeline is running is updated with additional permissions required to run a Vertex AI pipeline, (e.g. the Vertex AI Service Agent role).

This is the easiest way to configure the Vertex AI Orchestrator, but it has the following drawbacks:

• the setup is not portable on other machines and reproducible by other users.

• it uses the Compute Engine default service account, which is not recommended, given that it has a lot of permissions by default and is used by many other GCP services.

We can then register the orchestrator as follows:

zenml orchestrator register <ORCHESTRATOR_NAME> \
    --flavor=vertex \
    --project=<PROJECT_ID> \
    --location=<GCP_LOCATION> \
    --synchronous=true

Configuration use-case: GCP Service Connector with single service account

This use-case assumes you have already configured a GCP service account with the following permissions:

• permissions to create a job in Vertex Pipelines, (e.g. the Vertex AI User role).

• permissions to run a Vertex AI pipeline, (e.g. the Vertex AI Service Agent role).

• the Storage Object Creator Role to be able to write the pipeline JSON file to the artifact store directly.

It also assumes you have already created a service account key for this service account and downloaded it to your local machine (e.g. in a connectors-vertex-ai-workload.json file). This is not recommended if you are conscious about security. The principle of least privilege is not applied here and the environment in which the pipeline steps are running has many permissions that it doesn't need.

zenml service-connector register <CONNECTOR_NAME> --type gcp --auth-method=service-account --project_id=<PROJECT_ID> --service_account_json=@connectors-vertex-ai-workload.json --resource-type gcp-generic

zenml orchestrator register <ORCHESTRATOR_NAME> \
    --flavor=vertex \
    --location=<GCP_LOCATION> \
    --synchronous=true \
    --workload_service_account=<SERVICE_ACCOUNT_NAME>@<PROJECT_NAME>.iam.gserviceaccount.com 
    
zenml orchestrator connect <ORCHESTRATOR_NAME> --connector <CONNECTOR_NAME>

Configuration use-case: GCP Service Connector with different service accounts

This setup applies the principle of least privilege by using different service accounts with the minimum of permissions needed for the different components involved in running a Vertex AI pipeline. It also uses a GCP Service Connector to make the setup portable and reproducible. This configuration is a best-in-class setup that you would normally use in production, but it requires a lot more work to prepare.

This setup involves creating and configuring several GCP service accounts, which is a lot of work and can be error prone. If you don't really need the added security, you can use the GCP Service Connector with a single service account instead.

The following GCP service accounts are needed:

1. a "client" service account that has the following permissions:

   • permissions to create a job in Vertex Pipelines, (e.g. the Vertex AI User role).

   • permissions to create a Google Cloud Function (e.g. with the Cloud Functions Developer Role).

   • the Storage Object Creator Role to be able to write the pipeline JSON file to the artifact store directly (NOTE: not needed if the Artifact Store is configured with credentials or is linked to Service Connector).

2. a "workload" service account that has permissions to run a Vertex AI pipeline, (e.g. the Vertex AI Service Agent role).

A key is also needed for the "client" service account. You can create a key for this service account and download it to your local machine (e.g. in a connectors-vertex-ai-workload.json file).

With all the service accounts and the key ready, we can register the GCP Service Connector and Vertex AI orchestrator as follows:

zenml service-connector register <CONNECTOR_NAME> --type gcp --auth-method=service-account --project_id=<PROJECT_ID> --service_account_json=@connectors-vertex-ai-workload.json --resource-type gcp-generic

zenml orchestrator register <ORCHESTRATOR_NAME> \
    --flavor=vertex \
    --location=<GCP_LOCATION> \
    --synchronous=true \
    --workload_service_account=<WORKLOAD_SERVICE_ACCOUNT_NAME>@<PROJECT_NAME>.iam.gserviceaccount.com 

zenml orchestrator connect <ORCHESTRATOR_NAME> --connector <CONNECTOR_NAME>

Configuring the stack

With the orchestrator registered, we can use it in our active stack:

# Register and activate a stack with the new orchestrator
zenml stack register <STACK_NAME> -o <ORCHESTRATOR_NAME> ... --set

ZenML will build a Docker image called <CONTAINER_REGISTRY_URI>/zenml:<PIPELINE_NAME> which includes your code and use it to run your pipeline steps in Vertex AI. Check out this page if you want to learn more about how ZenML builds these images and how you can customize them.

You can now run any ZenML pipeline using the Vertex orchestrator:

python file_that_runs_a_zenml_pipeline.py

Vertex UI

Vertex comes with its own UI that you can use to find further details about your pipeline runs, such as the logs of your steps.

Vertex UI

For any runs executed on Vertex, you can get the URL to the Vertex UI in Python using the following code snippet:

from zenml.client import Client

pipeline_run = Client().get_pipeline_run("<PIPELINE_RUN_NAME>")
orchestrator_url = pipeline_run.run_metadata["orchestrator_url"].value

Run pipelines on a schedule

The Vertex Pipelines orchestrator supports running pipelines on a schedule using its native scheduling capability.

How to schedule a pipeline

from zenml.config.schedule import Schedule

# Run a pipeline every 5th minute
pipeline_instance.run(
    schedule=Schedule(
        cron_expression="*/5 * * * *"
    )
)

# Run a pipeline every hour
# starting in one day from now and ending in three days from now
pipeline_instance.run(
    schedule=Schedule(
        cron_expression="0 * * * *"
        start_time=datetime.datetime.now() + datetime.timedelta(days=1),
        end_time=datetime.datetime.now() + datetime.timedelta(days=3),
    )
)

The Vertex orchestrator only supports the cron_expression, start_time (optional) and end_time (optional) parameters in the Schedule object, and will ignore all other parameters supplied to define the schedule.

The start_time and end_time timestamp parameters are both optional and are to be specified in local time. They define the time window in which the pipeline runs will be triggered. If they are not specified, the pipeline will run indefinitely.

The cron_expression parameter supports timezones. For example, the expression TZ=Europe/Paris 0 10 * * * will trigger runs at 10:00 in the Europe/Paris timezone.

How to delete a scheduled pipeline

Note that ZenML only gets involved to schedule a run, but maintaining the lifecycle of the schedule is the responsibility of the user.

In order to cancel a scheduled Vertex pipeline, you need to manually delete the schedule in VertexAI (via the UI or the CLI).

Additional configuration

For additional configuration of the Vertex orchestrator, you can pass VertexOrchestratorSettings which allows you to configure node selectors, affinity, and tolerations to apply to the Kubernetes Pods running your pipeline. These can be either specified using the Kubernetes model objects or as dictionaries.

from zenml.integrations.gcp.flavors.vertex_orchestrator_flavor import VertexOrchestratorSettings
from kubernetes.client.models import V1Toleration

vertex_settings = VertexOrchestratorSettings(
    pod_settings={
        "affinity": {
            "nodeAffinity": {
                "requiredDuringSchedulingIgnoredDuringExecution": {
                    "nodeSelectorTerms": [
                        {
                            "matchExpressions": [
                                {
                                    "key": "node.kubernetes.io/name",
                                    "operator": "In",
                                    "values": ["my_powerful_node_group"],
                                }
                            ]
                        }
                    ]
                }
            }
        },
        "tolerations": [
            V1Toleration(
                key="node.kubernetes.io/name",
                operator="Equal",
                value="",
                effect="NoSchedule"
            )
        ]
    }
)

If your pipelines steps have certain hardware requirements, you can specify them as ResourceSettings:

resource_settings = ResourceSettings(cpu_count=8, memory="16GB")

These settings can then be specified on either pipeline-level or step-level:

# Either specify on pipeline-level
@pipeline(
    settings={
        "orchestrator.vertex": vertex_settings,
        "resources": resource_settings,
    }
)
def my_pipeline():
    ...

# OR specify settings on step-level
@step(
    settings={
        "orchestrator.vertex": vertex_settings,
        "resources": resource_settings,
    }
)
def my_step():
    ...

Check out the SDK docs for a full list of available attributes and this docs page for more information on how to specify settings.

For more information and a full list of configurable attributes of the Vertex orchestrator, check out the API Docs .

Enabling CUDA for GPU-backed hardware

Note that if you wish to use this orchestrator to run steps on a GPU, you will need to follow the instructions on this page to ensure that it works. It requires adding some extra settings customization and is essential to enable CUDA for the GPU to give its full acceleration.