Deploying LoRA Optimised BERT as a FastApi service on GKE | ML Engineering & MLOps
Hi All
Alright today on "Bored MLE" we're doing some inference, using a LoRA optimized BERT model on GKE. Now there's a lot more to MLOps than just this (like cluster level observability for example), but for this example I'm keeping it brief.
The full code is available on GitHub with added minikube deployment instructions (I only cover the Cloud deployment here).
I assume familiarity with Python, PyTorch, K8s, FastApi, Minikube and GKE. Let's get on with it.
View the full FastApi service source below, also available here:
Let's breakdown the above code, block by block.
0. Install Dependencies:
pip install fastapi uvicorn torch transformers peft accelerate
1. Imports
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from peft import PeftModel, PeftConfig
* fastapi: Frame for building the API
* FastApi: Main class to create the app.
* HTTPException: For handling HTTP errors.
* pydantic: Data validation and settings management.
* BaseModel: Base class for request/response models.
* typing: For type hints (eg. List)
* torch: PyTorch for tensor operations and GPU support.
* AutoModelForSequenceClassification: For classification tasks.
* AutoTokenizer: For tokenizing input text.
* peft: Library for Parameter-Effecient Fine-Tuning (LoRA).
* PeftModel: For loading LoRA-optimized models.
* PeftConfig: For loading LoRA configuration.
2. Model and Tokenizer Loading
base_model_name = "bert-base-uncased"
peft_model_id = "./lora-bert-finetuned" # Path to your LoRA-optimized model
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
# Load base model
base_model = AutoModelForSequenceClassification.from_pretrained(
base_model_name,
num_labels=2, # Adjust based on your task
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)
# Load LoRA weights
model = PeftModel.from_pretrained(base_model, peft_model_id)
* base_model_name: Name of the base BERT model (bert-base-uncased).
* peft_model_id: Path to ther directory containing your LoRA-optimized model.
* Tokenizer:
* Loads the tokenizer for the base BERT model.
* Base Model:
* Loads the base BERT model for sequence classification.
* num_labels=2: Adjust based on your task (eg. binary classification).
* Uses float16 for GPU or float32 for CPU to optimize memory usage.
* LoRA Model:
* Loads the LoRA-optimised weights into the base model using PeftModel.from_pretrained.
3. Device Setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model.eval()
* Device Selection:
* Uses GPU(cuda) if available, otherwise falls back to CPU
* Model Setup:
* Moves the model to the selected device.
* Sets the model to evaluation mode (model.eval()), disabling dropout and batch norm layers.
4. FastApi App Initialization
app = FastAPI()
* Initializes the FastAPI application.
5. Request Model (Pydantic)
class TextRequest(BaseModel):
texts: List[str]
max_length: int = 128
* Defines the request body schema using Pydantic's BaseModel.
* texts: List of input strings to process.
* max_length: Maximum token length for truncation (default: 128).
6. Prediction Endpoint
@app.post("/predict")
async def predict(request: TextRequest):
try:
inputs = tokenizer(
request.texts,
padding=True,
truncation=True,
max_length=request.max_length,
return_tensors="pt"
).to(device)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
predictions = torch.argmax(logits, dim=-1).cpu().numpy().tolist()
return {"predictions": predictions}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
* Endpoint: /predict (POST method).
* Request Handling:
* tokenizes the input text with padding, truncation, and dynamic max length.
* Moves the tokenized inputs to the same device as the model.
* Inference:
* Disables gradient computation (torch.no_grad()) for efficiency.
* Runs the model on the tokenized inputs.
* Extracts logits and converts them to predictions using argmax.
* Response:
* Returns predictions as a JSON response.
* Error Handling:
* Catches exceptions and returns a 500 error with the exception message.
7. Health Check Endpoint
@app.get("/health")
async def health_check():
return {"status": "healthy"}
* Endpoint: /health (GET method).
* Purpose: Simple endpoint to check if the service is running.
* Response: Returns a JSON object with {"status": "healthy"}.
The torch.nograd() context manage disables gradient computation, which is essential for inference to save memory and compute. Now let's have a look at the GKE ML Ops implementation.
1. Deployment Metadata
apiVersion: apps/v1
kind: Deployment
metadata:
name: lora-bert-deployment
* apiVersion: Specifies the Kubernetes API version (apps/v1 for Deployment objects).
* kind: Defines the type of K8s resource (Deployment)
* metadata:
* name: The name of the Deployment (lora-bert-deployment).
2. Deployment Spec
spec:
replicas: 3
selector:
matchLabels:
app: lora-bert
spec: Defines the desired state of the Deployment.
* replicas: Number of identical Pods to maintain (3 in this case).
* selector: Used to identify which Pods belong to this deployment.
* matchLabels: Selects Pods with the label app: lora-bert.
3. Pod Template
template:
metadata:
labels:
app: lora-bert
* template: Defines the Pod template that the Deployment will use to create Pods.
* metadata.labels: Labels applied to the Pods (app: lora-bert).
4. Container Specification
spec:
containers:
- name: lora-bert
image: gcr.io/${PROJECT_ID}/lora-bert-fastapi:latest
ports:
- containerPort: 8000
* containers: List of containers that will run in the Pod.
* Name of the container (lora-bert)
* image: Docker image to use for the container. Replace ${PROJECT_ID} with your actual project ID.
* ports:
* containerPort: Port exposed by the container (8000, matching the FastAPI app).
5. Resource Limits and Requests
resources:
limits:
cpu: "2"
memory: "4Gi"
nvidia.com/gpu: 1
requests:
cpu: "1"
memory: "2Gi"
* resources: Defines the compute resources allocated to the container.
* limits: Maximum resources the container can use.
* cpu: 2 CPU cores.
* memory: 4GiB of RAM.
* nvidia.com/gpu: 1 GPU (if using GPU acceleration).
* requests: Minimum resources guaranteed to the container.
* cpu: 1 CPU core.
* memory: 2 GiB of RAM.
6. Liveness Probe
livenessProbe:
httpGet:
path: /health
port: 8000
initialDelaySeconds: 30
periodSeconds: 10
* livenessProbe: Checks if the container is running and healthy.
* httpGet: Performa an HTTP GET request to the /health endpoint.
* initialDelaySeconds: Wait 30 seconds before starting the probe.
* periodSeconds: Checks every 30 seconds.
7. Readiness Probe
readinessProbe:
httpGet:
path: /health
port: 8000
initialDelaySeconds: 5
periodSeconds: 5
* readinessProbe: Checks if the container is ready to receive traffic.
* httpGet: Performs an HTTP GET request to the /health endpoint
* initialDelaySeconds: Wait 5 seconds before starting the probe.
* periodSeconds: Checks every 5 seconds.
8. Service Definition
---
apiVersion: v1
kind: Service
metadata:
name: lora-bert-service
spec:
selector:
app: lora-bert
ports:
- protocol: TCP
port: 80
targetPort: 8000
type: LoadBalancer
* --- : Seperator for multiple Kubernetes (K8s) resources in a single file.
* apiVersion: Specifies the K8s API version (v1 for Service objects)
* kind: Defines the type of K8s resource (Service).
* metadata:
* name: he name of the Service (lora-bert-service)
* spec:
* selector: Selects Pods with the label app: lora-bert
* ports:
* protocol: Protocol used (TCP)
* port: Port exposed by the Service (80)
* targetPort: Port on the Pods to forward traffic to (8000)
* type: Type of service (LoadBalancer), which provisions an external IP for the Service.
9. Deployment Script
docker build -t gcr.io/YOUR_PROJECT_ID/lora-bert-fastapi:latest .
gcloud auth configure-docker
docker push gcr.io/YOUR_PROJECT_ID/lora-bert-fastapi:latest
gcloud container clusters create lora-bert-cluster \
--num-nodes=3 \
--machine-type=n1-standard-4 \
--accelerator=type=nvidia-tesla-t4,count=1 \
--zone=us-central1-a
kubectl apply -f gke-deployment.yaml
kubectl get service lora-bert-service
* Builds and deploys Docker image to Google Container Registry
* Creates a GKE cluster with GPU nodes.
* Applies the K8s deployment yaml.
* Gets the deployment's external API.
The nvidia.com/gpu: 1 resource limit enables GPU acceleration. Ensure your GKE cluster has GPU nodes available. Full source available here. Please read the docs in the project for the full set of instructions, including minikube for local deployment (which is not covered here). That's all for now folks.
Later
- Ed
Comments
Post a Comment