Effortless MLflow Deployment With Terraform On GCP

Scenario 1: Deploying MLflow without Terraform

Without using Terraform, deploying MLflow in the cloud would require manually creating and configuring all the required cloud resources using the web console or CLI tools. This process would involve several steps, such as:

• Setting up a virtual machine (VM) to host the MLflow server
• Configuring the VM with the necessary dependencies and software
• Creating and configuring a Cloud Storage bucket for artifact storage
• Setting up a Cloud SQL instance for MLflow backend
• Configuring authentication and access management
• Manually connecting all the components and testing the setup

This approach can be time-consuming, error-prone, and difficult to reproduce 😵‍💫 especially if you need to deploy the same infrastructure multiple times or across different environments

Scenario 2: Deploying MLflow with Terraform

Using Terraform to deploy MLflow in the cloud simplifies the process by automating the creation and configuration of all necessary resources. With Terraform, you would:

• Write a Terraform configuration file that defines all the required infrastructure components (VM, Cloud Storage, Cloud SQL, etc.)
• Use the Terraform CLI to apply the configuration, which creates and configures the resources automatically
• Make any necessary adjustments to the infrastructure by updating the Terraform configuration file and reapplying it

In summary, while there is a learning curve when it comes to Terraform’s HCL language, once you become familiar with it, you can take advantage of the benefits of automated, streamlined, and customizable infrastructure deployments not only for your MLflow projects, but for other projects as well.

Fear not, however. You will not need to learn Terraform’s HCL language to deploy MLflow with this template 😌
On the contrary, you will need to dive into HCL language if you want to change the architecture in the future.

Setting up Cloud Identity and Access Management (IAM)

Skip this section if your GCP project is part of an organization. In fact, for creating the Identity Aware Proxy (IAP) brand, we need a Cloud Identity or a Google Workspace account.

To setup the OAuth 2.0 authentication to give your organization access to a MLflow Server, Terraform needs a gcp.iap.brand resource available only in project that belong to an organization. In this step, we will create a free Google Workspace account with our private domain.

If you don’t have a private domain, you can buy one in Google Domains or GoDaddy.

• Go to the IAM & Admin page.
• Go to Identity and Organization
• Click on Go to the checklist
• Click on Begin setup
• Tick the I'm a new customer checkbox
• Click on Sign up for Cloud Identity
• Follow the setup and add your private domain url when asked

You should have created an admin account on your private domain by now. Next step, we need to verify our private domain here:

• Go to the DNS Administration page of your domain provider (like GoDaddy)
• Copy the key that google sends you and create a TXT reccord with @ name and the key as value
• Click on Verify in the google page

You now have a free Google Workspace account and can head on to the next steps 👏

Google Cloud Platform Setup

If you already have gcloud CLI installed and an existing project, skip this section.

Otherwise:

• Follow the installation of the Google Cloud SDK here.
• Login to the Google Cloud Console with the admin account you created.
• Go to the Create Project page and select your organization.

You can now authenticate your gcloud client with the admin account you created:

gcloud auth login

Finally, set you default project:

gcloud auth application-default login --project $PROJECT

Terraform Installation

Install Terraform following this link.

Test your installation:

terraform version

Clone the repository containing the Terraform scripts to deploy the MLflow infrastructure:

git clone git@github.com:ConsciousML/one-click-mlflow.git

App Engine Flex for MLflow Web Server

The MLflow Web Server is like your personal assistant for managing machine learning workflows, providing a dashboard where you can track experiments, models, and more.

App Engine Flex is an excellent choice for hosting the MLflow Web Server on GCP due to its flexible and scalable architecture. As a managed platform, it abstracts away the underlying infrastructure and automatically scales the web server to meet the incoming traffic demand.

By leveraging App Engine Flex, you can focus on developing and deploying your MLflow workflows, leaving the infrastructure management to Google.

Cloud Storage as Artifact Backend

Let’s talk about the MLflow Artifact Backend, a fancy name for where we store our machine learning experiment outputs, such as trained models and artifacts 📜

Cloud Storage is a perfect fit for this role, offering virtually limitless, highly available, and durable object storage for your data, accessible from anywhere in the world. By using Cloud Storage as the Artifact Backend for your MLflow project, you can store and version your artifacts reliably, and access them with low latency from your MLflow Web Server 📈

Secret Manager for storing secrets

Secret Manager is a powerful GCP tool that helps you store and manage sensitive data, such as passwords, API keys, and certificates, securely 🔑

For MLflow, the Secret Manager comes in handy for managing credentials and other sensitive information required to access your cloud resources, such as Cloud SQL instances and Cloud Storage buckets. By storing your secrets in the Secret Manager, you can avoid hardcoding them in your code or configuration files, which can be a security risk.

With the GCP Secret Manager, you can rest assured that your sensitive information is kept safe and accessible only to authorized users and services.

Cloud SQL as MLflow Backend

Let’s talk about Cloud SQL, a fully-managed relational database service that runs on GCP. When it comes to MLflow, Cloud SQL serves as an excellent choice for your backend database, offering scalability, high availability, and automatic backups without any maintenance overhead.

The Cloud SQL instance will serve as the MLflow backend so that you can store your experiment logs. The MLflow UI will then communicate with the backend to offer visualization regarding hyperparameters, metrics, and tags.

By using Cloud SQL as the backend for your MLflow projects, you can focus on your data science work while leaving the database management to Google 😌

Want to learn more on the subject, feel free to explore my previous article on Cloud SQL vs BigQuery.

Log Pusher Service Account for MLflow Communication

The Log Pusher Service Account is a special account that enables secure communication between your MLflow components and Google Cloud Platform (GCP) services. It comes in handy when your MLflow code needs to access your Artifact Backend or the Web Server, requiring authentication and authorization.

By using a dedicated service account for this purpose, you can avoid using your personal credentials, which can be a security risk, and benefit from granular access control.

In practice, we also use it when a Data Scientist wants to push logs to MLflow. He need to download a service account key file from GCP and use it to authenticate with the MLflow server in his Python code.

Managing Terraform State in the Cloud

Hashicorp, the developer of Terraform, recommendeds hosting your state using Terraform Cloud, their proprietary SaaS offering. But we can also host our Terraform state on a Cloud Storage and benefit from this pratice that offers several benefits for managing your IAC projects.

Let’s take a simple example: let’s suppose you create a IAC project and you add ressources to it. When you deploy it, the current state of my GCP project will be tracked by a Terraform state file in my local directory.

What if a colleague of yours pull the repository and starts making some changes ? Our Terraform states will be out of sync and we will have to deal with conflicts 🤕

By setting up the backend of your Terraform project to a remote Cloud Storage, you can avoid this issue and have a single source of truth for your Terraform state.

Migrating Terraform State to Cloud SQL

Here how to do it in 7 steps:

1. Create a another project, if your MLflow is deployed on the mlflow project, name the new project as mlflow-tf-state.
2. Create a bucket on the new project and name it mlflow-tf-state.
3. Create a backend folder under your directory containing your terraform code.
4. Create a backend.tf file in backend with the following content:

terraform {
 backend "gcs" {
   bucket  = "mlflow-tf-state"
   prefix  = "terraform/state"
 }
}

5. Create a main.tf file in backend with the following content:

provider "google" {
  project = "mlflow-tf-state"
}

6. Run terraform init -migrate-state to initialize the backend.
7. Type yes to migrate the state to the new backend.