setup VM and networking for mlflow on GCP

Author: divakaivan

src/make_infra/main.tf

@@ -13,7 +13,7 @@ provider "google" {
   region  = var.region
 }
 
-resource "google_storage_bucket" "mlflow-101-raw-data" {
+resource "google_storage_bucket" "mlops-101-raw-data" {
   name          = var.gcs_raw_taxi_data_bucket_name
   location      = var.location
 
@@ -36,7 +36,7 @@ resource "google_storage_bucket" "mlflow-101-raw-data" {
   force_destroy = true
 }
 
-resource "google_storage_bucket" "mlflow-101-processed-data" {
+resource "google_storage_bucket" "mlops-101-processed-data" {
   name          = var.gcs_processed_taxi_data_bucket_name
   location      = var.location
 
@@ -58,3 +58,71 @@ resource "google_storage_bucket" "mlflow-101-processed-data" {
 
   force_destroy = true
 }
+
+resource "google_storage_bucket" "mlops_101_mlflow_artifacts" {
+  name          = var.mlops_101_mlflow_artifacts
+  location      = var.location
+
+  storage_class = var.gcs_storage_class
+  uniform_bucket_level_access = true
+
+  versioning {
+    enabled     = true
+  }
+
+  lifecycle_rule {
+    action {
+      type = "Delete"
+    }
+    condition {
+      age = 30  // days
+    }
+  }
+
+  force_destroy = true
+}
+
+resource "google_compute_instance" "mlops-101-mlflow-server" {
+  name         = "mlops-101-mlflow-server"
+  machine_type = "e2-small"
+  zone         = var.region
+
+  boot_disk {
+    initialize_params {
+      image = "projects/ml-images/global/images/c0-deeplearning-common-cpu-v20240613-debian-11"
+    }
+  }
+
+  network_interface {
+    network = "default"
+    access_config {
+      // Ephemeral IP
+    }
+  }
+
+  tags = ["mlflow"]
+
+  metadata_startup_script = <<-EOT
+    #!/bin/bash
+    echo "Startup script running..." >> /var/log/startup-script.log 2>&1
+    apt-get update >> /var/log/startup-script.log 2>&1
+    apt-get install -y python3-pip >> /var/log/startup-script.log 2>&1
+    pip3 install mlflow==2.20.1 >> /var/log/startup-script.log 2>&1
+    mlflow server -h 0.0.0.0 -p 5000 --backend-store-uri ${var.mlflow_backend_store_uri} --default-artifact-root ${var.mlflow_artifact_location} >> /var/log/startup-script.log 2>&1
+    echo "MLflow server started" >> /var/log/startup-script.log 2>&1
+    EOT
+}
+
+resource "google_compute_firewall" "mlflow-server-firewall" {
+  name    = "mlflow-server-firewall"
+  network = "default"
+  direction = "INGRESS"
+
+  allow {
+    protocol = "tcp"
+    ports    = ["5000"]
+  }
+
+  target_tags = ["mlflow"]
+  source_ranges = [ "0.0.0.0/0" ]
+}

src/make_infra/variables.tf

@@ -31,3 +31,18 @@ variable "gcs_storage_class" {
     description = "Bucket Storage Class"
     default = "STANDARD"
 }
+
+variable "mlops_101_mlflow_artifacts" {
+    description = "Bucket name for MLflow artifacts"
+    default = "mlops_101_mlflow_artifacts"
+}
+
+variable "mlflow_backend_store_uri" {
+    description = "MLflow Backend Store URI"
+    default = "sqlite:///mlops_101.sqlite"
+}
+
+variable "mlflow_artifact_location" {
+    description = "MLflow Artifact Location"
+    default = "gs://mlops_101_mlflow_artifacts"
+}