Securing Vertex AI Pipeline Scheduling

Vertex AI Pipelines provides a great way to orchestrate your Machine Learning workloads in serverless fashion on Google Cloud. It is based on KubeFlow Pipelines, an open-source platform for building Machine Learning (ML) pipelines using containers, making the pipelines in Vertex AI open-source compatible.

However, unlike KubeFlow Pipelines, Vertex Pipelines does not have a built-in mechanism for scheduling pipeline runs. For many use cases, the ability to schedule pipeline runs is a key element of ML automation (e.g., scheduled batch prediction pipelines, scheduled model retraining, etc).

Publicly published Vertex AI documents provides an approach for scheduling Vertex AI Pipeline jobs using Cloud Scheduler service.

1. Cloud Scheduler for scheduling an HTTP/S request, that will trigger a Cloud Function
2. Cloud Function will trigger the Vertex pipeline using the Vertex AI SDK

Fig 1: Google’s proposed approach for scheduling Vertex Pipelines

In this publicly documented approach, we can use Cloud Scheduler to securely call Cloud Functions with OIDC token.

In this approach to use a Cloud Function, you need to send the request with an authentication identification token. Cloud Scheduler automates this for you by allowing you to specify multiple kinds of authentication headers. An OpenID Connect (OIDC) token is the most general way to provide a token in the request header.

A Cloud Function with a HTTP trigger will be created. Cloud Scheduler will use the same http endpoint to trigger the Cloud Function using OIDC token as shown below(Fig 2).

Fig 2: Schedule definition with OIDC settings

On the Cloud Function side, you can use VPC Service Controls to add additional layer of network level security to your Cloud Functions. VPC Service Controls enforces calls to the Cloud Functions API will fail, unless they originate from within the service perimeter.

However keep a note that VPC Service Controls does not support Cloud Scheduler jobs with the following targets:

• App Engine
• HTTP

So, all the calls coming from the Cloud Scheduler will fail as they will be treated as coming from outside the service perimeter.

Therefore, we have to make our Cloud Function public facing as shown below(Fig 3):

Fig 3: Network configuration for a Cloud Function

Security Limitation of the above approach:

• Public facing cloud function endpoint makes it vulnerable to attacks(like DDOS).

To overcome this security limitation you can use GCP services as mentioned below in the specific order.

• Cloud Scheduler job sends a message to the Pub/Sub
• Pub/Sub topic to trigger a Cloud Function
• A Cloud Function, that will trigger the Vertex AI Pipeline using the Vertex AI SDK

Fig 4: Secured approach for triggering the Vertex AI Pipeline

In Cloud Functions, a Pub/Sub trigger enables a function to be called in response to Pub/Sub messages passed via Cloud Scheduler. Your Cloud Function will be triggered whenever a message is published to the specified topic.

Fig 5: Cloud Function with Pub/Sub as trigger

Pub/Sub is fully supported by VPC Service Controls. VPC Service Controls protection applies to all administrator operations, publisher operations, and subscriber operations in case of Pub/Sub.

All the events originating from Pub/Sub will be treated as coming from within the service perimeter, hence not requiring to make our Cloud Function public facing as shown in Fig 6.

Fig 6: Network configuration for Cloud function

This is how you can make the call to Vertex AI pipelines in more secured fashion by leveraging Cloud Scheduler, Pub/Sub and Cloud Function Service.

………..Thank you for reading the blog and have an amazing day!…………