Mlflow end to end deployment

Author: IshmeetMehta

platforms/gke-aiml/playground/mvp_resources.tf

@@ -304,7 +304,7 @@ resource "google_project_iam_member" "mlflow_ksa_user" {
 
   project = data.google_project.environment.project_id
   member  = "${local.wi_member_principal_prefix}/${local.mlflow_ksa}"
-  role    = "roles/alloydb.client"
+  role    = ["roles/alloydb.admin", "roles/serviceusage.serviceUsageConsumer"]
 }
 
 # DATA BUCKET

use-cases/mlflow-inference-pipeline/mlflow-deployment/mlflow/Dockerfile

@@ -10,9 +10,7 @@ RUN pip3 install --no-cache-dir -r requirements.txt
 
 ENV PYTHONUNBUFFERED 1
 ENV PGPASSWORD=""
-# ENV MLFLOW_DATABASE_URI=V_MLFLOW_DATABASE_URI
-# ENV MLFLOW_ARTIFACT_LOCATION=V_MLFLOW_ARTIFACT_LOCATION
 EXPOSE 5000
 
 # Run the MLflow server 
-CMD ["mlflow", "server", "--host", "0.0.0.0", "--port", "5000", "--backend-store-uri","V_MLFLOW_DATABASE_URI" , "--default-artifact-root", "V_MLFLOW_ARTIFACT_LOCATION"]
+CMD ["mlflow", "server", "--host", "0.0.0.0", "--port", "5000", "--backend-store-uri","V_MLFLOW_DATABASE_URI" , "--default-artifact-root", "gs://V_MLFLOW_ARTIFACT_LOCATION"]

use-cases/mlflow-inference-pipeline/mlflow-deployment/mlflow/manifests/deployment.yaml

@@ -14,27 +14,31 @@ spec:
       labels:
         app: mlflow-tracking
     spec:
-      serviceAccountName: mlflow-prod-ml-team-mlflow
-      initContainers:
-        - name: alloydb-auth-proxy
-          restartPolicy: Always
+      serviceAccountName: V_MLFLOW_KSA
+      containers:
+        - name: alloydb-proxy
           image: gcr.io/alloydb-connectors/alloydb-auth-proxy:latest
           args:
-            - "projects/gkebatchenv3a4ec43f/locations/us-central1/clusters/mlp-mlflow-prod/instances/mlp-mlflow-prod-primary"
+            - "V_DB_INSTANCE_URI"
             - "--psc"
             - "--port=5432"
             - "--auto-iam-authn"
             - "--structured-logs"
+          ports:
+            - containerPort: 5432
+              name: alloydb
           securityContext:
             runAsNonRoot: true
           resources:
             requests:
               memory: "2Gi"
               cpu: "1"
-      containers:
         - name: mlflow
-          image: us-central1-docker.pkg.dev/gkebatchenv3a4ec43f/mlflow-prod/mlflow:1.0.0
+          image: V_MLFLOW_IMAGE
           imagePullPolicy: Always
+          ports:
+            - containerPort: 5000
+              name: mlflow-port
           resources:
             requests:
               cpu: "2"