Deploying MLflow in GCP Using Terraform: A Step-by-Step Guide

Setup *docker* auth once:

gcloud auth configure-docker europe-west2-docker.pkg.dev

build imagine :

docker build -t mlflow-gcp .

tag :

docker tag mlflow-gcp europe-west2-docker.pkg.dev/test-ml-flow-terraform/test-ml-flow-terraform-repo/mlflow-imagine:latest

and push :

docker push europe-west2-docker.pkg.dev/test-ml-flow-terraform/test-ml-flow-terraform-repo/mlflow-imagine:latest

Step 8: Creating a Database Instance

1. First, we need to create an instance that will then hold our database (a new .tf file for the database should be created).
   
   1. • The version used is PostgreSQL 14 (the latest available currently), minimal configuration for this version is 1 CPU and 3.75GB RAM, so a custom tier needs to be created (read more about creating custom tiers here).
      • The instance needs to have some IP, so either use ipv4_enabled or create a private network if that is a better option for you.

        data "google_project" "project" {
          project_id = var.project_name
        }
        
        resource "google_sql_database_instance" "main" {
          name             = "${var.project_name}-${var.env}-db"
          database_version = "POSTGRES_14"
          project          = var.project_name
        
          settings {
            tier            = "db-custom-1-3840" # vCPUs: 1 Memory: 3.75 GB
            disk_size       = "10"               # HDD storage: 10 GB
            disk_type       = "PD_HDD"
            disk_autoresize = false
            ip_configuration {
              ipv4_enabled  = true
              authorized_networks {
                name  = "Your VPN"
                value = var.vpn_to_access_db
              }
            }
          }
        }

        Use the command Terraform plan. If there are no issues, execute Terraform apply. Be aware that creating database instances may take anywhere from 5 to 20 minutes.

        It should look like this:

2. Once the instance is up and ready, we can create the database. It will attach to the instance we created previously.

   resource "google_sql_database" "ml_flow_db" {
     name     = "mlflow-${var.project_name}-${var.env}"
     instance = google_sql_database_instance.main.name
     project  = var.project_name
   }

3. The next step is to create a user. The password is created using random_password resource to generate random symbols for you. The login, password, and database URL are added to secrets, as they will be needed later.

   resource "google_sql_user" "ml_flow_user" {
     name     = "mlflow-${var.project_name}-${var.env}"
     instance = google_sql_database_instance.main.name
     password = random_password.db_password_ml_flow.result
     project  = var.project_name
   }

   Creating a secret that will hold the login:

   resource "google_secret_manager_secret" "db_login_ml_flow" {
     secret_id = "mlflow_tracking_username"
     project   = var.project_name
     replication {
       automatic = true
     }
   }
   
   resource "google_secret_manager_secret_version" "ml-flow-username-version-basic" {
     secret      = google_secret_manager_secret.db_login_ml_flow.id
     secret_data = google_sql_user.ml_flow_user.name
   }
   

   Creating a password and secret that will hold the password:

   resource "random_password" "db_password_ml_flow" {
     length           = 16
     special          = true
     override_special = "!#$%&*()-_=+[]{}<>:?"
   }
   
   resource "google_secret_manager_secret" "db_password_ml_flow" {
     secret_id = "mlflow_tracking_password"
     project   = var.project_name
     replication {
       automatic = true
     }
   }
   
   resource "google_secret_manager_secret_version" "ml-flow-pass-version-basic" {
     secret      = google_secret_manager_secret.db_password_ml_flow.id
     secret_data = random_password.db_password_ml_flow.result
   }

   Creating a secret with a database connection link that we will use later:

   resource "google_secret_manager_secret" "mlflow_database_url" {
     secret_id = "mlflow_database_url"
     project   = var.project_name
     replication {
       automatic = true
     }
   }
   
   resource "google_secret_manager_secret_version" "mlflow_database_url-version-basic" {
     secret      = google_secret_manager_secret.mlflow_database_url.id
     secret_data = "postgresql+psycopg2://${google_sql_user.ml_flow_user.name}:${random_password.db_password_ml_flow.result}@/${google_sql_database.ml_flow_db.name}?host=/cloudsql/${google_sql_database_instance.main.connection_name}"
   }

   Use a Terraform plan and then Terraform apply.

