Google Cloud VertexAI

Executing individual steps in Vertex AI.

Vertex AI offers specialized compute instances to run your training jobs and has a comprehensive UI to track and manage your models and logs. ZenML's Vertex AI step operator allows you to submit individual steps to be run on Vertex AI compute instances.

When to use it

You should use the Vertex step operator if:

one or more steps of your pipeline require computing resources (CPU, GPU, memory) that are not provided by your orchestrator.

you have access to Vertex AI. If you're using a different cloud provider, take a look at the SageMaker or AzureML step operators.

How to deploy it

Enable Vertex AI here.

Create a service account with the right permissions to create Vertex AI jobs (roles/aiplatform.admin) and push to the container registry (roles/storage.admin).

How to use it

The GCP step operator (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 step operator, we need:

The ZenML gcp integration installed. If you haven't done so, runCopyzenml integration install gcp

Docker installed and running.

Vertex AI enabled and a service account file. See the deployment section for detailed instructions.

A GCR container registry as part of our stack.

(Optional) A machine type that we want to execute our steps on (this defaults to n1-standard-4). See here for a list of available machine types.

A remote artifact store as part of your stack. This is needed so that both your orchestration environment and VertexAI can read and write step artifacts. Check out the documentation page of the artifact store you want to use for more information on how to set that up and configure authentication for it.

You have three different options to provide GCP credentials to the step operator: