Not just another MLflow on Kubernetes article
Published: 2022-06-28 06:16:49

There's a ton of content on the internet about deploying an MLflow Tracking Server to Kubernetes and it's great! But if that's the case, then why does this article exist?

When I was looking for a way to run MLflow on our cloud, I wanted a solution that offers **configurability** in terms of what cloud and storage I wanted to use and absolute **flexibility** in the mode of operation of the server (more on that ahead). In this article, I show you what I learned while developing such a solution and how you can adapt and use it for your own specific needs. Come, join me 🙋‍♂️


## Contents

* [A refresher on MLflow](#intro-mlflow)
* [Why is it a challenge?] (#challenges)
* [The outline for the solution](#outline)
* [Preparing the Helm Release](#prepare_release)
* [Setting up the Ingress](#ingress)
* [Test your MLflow server](#test)
* [YAYY! 🥳](#yay)

## A refresher on MLflow <a name="intro-mlflow"></a>

MLflow is an essential component to a great number of ML practitioners and ML Engineers. If you have ever felt the need to track your experiments, either in the form of logs and metadata or as artifacts from your pipeline runs, you know MLflow Tracking is the go-to tool! 😎

> Discussing what MLflow is and its pros and cons is beyond the scope of this article but there's wonderful [documentation](https://www.mlflow.org/docs/latest/tracking.html) available along with many more third-party resources online that you can use to get started!

Even with its popularity and utility, it is not always straightforward to have it deployed on your systems 😕, especially on the cloud, with **liberal customizations and authentication** built in.  

## Why is it a challenge? <a name="challenges"></a>

To make the situation and the challenges clearer, take a look at the number of [scenarios](https://www.mlflow.org/docs/latest/tracking.html#how-runs-and-artifacts-are-recorded) possible 😯 for your MLflow tracking server deployment. 

We won't jump into details as you can read the documentation page which explains them very well. Our concern here is to understand that each of these scenarios requires you to run the server with certain flags. As such, our solution should fundamentally allow easy customization.

In addition to being customizable in the cloud environment, we also need a way to control traffic around our service to protect the data that is stored on the server 🥷.


## The outline <a name="outline"></a>

There are three broad steps that I can think of when it comes to getting MLflow runnning on your cluster.

✅ Prepare a Helm chart for the tracking server that allows maximum customization.
✅ Set up an ingress controller and an ingress resource to control traffic to your MLflow service.
✅ Add authentication to secure your service and test if it works!

Now, for someone who hasn't dealt much with Kubernetes before, some of these tasks may sound too daunting. Fret not! In this article, I use a tool called Devtron to make Kubernetes deployments much simpler. Setting up and managing full Helm releases can be done very intuitively through the Devtron UI. 

However, that is **not** a pre-requisite. I have also explained how you could achieve the same outcomes in the traditional way. 💫 

The code in this article is taken from my [repository](https://github.com/wjayesh/mlflow-tracking-server) at GitHub. Feel free to experiment with it 🧪

## Preparing the Helm Release <a name="prepare_release"></a>

When it comes to deploying apps to your cluster, nothing comes close to using Helm, in terms of simplicity and ease of execution. Therefore, I will walk you through the entire journey, starting from scratch to having a full-fledged chart ready for sailing! ⛵
 
If that sounded too challenging or if you need only some basic customizations, I have a fast track just for you 😉

### Going All Gas, No Brakes!

<img width="50" height="50" src="https://media.giphy.com/media/e9CYq4R7AMwqk/giphy.gif">

In this process, we'll use a community Helm chart which allows setting simple configurations as our choice to deploy the MLflow tracking server. It's hosted [here](https://artifacthub.io/packages/helm/community-charts/mlflow)!

To deploy it to our cluster, we are going to make use of the Devtron Dashboard. However, you can also do a simple `helm install`. 

Setting up Devtron is simple - I'll let its [documentation](https://docs.devtron.ai/devtron/setup/install) guide you better. Once deployed, follow the steps below.

1. We'll first add the repository for this chart to our Devtron configuration. Go into the Global Configurations tab and select "Add Repository".
![Global Configuration page](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/gqf0egtyvb2evbbxjx1c.png) 

2. Fill in the details as above and save it.
3. Now you can go to the chart store and click "Deploy" on the MLflow chart, and that's it.

You can now skip directly to the next section 😎 on [setting up the ingress](#ingress)!

---

For folks who want a finer control on the tracking server and more customizations, follow along. This is going to be fun!

![A person arranging lots of cute kitten](https://media.giphy.com/media/a9wFQSc0oRQGI/giphy.gif)

What we're going to do can be split into three smaller steps.
1. Prepare a Docker Image of the MLflow tracking Server.
2. Create a Helm project.
3. Examine the manifests and create customizable values as necessary.

Let's jump into each of the them to get a grip of what's brewing ♨️.

### Prepare a Docker Image
The heart of this step lies in the Dockerfile. It's a file that acts as the recipe 🥗 for your image. You can define what goes into your image, what commands are run when this image is executed, the environment of this image and much more. 

To know what to put inside the Dockerfile, we have to think about what happens when you run an MLflow tracking server locally. To put it very simply, there's this main command that needs to be executed in your environment and it takes care of setting up the server and making it listen on a port.

```bash
mlflow server  \
--backend-store-uri "./mlflow/..." \
--default-artifact-root "./mlflow/..." \
--host "127.0.0.1" \
--port 5000
```
> **Note** 💁
> In order for this command to run inside a container environment, you need to choose the host as "0.0.0.0" so that it can run on all IP addresses (including the one where you are deploying this code).

Now, what we need is for this command to be run inside our container once the image starts, but with the added option to modify the value of any of its flags. Therefore we'll require a bash script that can take a variety of input flags and apply them accordingly to the base command.

I have created a script ✍️ that can take inputs for an artifact store and metadata store which you can use. It's a basic one but you can extend it to include more options like the [proxied access scenario](https://www.mlflow.org/docs/latest/tracking.html#scenario-5-mlflow-tracking-server-enabled-with-proxied-artifact-storage-access). 

We will now put this script inside our Dockerfile to be run when the container gets created.
```Dockerfile
FROM python:3.9

RUN pip install mlflow && \
    pip install awscli --upgrade --user && \
    pip install boto3==1.24.10

ENV PORT 5000

COPY scripts/run.sh /

WORKDIR /
RUN chmod +x run.sh

ENTRYPOINT ["./run.sh"]
```
Some observations 🤓:
- It uses a pre-built Python image as base so that we don't have to install pip and other tools ourselves.
- We then install the `mlflow` library along with `awscli` and `boto`. I'm using AWS since I plan to use S3 as the artifact store.
- You then copy over the script (which contains your command) into your container, make it an executable and set it to run once the container is spun up.

> **Note** 💁
> You may also want to add ENV variables as part of the Dockerfile for storing credentials for storage access, depending on what scenario you wish to deploy for.

#### Building and pushing the image
You can make use of a GitHub action to help build your Dockerfile into an image and push it to the registry of your choice. Take a look at my workflow file [here](https://github.com/wjayesh/mlflow-tracking-server/blob/main/.github/workflows/build.yml) for inspiration!

If you prefer to follow along with a video, I have just the thing you need 👇

{% embed https://youtu.be/oN2HhSZWcRs %}

### Create a Helm project

Once the image is built, we are ready to deploy our applications to a Kubernetes cluster. The fastest way to set up the deployment, the service, and all other related resources is to create a Helm chart. 

Run this command inside your desired directory to get some starter code for your chart.

```
helm create mlflow-tracking
```

![Showing file directory after helm create is run](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/xe52jsh665hviu29dyyt.png)

You will now see a whole bunch of files that get created in the directory. Our interest lies in the `deployment.yaml` file inside the `templates` folder, specifically in the following lines of code.

```yaml
containers:
        - name: {{ .Chart.Name }}
          securityContext:
            {{- toYaml .Values.securityContext | nindent 12 }}
          image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
          imagePullPolicy: {{ .Values.image.pullPolicy }}
          command: ["./run.sh"]
          args: ["-m", "{{ .Values.metadata_store }}", "-a", "{{ .Values.artifact_store }}"]
          
```
Some observations 🤓:
- We've used some variables to denote the image repository and tag. 
- Also modified are the `command` and the `args` parameters. 
  - The command is what will be executed when the container runs. This overrides what we had set in the Dockerfile.
  - The `args` parameter specifies the metadata store and the artifact store, in accordance with the the flags we've set while building our script.

#### Values

Where do you set these values, though? Important obvious question but there's a simple answer. 🥁🥁Inside the `values.yaml`! 

```yaml
image:
  repository: wjayesh/mlflow-tracking-server
  pullPolicy: Always
  # Overrides the image tag whose default is the chart appVersion.
  tag: latest

containers:
  port: 5000

metadata_store: ""    # add the metadata store URL
artifact_store: ""    # add the artifact store URL
```
As you can see, I've added my image name and tag here for the chart to access while deploying the server. 

> **Note** 💁
> You can also consider changing the service type inside `values.yaml` to `LoadBalancer` instead of the default `ClusterIP`, if you want to access it on a public IP address outside of your cluster environment.

That's all you need to do to set up your chart. We're now ready for launch! 🧑‍🚀

### Deploy to cluster 🚀

![A guy on computer with a thumbs up](https://media.giphy.com/media/8xgqLTTgWqHWU/giphy.gif)

Run this command to install the server to the cluster.

```bash
helm install <release-name> mlflow-tracking/ \
    --values mlflow-tracking/values.yaml
```
Now, verify if the MLflow service is up and running by using the following command ✅.

```bash
kubectl get svc
```

## Setting up the Ingress <a name="ingress"></a>

Once the MLflow tracking server is deployed, you can start making requests to its endpoint. All good, then why this extra step? Security!💂
We don't want to keep the server exposed to the outside world without any sort of checks in place since it holds critical information like your models, artifacts and logs. 

We'll set up something called an [Ingress](https://kubernetes.io/docs/concepts/services-networking/ingress/) in Kubernetes. It essentially controls the traffic, to and from a service. You can add rules, paths and authentication very easily by defining an ingress resource. 

Three more steps and I promise that will be the end 🤐😂

- Creating an Ingress Controller 
- Creating a secret which stores the credentials you want to apply to your server. 
- Finally, the ingress resource with your MLflow service as its backend. 

### Ingress Controller

An ingress controller is what actually accepts traffic from outside Kubernetes and sends it to the services, following the rules set in the ingress resources that you define. Naturally, we need to have a controller running first for the ingress definitions to work.

Let's head back to our Devtron dashboard (If you wish to not use it you can also do a simple Helm install)
1. Go into the Charts Store
![Charts Store page](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/n7535gyfiy58orxtouq8.png)
2. You can see the "NGINX" logo. That's the ingress controller we'll install for this example. Click on "Deploy"!
![NGINX Helm chart page](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/ushrvyt2hi0ru2ytr5xa.png)

The controller should be running inside the `ingress-nginx` namespace. Confirm by executing this command.

```bash
kubectl get pods -n ingress-nginx
```

### Creating a username and password for the server 🗝️

1. Head over to any website that allows creation of an `.htpasswd` file. Enter your desired username and password and you'll get get a similar output to the following. 
```
jayesh:$apr1$i4yl6mjs$aq3wqiZGbiJYqgypQeYGK/
```

2. Store the contents in a file called `auth`. It's important 🧐 that it be called **auth** otherwise NGINX would throw a 503 error. 

3. Now execute the following command to create a Kubernetes secret. 

```bash
kubectl create secret generic basic-auth --from-file=auth
```

### Creating the ingress rules

We can now finally create the rules to govern control to our MLflow service. 

Create an ingress resource by applying the following YAML.

```yaml
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
  name: mlflow-ingress
  annotations:
    # type of authentication
    nginx.ingress.kubernetes.io/auth-type: basic
    # your secret with user credentials
    nginx.ingress.kubernetes.io/auth-secret: basic-auth
    # message to display 
    nginx.ingress.kubernetes.io/auth-realm: 'Please authenticate first'
spec:
  rules:
  - http:
      paths:
      - path: /
        pathType: Prefix
        backend:
          service:
            name: mlflow-tracking
            port:
              number: 5000
  ingressClassName: nginx
``` 

We're now ready to test the MLflow tracking server with our username and password 🥳

## Test the MLflow server <a name="test"></a>

![finally](https://media.giphy.com/media/dsGcvvBPzXJ9KSnLPg/giphy.gif)

1. The folowing command can be used to get the tracking URL for the MLflow server. The EXTERNAL_IP field is the IP of the ingress controller and the path "/" is configured already (inside our ingress resource) to direct to the MLflow tracking server.

```bash
kubectl get service ingress-nginx-controller -n ingress-nginx
```

2. Run the following command, first without the "-u" flag and then with it to confirm that the MLflow tracking server is now authenticated!

```bash
curl -v <hostname> -u '<username>:<password>'
```

## Yay! 👏👏 <a name="yay"></a>

![pat on the back](https://media.giphy.com/media/8KkURDqxDWbLsMX99G/giphy.gif)

If you have come this far, pat yourself on the back 💪. You have successfully created a customized image for the MLflow tracking server, built a Helm chart for it, defined traffic access control for the service and tested it with your custom username and password combination. 

Although this is a fun way to get MLflow running on your cluster, it is certainly not the simplest 😂. In my next article, I'll dramatically reduce all of this work down to a **handful of commands** for you to run and have everything that we discussed running on your setup like new! Keep an eye out; I'm just as excited as you 👀

![shocked](https://media.giphy.com/media/vzoYZ7Ef9clbO/giphy.gif)

In the meantime, feel free to get in touch with me if you have any questions, or any ambitious ideas about what we should hack about next 👷
Here's everything you need: [bio.link/wjayesh](https://bio.link/wjayesh) 

See ya! And happy coding 🙋‍♂️

---

_The above blog is submitted as part of 'Devtron Blogathon 2022' - https://devtron.ai/
Check out Devtron's GitHub repo - https://github.com/devtron-labs/devtron/ and give a ⭐ to show your love & support.
Follow Devtron on LinkedIn - https://www.linkedin.com/company/devtron-labs/ and Twitter - https://twitter.com/DevtronL/, to keep yourself updated on this Open Source project._