Step 9: Creating Service Accounts for CloudRun and MLflow

Next, we have two service accounts to create, one with CloudRun and the second to connect to MLflow. We will assign the rights to them in a separate file (a new .tf file for service accounts should be created).

1. For CloudRun – it will only be used by CloudRun:

   resource "google_service_account" "cloudrun-mlflow" {
     account_id   = "cloudrun-mlflow-${var.env}"
     display_name = "Service Account for Cloud Run running ML Flow"
   }

2. For connecting to MLflow and creating objects in Bucket, you will need the key to connect to the MLflow using a script.

   resource "google_service_account" "bucket_object_creator" {
     account_id   = "mlflow-connect-${var.env}"
     display_name = "Service Account for create objects in ML Flow bucket."
   }
   
   resource "google_service_account_key" "key_to_bucket_object_creator" {
     service_account_id = google_service_account.bucket_object_creator.name
   }
   
   resource "google_secret_manager_secret" "bucket_object_creator_account_key" {
     secret_id = "mlflow_service_account_key"
     project   = var.project_name
     replication {
       automatic = true
     }
   }
   
   resource "google_secret_manager_secret_version" "mlflow_key-version-basic" {
     secret      = google_secret_manager_secret.bucket_object_creator_account_key.id
     secret_data = base64decode(google_service_account_key.key_to_bucket_object_creator.private_key)
   }

After Terraform plan and Terraform apply, you should be able to see two new service accounts; one should have a key.

Step 10: Creating Storage Bucket for MLflow Artifacts

At this point, we will be creating a storage bucket for our MLflow artifacts, the bucket is created as regional, but if you need higher availability, then you can use other types (a new .tf file for the bucket should be created).

resource "google_storage_bucket" "mlflow_artifacts_bucket" {
  name                        = "${var.project_name}-mlflow-${var.env}-${var.region}"
  project                     = var.project_name
  location                    = "EUROPE-WEST2"
  storage_class               = "REGIONAL"
  uniform_bucket_level_access = true
  public_access_prevention    = "enforced"
  force_destroy               = false
}

Bucket url needs to be inserted into the secret as we will need it later.

resource "google_secret_manager_secret" "mlflow_artifact_url" {
  secret_id = "mlflow_artifact_url"
  project   = var.project_name
  replication {
    automatic = true
  }
}

resource "google_secret_manager_secret_version" "ml-flow-artifact-url-version-basic" {
  secret      = google_secret_manager_secret.mlflow_artifact_url.id
  secret_data = google_storage_bucket.mlflow_artifacts_bucket.url
}

Step 11: IAM Configuration

1. Add rights to the CloudRun service account (created earlier) to be able to connect to the database (new .tf file for IAM policies should be created).

   resource "google_project_iam_member" "cloudsql_connect" {
     project = var.project_name
     role    = "roles/cloudsql.client"
     member  = "serviceAccount:${google_service_account.cloudrun-mlflow.email}"
   }

2. Add rights to the CloudRun service account, so it can read secrets.

   resource "google_secret_manager_secret_iam_member" "db_login_ml_flow_accessor" {
     secret_id = google_secret_manager_secret.db_login_ml_flow.id
     role      = "roles/secretmanager.secretAccessor"
     member    = "serviceAccount:${google_service_account.cloudrun-mlflow.email}"
   }
   
   resource "google_secret_manager_secret_iam_member" "db_password_ml_flow_accessor" {
     secret_id = google_secret_manager_secret.db_password_ml_flow.id
     role      = "roles/secretmanager.secretAccessor"
     member    = "serviceAccount:${google_service_account.cloudrun-mlflow.email}"
   }
   
   resource "google_secret_manager_secret_iam_member" "mlflow_artifact_url_accessor" {
     secret_id = google_secret_manager_secret.mlflow_artifact_url.id
     role      = "roles/secretmanager.secretAccessor"
     member    = "serviceAccount:${google_service_account.cloudrun-mlflow.email}"
   }
   
   resource "google_secret_manager_secret_iam_member" "mlflow_database_url_accessor" {
     secret_id = google_secret_manager_secret.mlflow_database_url.id
     role      = "roles/secretmanager.secretAccessor"
     member    = "serviceAccount:${google_service_account.cloudrun-mlflow.email}"
   }

   All secrets should look like this:
   

