Grigor Khachatryan

Director of Engineering, Platform | Los Angeles, CA


How to Deploy Machine Learning Models Using Docker: A Step-by-Step Guide

Published October 6, 2024

Deploying machine learning models can be a complex endeavor, especially when aiming for consistency across various environments. Docker simplifies this process by allowing you to package your models and their dependencies into portable containers. In this article, we’ll explore how to deploy machine learning models using Docker, making it accessible for both newcomers and seasoned developers.

Why Use Docker for Machine Learning Deployment?

What Is Docker?

Docker is an open-source platform designed to automate the deployment of applications within lightweight, portable containers. These containers bundle your application with all its dependencies, ensuring it runs seamlessly in any environment.

Benefits of Using Docker

  • Portability: Run your containerized applications anywhere Docker is installed.
  • Consistency: Eliminate environment-related issues by packaging dependencies together.
  • Scalability: Easily scale applications horizontally by deploying multiple container instances.

Getting Started with Docker

Installing Docker

First, install Docker on your system by downloading Docker Desktop from the official website. Follow the installation instructions specific to your operating system.

Basic Docker Commands

  • docker build: Build a Docker image from a Dockerfile.
  • docker run: Run a container from a Docker image.
  • docker images: List all Docker images on your system.
  • docker ps: List all running Docker containers.

Containerizing a Machine Learning Model

We’ll walk through containerizing a simple machine learning model built with Python and scikit-learn.

Step 1: Create a Simple Machine Learning Model

Create a Python script named model.py:

# model.py  
from sklearn.datasets import load_iris  
from sklearn.ensemble import RandomForestClassifier  
import pickle  
  
# Load dataset  
iris = load_iris()  
X, y = iris.data, iris.target  
  
# Train model  
clf = RandomForestClassifier()  
clf.fit(X, y)  
  
# Save the model  
with open('model.pkl', 'wb') as f:  
    pickle.dump(clf, f)

Step 2: Write a Flask Application to Serve the Model

Create an app.py file:

# app.py  
from flask import Flask, request, jsonify  
import pickle  
  
# Load the model  
with open('model.pkl', 'rb') as f:  
    model = pickle.load(f)  
  
app = Flask(__name__)  
  
@app.route('/predict', methods=['POST'])  
def predict():  
    data = request.get_json(force=True)  
    prediction = model.predict([data['input']])  
    return jsonify({'prediction': prediction.tolist()})  
  
if __name__ == '__main__':  
    app.run(host='0.0.0.0')

Step 3: Write a Dockerfile

Create a file named Dockerfile:

# Use an official Python runtime as a parent image  
FROM python:3.13-slim  
  
# Set the working directory  
WORKDIR /app  
  
# Copy current directory contents into the container  
COPY . /app  
  
# Install dependencies  
RUN pip install --no-cache-dir flask scikit-learn  
  
# Expose port 5000  
EXPOSE 5000  
  
# Set environment variables  
ENV FLASK_APP=app.py  
  
# Run the application  
CMD ["flask", "run", "--host=0.0.0.0"]

Understanding the Dockerfile

  • FROM: Specifies the base image.
  • WORKDIR: Sets the working directory inside the container.
  • COPY: Copies files from the host to the container.
  • RUN: Executes commands during image build.
  • EXPOSE: Informs Docker which port the container listens on.
  • ENV: Sets environment variables.
  • CMD: Specifies the default command to run when starting the container.

Building and Running the Docker Image

Step 1: Build the Docker Image

In your terminal, run:

docker build -t ml-model:latest .

This command builds an image named ml-model using the Dockerfile in the current directory.

Step 2: Run the Docker Container

Run the container with:

docker run -p 5000:5000 ml-model:latest

This maps port 5000 on your host to port 5000 in the container.

Step 3: Test the Deployed Model

In a new terminal window, send a test request:

curl -X POST -H "Content-Type: application/json" \  
-d '{"input": [5.1, 3.5, 1.4, 0.2]}' \  
http://localhost:5000/predict

You should receive a JSON response with the prediction.

Best Practices

Optimize Docker Images

  • Use a .dockerignore File: Exclude unnecessary files to reduce image size.
  • Minimize Layers: Combine commands when possible to reduce the number of layers.
  • Choose Lightweight Base Images: Use slim or Alpine versions of base images.

Docker Best Practices: Images

Security Considerations

  • Run as Non-Root User: Avoid running applications as the root user inside containers.
  • Keep Dependencies Updated: Regularly update base images and dependencies to patch vulnerabilities.

Docker Best Practices: Security

Conclusion

Docker streamlines the deployment of machine learning models by providing a consistent and portable environment. By containerizing your models, you can ensure they run reliably in any setting, from local machines to cloud servers. Whether you’re new to Docker or looking to enhance your deployment process, integrating Docker into your workflow offers significant benefits.