Sparkling Vertex AI Pipelines
By: Ivan Nardini
Tags: Spark, Vertex, Google Cloud Platform, Pipeline, Mlops

Disclaimer: This article assumes that you already have some prior knowledge of MLOps, machine learning pipelines and Vertex AI. In case you don’t, I would recommend watching these videos from Priyanka Vergadia for a quick overview: Introduction to MLOps and Vertex Pipelines and End-to-end MLOps with Vertex AI.

Premises

Vertex AI Pipelines is one of the most powerful services of the Vertex AI MLOps features launched this year on Google Cloud.They make it really easy to orchestrate machine learning workloads using Kubeflow. Indeed, you can do almost anything inside a pipeline component. And you’re not limited to certain Vertex functions.

So what if you have a PySpark ML pipeline?

In this article, we’ll illustrate that generality by using Spark inside some Vertex Pipeline components. In particular, I will show how to use Vertex AI Pipelines in conjunction with Dataproc to train and deploy a ML model for near-real time predictive maintenance application.

Our scenario

Imagine you are a ML engineer, and now that your company is moving to the cloud, you’ve been told to lead a proof of concept and prove how to train a predictive maintenance pipeline. Basically, the goal is to process machine log files loaded from on-prem to cloud continuously for training and deploying a ML pipeline you already built on prem.

Regarding the process, you need to become more independent on running (ML) workloads compared in order to speed up the experimental phases¹. Also we need to adopt MLOps practices to make the process reliable, scalable and reproducible across the team.

At the end, you found that Vertex AI Pipelines accomplishes all of this.

Indeed, you will be able to submit ML jobs with Kubeflow based technology, which we assume you are already familiar with compared to other cluster orchestration tools like Workflow. Also, they enable you to track artifacts and lineage of data, feature, model and experiment metrics across your ML workflow with the associated metadata store and the other Vertex AI services such as Datasets, Feature Store and Experiments.

The Dataset

In order to implement this scenario, we use the UCI Machine Learning — AI4I 2020 Predictive Maintenance dataset, which is a synthetic dataset that reflects real predictive maintenance data encountered in industry. It consists of 14 features associated with the machine process such as process temperature, rotational speed and torque. And the target machine failure label indicates whether the machine has failed for at least one of the defined failure modes.

The Anatomy of our Pipeline: Architecture, Components and Representation

About Vertex AI Pipeline, we assumed that

each component of the pipeline has to create a Dataproc cluster, process a PySpark job and destroy the cluster.

Someone could argue that this pattern adds extra running time. That’s true, but each step of the pipeline may need different amounts of resources and then it would need a different cluster configuration. One example could be the hyperparameter tuning step which generally needs more processing power. Also, one of the advantages of using Dataproc is related to the idea of ephemeral clusters which gives more control and more flexibility at the same time.

With that said, below you can find the Pipeline Architecture