3. Create a policy for the bucket for your team to be able to look at the objects, and for the service account to create the objects.

data "google_iam_policy" "mlflow-bucket-policy" {
  binding {
    role    = "roles/storage.objectViewer"
    members = concat(var.bucket_users_list, ["serviceAccount:${google_service_account.cloudrun-mlflow.email}"])
  }
  binding {
    role = "roles/storage.objectCreator"
    members = [
      "serviceAccount:${google_service_account.bucket_object_creator.email}",
    ]
  }
}

resource "google_storage_bucket_iam_policy" "policy_bucket_object_create" {
  bucket      = google_storage_bucket.mlflow_artifacts_bucket.name
  policy_data = data.google_iam_policy.mlflow-bucket-policy.policy_data

  depends_on = [google_service_account.cloudrun-mlflow, google_storage_bucket.mlflow_artifacts_bucket]
}

You should be able to see these permissions in the bucket:

Step 12: Accessing Secret Manager 

If everything was done correctly, you should see this:

Step 13: Cloud Run Configuration

This resource will be the most complex of all. Here is the full resource (a new .tf file for cloud run should be created): 

• template/service_account – the service account that Terraform will use
• template/volumes – tells CloudRun that it needs to connect to the database
• template/containers – holds information on what container should be used; you can also declare env values for the container here:

  data "google_project" "project" {
    project_id = var.project_name
  }
  
  resource "google_cloud_run_v2_service" "mlflow_on_cloudrun" {
    name     = "mlflow-${var.project_name}-${var.env}"
    location = var.region
    ingress  = "INGRESS_TRAFFIC_ALL"
    client   = "cloud-console"
  
    depends_on = [google_service_account.cloudrun-mlflow, google_artifact_registry_repository.ar_mlflow]
  
    template {
      service_account = google_service_account.cloudrun-mlflow.email
      volumes {
        name = "cloudsql"
  
        cloud_sql_instance {
          instances = [
            google_sql_database_instance.main.connection_name,
          ]
        }
      }
      scaling {
        max_instance_count = 10
      }
      containers {
        image = "${var.region}-docker.pkg.dev/${var.project_name}/${google_artifact_registry_repository.ar_mlflow.name}/${var.image_name}:latest"
        env {
          name  = "GCP_PROJECT"
          value = data.google_project.project.number
        }
        volume_mounts {
          mount_path = "/cloudsql"
          name       = "cloudsql"
        }
  
        resources {
          cpu_idle = true
          limits = {
            cpu    = "1000m" # 1 vCPU
            memory = "1024Mi"
          }
        }
      }
    }
  }

  This part of the code lets everyone access the link to the MLflow, which is then secured by the login and password.

  resource "google_cloud_run_service_iam_binding" "access_to_mlflow" {
    location = google_cloud_run_v2_service.mlflow_on_cloudrun.location
    service  = google_cloud_run_v2_service.mlflow_on_cloudrun.name
    role     = "roles/run.invoker"
    members = [
      "allUsers"
    ]
  }

If you plan on uploading sensitive data to the MLflow, I suggest adding another security layer.

Conclusion

And with that, we have reached the end of this guide. I hope that you found this tutorial insightful and beneficial.

As you can see, setting up MLflow on GCP using Terraform brings about enhanced automation and streamlining of ML infrastructure. But Terraform offers more than just automation – it also provides key benefits like code reusability, simplified management, and reduced provisioning time, all of which are vital in a fast-paced ML environment.

As machine learning advances, staying informed and adaptable to emerging trends and technologies is crucial. If you found this guide helpful, consider subscribing to our newsletter. We plan to publish more insightful tutorials that could significantly enhance the efficiency of your work.

Fill in the bucket data (on the screen, you can see sample data). Consider choosing ‘Multiregion location type’ if your project needs higher availability. You can adjust other options as required.