Integrating MLflow into Kubeflow: A Solution for Model Management

In the rapidly evolving field of machine learning, the ability to efficiently manage and track models is essential for success. Kubeflow has emerged as a powerful platform designed to streamline the deployment and orchestration of machine learning workflows. However, one notable limitation of Kubeflow is its lack of a user-friendly interface for model registration and management. This gap can hinder data scientists and machine learning engineers in effectively tracking their experiments and models.

While a user-friendly UI installation is scheduled for future releases of Kubeflow (as noted in this GitHub issue), we will explore a practical solution to address this challenge in the meantime.

Enter MLflow—a robust open-source platform that simplifies the management of the machine learning lifecycle, including experimentation, reproducibility, and deployment. By integrating MLflow into Kubeflow, users can leverage MLflow’s intuitive UI and comprehensive model registry capabilities to enhance their machine learning workflows.

---

Our initial objective is to deploy MLflow on Kubernetes, a process that follows the same principles as any standard deployment. We will start by creating a Dockerfile. This step is essential because the default MLflow image lacks the boto3 and psycopg2-binary packages, which are necessary for connecting MLflow to MinIO and PostgreSQL:

FROM ghcr.io/mlflow/mlflow:latest

RUN pip install psycopg2-binary boto3

CMD ["mlflow", "server"]

Generate a deployment.yaml file:

Make sure to provide the environment variables using a config map or another appropriate method. We will also include the credentials needed to connect MLflow to MinIO and PostgreSQL in this file:

metadata:
  name: mlflow-service
  namespace: model-registry
spec:
  selector:
    app: mlflow
  ports:
    - protocol: TCP
      port: 5000
      targetPort: 5000
---
apiVersion: v1
kind: ServiceAccount
metadata:
  name: mlflow-sa
  namespace: model-registry
---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: mlflow
  namespace: model-registry
spec:
  replicas: 1
  selector:
    matchLabels:
      app: mlflow
  template:
    metadata:
      labels:
        app: mlflow
    spec:
      serviceAccountName: mlflow-sa
      containers:
      - name: mlflow
        image: harbor.partdp.ir/devops/custom-mlflow:v2.17.2
        ports:
          - containerPort: 5000
        env:
          - name: MLFLOW_TRACKING_URI
            value: "postgresql+psycopg2://mlflow:<PASSWORD>@192.168.34.85:5000/mlflow_db"
          - name: MLFLOW_S3_ENDPOINT_URL
            value: "https://minio.partdp.ir:9000"
          - name: AWS_ACCESS_KEY_ID
            valueFrom:
              secretKeyRef:
                name: mlflow-secret
                key: AWS_ACCESS_KEY_ID
          - name: AWS_SECRET_ACCESS_KEY
            valueFrom:
              secretKeyRef:
                name: mlflow-secret
                key: AWS_SECRET_ACCESS_KEY
        command: ["mlflow", "server"]
        args:
          - "--host=0.0.0.0"
          - "--port=5000"
          - "--backend-store-uri=$(MLFLOW_TRACKING_URI)"
          - "--default-artifact-root=s3://pejman"
        resources:
          requests:
            memory: "512Mi"
            cpu: "500m"
          limits:
            memory: "4Gi"
            cpu: "2"

and a secret:

apiVersion: v1
kind: Secret
metadata:
  name: mlflow-secret
  namespace: model-registry
type: Opaque
data:
  AWS_ACCESS_KEY_ID: <BASE64 encoded>
  AWS_SECRET_ACCESS_KEY: <BASE64 encoded>

---

Next, we need to integrate an MLflow tab into the central dashboard of Kubeflow. To accomplish this, we need to set up a virtual service that will expose the MLflow service through Istio Ingress.

apiVersion: networking.istio.io/v1beta1
kind: VirtualService
metadata:
  name: mlflow
  namespace: model-registry
spec:
  gateways:
  - kubeflow/kubeflow-gateway
  hosts:
  - '*'
  http:
  - match:
    - uri:
        prefix: /mlflow/
    rewrite:
      uri: /
    route:
    - destination:
        host: mlflow-service.model-registry.svc.cluster.local
        port:
          number: 5000

---

Next, we will modify the ConfigMap for Kubeflow's dashboard to make MLflow visible:

kubectl edit cm centraldashboard-config -n kubeflow

and adding this config in menuLinks section:

            { 
                "type": "item",
                "link": "/mlflow/",
                "text": "MlFlow",
                "icon": "icons:cached"
            },

Now just restart the Dashboard:

kubectl rollout restart deploy centraldashboard -n kubeflow


And we have MLflow there: