Build and Deploy an ML Model App using Streamlit, Docker and GKE
Introduction
You‘ve built a great machine learning model – now how do you get it in the hands of users? Deploying ML models as web applications is a common requirement, but it can be challenging, especially when you need a user-friendly front-end and want to leverage containerization and cloud platforms for scalability and maintainability.
In this step-by-step guide, we‘ll walk through the process of building a complete ML model app and deploying it using some of the most popular and powerful tools available today:
- Streamlit: A Python library for quickly creating interactive web apps for data science and ML
- FastAPI: A modern Python web framework for building APIs
- Docker: A platform for developing, shipping, and running applications using containers
- Google Kubernetes Engine (GKE): A managed environment for deploying, scaling, and managing containerized applications using Google Cloud
We‘ll use a real example of predicting employee churn and cover the end-to-end workflow, including:
- Building the predictive model
- Wrapping the model in a REST API with FastAPI
- Creating a web app UI for the model with Streamlit
- Containerizing the Streamlit app and API using Docker
- Deploying and orchestrating the containers in GKE
- Integrating and communicating between the services
- Updating the deployed app
Whether you‘re a data scientist, ML engineer, or software developer, by the end of this article you‘ll have a solid understanding of the architecture and tools required to deploy ML apps into production using industry best practices. Let‘s dive in!
The Machine Learning Model
For our example, we‘ll tackle the problem of predicting employee churn – a critical issue for many companies. We have a dataset of employee records including features like satisfaction level, number of projects, average monthly hours, and more. Our goal is to build a binary classification model to predict if an employee is likely to leave the company or not.
Here‘s a sample of the code to build the model using scikit-learn:
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score
df = pd.read_csv(‘employee_churn.csv‘)
y = df[‘left‘]
X = df.drop(‘left‘, axis=1)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
model = RandomForestClassifier(n_estimators=100)
model.fit(X_train, y_train)
y_pred = model.predict(X_test)
roc_auc = roc_auc_score(y_test, y_pred)
print(f‘The ROC AUC on test data is: {roc_auc:.3f}‘)
After some experimentation and hyperparameter tuning, we arrive at a Random Forest model that achieves an ROC AUC of 0.95 on the test set, which is quite good. We‘re now ready to deploy this model, so users can input employee data and get a prediction.
The Prediction API
To make the model available to serve predictions, we‘ll create a REST API using FastAPI. This provides a standard interface for other services to request predictions by sending employee data to an HTTP endpoint.
Here‘s the key parts of the FastAPI code:
from fastapi import FastAPI
from pydantic import BaseModel
from fastapi.encoders import jsonable_encoder
import joblib
# Initialize API
app = FastAPI()
# Load model
model = joblib.load(‘employee_churn_model.pkl‘)
# Define input data schema
class EmployeeData(BaseModel):
satisfaction_level: float
last_evaluation: float
number_project: int
average_montly_hours: int
years_at_company: int
work_accident: int
promotion_last_5years: int
department: str
salary: str
# Define prediction endpoint
@app.post(‘/predict‘)
async def predict(employee: EmployeeData):
input_data = jsonable_encoder(employee)
prediction = model.predict_proba([[
input_data[‘satisfaction_level‘],
input_data[‘last_evaluation‘],
input_data[‘number_project‘],
input_data[‘average_montly_hours‘],
input_data[‘years_at_company‘],
input_data[‘work_accident‘],
input_data[‘promotion_last_5years‘],
input_data[‘department‘],
input_data[‘salary‘]
]])[:,1]
return {‘churn_probability‘: prediction[0]}
This API has a single endpoint /predict that accepts a POST request containing the employee data. It uses pydantic to define and validate the schema of the input data, which is a best practice for APIs. The employee data is then passed to the model to generate a prediction of the probability of churn, which is returned in the response.
To run this API locally, we can simply use the uvicorn server:
uvicorn app:app --reload
We can then test the API by sending a request using a tool like curl or Postman.
The Streamlit Web App
To enable users to easily interact with our model, we‘ll create a simple web app using Streamlit. Streamlit is an excellent choice for data science and ML apps because it allows you to create an interactive UI directly from Python code, without needing to know web development.
Here‘s the code for our Streamlit app:
import streamlit as st
import requests
import json
st.title(‘Employee Churn Predictor‘)
st.sidebar.header(‘Employee Data‘)
satisfaction_level = st.sidebar.slider(‘Satisfaction Level‘, 0.0, 1.0, 0.5)
last_evaluation = st.sidebar.slider(‘Last Evaluation‘, 0.0, 1.0, 0.5)
number_project = st.sidebar.slider(‘Number of Projects‘, 0, 10, 5)
average_montly_hours = st.sidebar.slider(‘Average Monthly Hours‘, 0, 500, 250)
years_at_company = st.sidebar.slider(‘Years at Company‘, 0, 20, 10)
work_accident = st.sidebar.radio(‘Had Work Accident‘, [0, 1])
promotion_last_5years = st.sidebar.radio(‘Promoted in Last 5 Years‘, [0, 1])
department = st.sidebar.selectbox(‘Department‘, [‘sales‘, ‘IT‘, ‘hr‘, ‘technical‘, ‘support‘, ‘management‘])
salary = st.sidebar.selectbox(‘Salary‘, [‘low‘, ‘medium‘, ‘high‘])
data = {
‘satisfaction_level‘: satisfaction_level,
‘last_evaluation‘: last_evaluation,
‘number_project‘: number_project,
‘average_montly_hours‘: average_montly_hours,
‘years_at_company‘: years_at_company,
‘work_accident‘: work_accident,
‘promotion_last_5years‘: promotion_last_5years,
‘department‘: department,
‘salary‘: salary
}
if st.button(‘Predict‘):
response = requests.post(‘http://fastapi:8000/predict‘, json=data)
prediction = response.json()
churn_probability = prediction[‘churn_probability‘]
st.write(f‘The probability of churn for this employee is {churn_probability:.2%}‘)
This app has a sidebar with input widgets for each feature in the employee data schema, matching the input to the API. When the user clicks the "Predict" button, the app sends the data to the FastAPI server to get a prediction and displays the result.
To run the app locally:
streamlit run app.py
Containerizing with Docker
To deploy our Streamlit app and FastAPI server, we‘ll containerize them using Docker. This ensures they have all the required dependencies, can be consistently deployed across environments, and can scale independently.
We‘ll create two containers:
- The FastAPI container with the model and API server
- The Streamlit container with the web app UI
Here are the two Dockerfiles:
FastAPI Dockerfile:
FROM python:3.8
WORKDIR /app
COPY requirements.txt .
RUN pip install -r requirements.txt
COPY . .
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]
Streamlit Dockerfile:
FROM python:3.8
WORKDIR /app
COPY requirements.txt .
RUN pip install -r requirements.txt
COPY . .
CMD ["streamlit", "run", "app.py"]
And here‘s the docker-compose file to build and run both containers:
version: "3"
services:
fastapi:
build: ./fastapi
container_name: fastapi
ports:
- 8000:8000
streamlit:
build: ./streamlit
container_name: streamlit
ports:
- 8501:8501
To run this locally:
docker-compose up -d
The FastAPI server will be available at localhost:8000 and the Streamlit app at localhost:8501.
Deploying to Google Kubernetes Engine
The final step is to deploy our containers to a Kubernetes cluster in GKE.
First, create a new project in Google Cloud Console and enable billing. Then create a new Kubernetes cluster in GKE, either using the GUI or gcloud CLI tool.
Next, we need to create Kubernetes manifests to define the deployments and services for our FastAPI and Streamlit containers.
Here‘s the manifest for the FastAPI deployment and service:
apiVersion: apps/v1
kind: Deployment
metadata:
name: fastapi
spec:
replicas: 2
selector:
matchLabels:
app: fastapi
template:
metadata:
labels:
app: fastapi
spec:
containers:
- name: fastapi
image: gcr.io/my-project/fastapi:v1
ports:
- containerPort: 8000
---
apiVersion: v1
kind: Service
metadata:
name: fastapi
spec:
selector:
app: fastapi
ports:
- protocol: TCP
port: 8000
targetPort: 8000
And the Streamlit manifest:
apiVersion: apps/v1
kind: Deployment
metadata:
name: streamlit
spec:
replicas: 1
selector:
matchLabels:
app: streamlit
template:
metadata:
labels:
app: streamlit
spec:
containers:
- name: streamlit
image: gcr.io/my-project/streamlit:v1
ports:
- containerPort: 8501
---
apiVersion: v1
kind: Service
metadata:
name: streamlit
spec:
type: LoadBalancer
selector:
app: streamlit
ports:
- protocol: TCP
port: 80
targetPort: 8501
You‘ll notice a few key points:
- The FastAPI deployment has multiple replicas for scaling and availability
- The Streamlit service is exposed externally using a LoadBalancer
- The container images are stored in Google Container Registry
After pushing your container images to GCR, you can deploy the manifests to your GKE cluster:
kubectl apply -f fastapi.yaml
kubectl apply -f streamlit.yaml
Once the services are running, you‘ll be able to access the Streamlit web app using the external IP of its LoadBalancer service.
However, the Streamlit app still needs to communicate with the FastAPI server in GKE. To enable this, update the Streamlit code to use the internal DNS name of the FastAPI service instead of localhost:
response = requests.post(‘http://fastapi:8000/predict‘, json=data)
This works because GKE automatically provides service discovery within the cluster. Each service can be accessed using its name and default port.
Finally, anytime you make changes, you can update the deployments in GKE by building new container images, pushing them to GCR, and running:
kubectl rollout restart deployment fastapi
kubectl rollout restart deployment streamlit
Conclusion
Let‘s review what we‘ve accomplished:
- Built an ML model to predict employee churn
- Created a FastAPI server to generate predictions from the model
- Developed a user-friendly web interface for the model using Streamlit
- Containerized the FastAPI server and Streamlit app using Docker
- Deployed the containers to Google Kubernetes Engine
- Integrated the services to communicate within the cluster
- Walked through the process for updating the app
We now have a scalable, maintainable, and reproducible architecture for deploying a machine learning application, using some of the most powerful open-source tools available.
This workflow has several key benefits:
- Streamlit enables rapid prototyping of web apps for ML without in-depth web dev knowledge
- FastAPI provides an easy yet robust way to serve models through APIs
- Docker containers encapsulate dependencies and allow the app components to run consistently anywhere
- Kubernetes enables declarative deployment, scaling, and management of the application
- GKE provides a managed platform for running the app in the cloud without managing infrastructure
That said, there are other tools and approaches you can consider as well, such as Flask or Django for the web framework, AWS or Azure for the cloud platform, or serverless compute options. The right choice depends on your specific needs and constraints.
The key lessons to take away are the importance of:
- Modularity in your application design
- Using containers for portability and consistency
- Leveraging managed services to reduce overhead
- Enabling communication between components through APIs and service discovery
I hope this guide has been helpful in demonstrating a practical example of productionizing ML using Streamlit, Docker, and GKE. You can find all the code for this project in the links below. Feel free to use it as a starting point for your own applications!