An example MLOps code demo that deploys a model and visualizes the model predictions. This was a task for a job interview that I had to finish in 6 hours.

This project implements a simple web application that demonstrates an end-to-end MLOps workflow with the following components:

• Machine Learning Model: A linear regression model trained on the diabetes dataset from scikit-learn
• Backend API: A FastAPI-based REST API that serves the ML model
• Frontend Application: A Dash-based web application that interacts with the backend API
• Containerization: Docker Compose orchestrates the backend and frontend containers
• Visualization: Interactive Plotly visualizations of model predictions

The backend is built with FastAPI and provides the following endpoints:

• GET /: Health check endpoint
• GET /model/load: Loads the ML model into memory
• POST /model/predict: Accepts input data and returns model predictions

The ML model is a simple linear regression model trained on the diabetes dataset, using a single feature for simplicity and speed. The model is trained when first loaded and then used for all subsequent predictions.

The frontend is built with Dash and Dash Mantine Components, providing:

• A clean, responsive user interface
• A button to trigger the ML workflow
• Loading indicators to show processing status
• Interactive Plotly visualizations showing:
  • Training data points and model predictions
  • Test data points and model predictions

The application is containerized using Docker with separate containers for the frontend and backend services:

• Backend container: Runs the FastAPI application on port 8000
• Frontend container: Runs the Dash application on port 8050
• Docker Compose orchestrates the containers and sets up the necessary networking

To develop the application locally (mac arm and linux x86_64 only):

1. Install pixi

2. Start the backend server:

   pixi run backend-dev

3. In a separate terminal, start the frontend server:

   pixi run frontend-dev

4. Open your browser and navigate to http://localhost:8050 to access the application.

5. In another terminal, run the tests:

   pixi run tests

To develop the application locally using Docker:

1. Make sure you have Docker and Docker Compose installed on your system.

2. Build and start the containers:

   pixi run dcp-up --build --force-recreate

3. Open your browser and navigate to http://localhost:8050 to access the application.

4. run the tests:

   pixi run tests

5. Bring down the containers once you are done:

   pixi run dcp-down

You can deploy this application to Heroku using their Container Registry feature. Follow these steps to deploy both the backend and frontend services:

1. Install the Heroku CLI
2. Log in to Heroku:

   heroku login

3. Log in to the Heroku Container Registry:

   heroku container:login

1. Create a new Heroku app for the backend:

   heroku create mlops-demo-backend
   heroku stack:set container --app mlops-demo-backend

2. Build and push the backend Docker image: (see https://help.heroku.com/74S3XXMP/workaround-for-docker-push-error-405-method-not-allowed-using-buildkit-0-11 if you encounter an error, I had to disable containerd in my docker settings.)

   # due to limitations in the Heroku CLI, we need to rename the Dockerfile to Dockerfile.web
   mv Dockerfile.backend Dockerfile.web
   heroku container:push web --app mlops-demo-backend --context-path=. --recursive
   mv Dockerfile.web Dockerfile.backend

3. Release the container:

   heroku container:release web --app mlops-demo-backend

4. Ensure the web dyno is running:

   heroku ps:scale web=1 --app mlops-demo-backend

5. Verify the deployment:

   heroku open --app mlops-demo-backend

   You should see the the following json response:

   { "message": "Welcome to the ML model API" }

1. Create a new Heroku app for the frontend:

   heroku create mlops-demo-frontend
   heroku stack:set container --app mlops-demo-frontend

2. Set the backend URL environment variable to point to your deployed backend:

   heroku config:set BACKEND_URL=https://mlops-demo-backend-9b7fa0714fe4.herokuapp.com/ --app mlops-demo-frontend

3. Build and push the frontend Docker image:

   # due to limitations in the Heroku CLI, we need to rename the Dockerfile to Dockerfile.web
   mv Dockerfile.frontend Dockerfile.web
   heroku container:push web --app mlops-demo-frontend --context-path=. --recursive
   mv Dockerfile.web Dockerfile.frontend

4. Release the container:

   heroku container:release web --app mlops-demo-frontend

5. Ensure the web dyno is running:

   heroku ps:scale web=1 --app mlops-demo-frontend

6. Open the frontend application:

   heroku open --app mlops-demo-frontend

7. To view logs for debugging:

   heroku logs --tail --app mlops-demo-backend
   heroku logs --tail --app mlops-demo-frontend

View the application at https://mlops-demo-frontend-75cad5f56be2.herokuapp.com/

BACKEND_URL=https://mlops-demo-backend-9b7fa0714fe4.herokuapp.com/ pixi run tests -k e2e_backend

While the current implementation meets all the requirements, there are several areas that could be improved:

1. Model Persistence: Currently, the model is retrained each time the backend starts. Using a model registry (e.g., MLflow, AWS SageMaker, Azure Machine Learning) would improve startup time and allow for model versioning.

2. Error Handling: The error handling is basic and could be enhanced with more detailed error messages and graceful fallbacks when services are unavailable.

3. Testing Coverage: While there are test files in place, expanding test coverage would improve reliability, especially for edge cases and integration tests.

4. Monitoring and Logging: I skipped this due to the time constraint. Logging could be implemented to track model performance metrics, user interactions, and system health over time.

5. UI Enhancements: The current UI is functional but minimal. After a discussion with stakeholders, we can decide how to improve the UI.

6. Cloud Deployment: I chose heroku for the deployment because it was simple but I would have liked to deploy on AWS Fargate which offers serverless compute for containers.

7. CI/CD Pipeline: Adding a CI/CD pipeline would automate testing and deployment, ensuring consistent quality and easier updates. You may look at the CI I setup for another project: idiap/gridtk#1

8. Infrastructure as Code: The project can be deployed using infrastructure as code. For example, using Terraform to deploy the infrastructure on AWS.

9. Security: The application doesn't implement any authentication or authorization mechanisms, which would be necessary for a production environment.

10. Frontend Workers: The preparation of the data for the frontend could be a time consuming task. Using workers (e.g., Celery) to offload this task could improve the user experience.

• The frontend only depends on scikit-learn to be able to load the data. This will be different if the frontend is loading data from a database or another source.
• The whole development process is automated with pixi, but due to the time limit I chose to install dependencies using uv in the Dockerfiles. Doing this in pixi is possible too: https://pixi.sh/latest/deployment/container/ but it needed a bit more work. This would have avoided the need to declare the dependencies in the Dockerfiles.
• The demo is based on this guide: https://scikit-learn.org/stable/auto_examples/linear_model/plot_ols.html