LitServe Tutorial: Fast and Easy AI Model Serving Framework

Introduction

LitServe is an open-source framework from Lightning AI designed for serving AI/ML models with high performance and ease of use. Built on top of FastAPI, LitServe provides enterprise-grade features such as batching, streaming, GPU autoscaling, and multi-model endpoints without complex configuration.

Unlike other serving frameworks that require extensive boilerplate code and infrastructure setup, LitServe lets you turn any AI model into a production-ready API with just a few lines of code. The framework supports all model types including LLMs, computer vision, NLP, audio, and classic scikit-learn models.

In this tutorial, we will cover everything from installation, creating a simple API, to advanced features like batching, streaming, and authentication.

Why LitServe?

Before diving into implementation, here is why LitServe stands out:

• High Performance: Up to 2x faster than plain FastAPI thanks to internal optimizations
• Minimal Boilerplate: Just define setup(), decoderequest(), predict(), and encoderesponse()
• Automatic Batching: Combines multiple requests for higher GPU throughput
• Streaming Response: Native support for streaming output (ideal for LLMs)
• Multi-Model: Single server can serve multiple models simultaneously
• GPU Management: Automatic GPU allocation and multi-worker handling
• OpenAPI Docs: Auto-generated API documentation (Swagger UI)

Installation

Prerequisites

Make sure you have Python 3.8 or later. Using a virtual environment is recommended.

python -m venv litserve-env
source litserve-env/bin/activate  # Linux/Mac
or
litserve-env\Scripts\activate     # Windows

Install LitServe

pip install litserve

For additional features, you can install optional dependencies:

# For all features
pip install litserve[all]

For PyTorch models
pip install torch torchvision

For Hugging Face models
pip install transformers

For scikit-learn models
pip install scikit-learn

Verify installation:

import litserve as ls
print(ls.version)

Basic Usage: Creating Your First API

Core Concept: LitAPI

The core of LitServe is the LitAPI class. You need to implement at minimum 4 methods:

Example 1: Simple API with Scikit-Learn

Let's build an API for an iris classification model using scikit-learn.

# server.py
import litserve as ls
from sklearn.datasets import loadiris

from sklearn.ensemble import RandomForestClassifier
import numpy as np

class IrisAPI(ls.LitAPI):
    def setup(self, device):
        """Load and train model at server startup."""
        iris = loadiris()
        self.model = RandomForestClassifier(nestimators=100, randomstate=42)
        self.model.fit(iris.data, iris.target)
        self.classnames = iris.targetnames.tolist()

    def decoderequest(self, request):

        """Parse input from HTTP request."""
        return np.array(request["features"]).reshape(1, -1)

    def predict(self, x):
        """Run prediction."""
        prediction = self.model.predict(x)
        probability = self.model.predictproba(x)
        return {
            "classidx": int(prediction[0]),

            "probabilities": probability[0].tolist()
        }

    def encoderesponse(self, output):
        """Format response for client."""
        return {
            "predictedclass": self.classnames[output["classidx"]],

            "probabilities": {
                name: round(prob, 4)
                for name, prob in zip(self.classnames, output["probabilities"])
            }
        }

if name == "main":
    api = IrisAPI()
    server = ls.LitServer(api)
    server.run(port=8000)

Run the server:

python server.py

Test with curl:

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

Response:

{
  "predictedclass": "setosa",

  "probabilities": {
    "setosa": 0.9800,
    "versicolor": 0.0150,
    "virginica": 0.0050
  }
}

# visionserver.py
import litserve as ls
import torch
from torchvision import transforms, models
from PIL import Image
import io
import base64

class ImageClassifierAPI(ls.LitAPI):
    def setup(self, device):
        """Load pretrained ResNet model."""
        self.device = device
        self.model = models.resnet50(weights=models.ResNet50Weights.DEFAULT)

        self.model.to(device)
        self.model.eval()

        self.transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            ),
        ])

        # Load ImageNet labels
        self.labels = models.ResNet50Weights.DEFAULT.meta["categories"]

    def decoderequest(self, request):

        """Decode base64 image from request."""
        imagebytes = base64.b64decode(request["image"])
        image = Image.open(io.BytesIO(imagebytes)).convert("RGB")

        return self.transform(image).unsqueeze(0).to(self.device)

    def predict(self, x):
        """Run inference."""
        with torch.nograd():
            outputs = self.model(x)
            probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
        return probabilities

    def encoderesponse(self, output):

        """Return top-5 predictions."""
        top5prob, top5idx = torch.topk(output, 5)

        results = []
        for i in range(5):
            results.append({
                "label": self.labels[top5idx[i].item()],
                "confidence": round(top5prob[i].item(), 4)

            })
        return {"predictions": results}

if name == "main":
    api = ImageClassifierAPI()
    server = ls.LitServer(api, accelerator="auto")
    server.run(port=8000)

import litserve as ls
import torch
from transformers import AutoTokenizer, AutoModel

class EmbeddingAPI(ls.LitAPI):
    def setup(self, device):
        self.device = device
        modelname = "sentence-transformers/all-MiniLM-L6-v2"
        self.tokenizer = AutoTokenizer.frompretrained(modelname)
        self.model = AutoModel.frompretrained(modelname).to(device)
        self.model.eval()

    def decoderequest(self, request):

        return request["text"]

    def predict(self, batch):
        """With batching, batch is a list of all inputs."""
        encoded = self.tokenizer(
            batch,
            padding=True,
            truncation=True,
            maxlength=512,
            returntensors="pt"

        ).to(self.device)

        with torch.nograd():
            outputs = self.model(*encoded)


        # Mean pooling
        attentionmask = encoded["attentionmask"]
        tokenembeddings = outputs.lasthiddenstate

        inputmask = attentionmask.unsqueeze(-1).expand(tokenembeddings.size()).float()
        embeddings = torch.sum(tokenembeddings  inputmask, 1) / torch.clamp(
            inputmask.sum(1), min=1e-9

        )

        return embeddings.cpu().tolist()

    def encoderesponse(self, output):
        return {"embedding": output}

if name == "main":
    api = EmbeddingAPI()
    server = ls.LitServer(
        api,
        accelerator="auto",
        maxbatchsize=32,      # Maximum 32 requests per batch
        batchtimeout=0.05,     # Wait 50ms to collect batch

    )
    server.run(port=8000)

import litserve as ls
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

class LLMStreamingAPI(ls.LitAPI):
    def setup(self, device):
        self.device = device
        modelname = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
        self.tokenizer = AutoTokenizer.frompretrained(modelname)
        self.model = AutoModelForCausalLM.frompretrained(

            modelname,
            torchdtype=torch.float16

        ).to(device)

    def decoderequest(self, request):
        prompt = request["prompt"]
        maxtokens = request.get("maxtokens", 256)
        return {"prompt": prompt, "maxtokens": maxtokens}

    def predict(self, x):
        """Generator function for streaming."""
        inputs = self.tokenizer(
            x["prompt"],
            returntensors="pt"

        ).to(self.device)

        generatedids = inputs["inputids"]


        for  in range(x["maxtokens"]):

            with torch.nograd():
                outputs = self.model(generatedids)


            nexttokenlogits = outputs.logits[:, -1, :]

            nexttokenid = torch.argmax(nexttokenlogits, dim=-1, keepdim=True)


            if nexttokenid.item() == self.tokenizer.eostokenid:

                break

            generatedids = torch.cat([generatedids, nexttokenid], dim=-1)

            tokentext = self.tokenizer.decode(
                nexttokenid[0],
                skipspecialtokens=True
            )

            yield tokentext


    def encoderesponse(self, output):
        for token in output:
            yield {"token": token}

if name == "main":
    api = LLMStreamingAPI()
    server = ls.LitServer(api, stream=True, accelerator="auto")
    server.run(port=8000)

Streaming client:

import requests

response = requests.post(
    "http://localhost:8000/predict",
    json={"prompt": "Explain quantum computing in simple terms:", "maxtokens": 100},

    stream=True
)

for chunk in response.itercontent(chunksize=None):

    if chunk:
        print(chunk.decode(), end="", flush=True)

import litserve as ls
from transformers import pipeline

class SentimentAPI(ls.LitAPI):
    def setup(self, device):
        self.pipe = pipeline(
            "sentiment-analysis",
            model="distilbert-base-uncased-finetuned-sst-2-english",
            device=device
        )

    def decoderequest(self, request):
        return request["text"]

    def predict(self, x):
        return self.pipe(x)[0]

    def encoderesponse(self, output):

        return {"label": output["label"], "score": round(output["score"], 4)}


class SummarizationAPI(ls.LitAPI):
    def setup(self, device):
        self.pipe = pipeline(
            "summarization",
            model="sshleifer/distilbart-cnn-12-6",
            device=device
        )

    def decoderequest(self, request):
        return {
            "text": request["text"],
            "maxlength": request.get("maxlength", 130)
        }

    def predict(self, x):
        return self.pipe(
            x["text"],
            maxlength=x["maxlength"],
            minlength=30

        )[0]

    def encoderesponse(self, output):
        return {"summary": output["summarytext"]}



if name == "main":
    server = ls.LitServer(
        [SentimentAPI(), SummarizationAPI()],
        apipath=["/sentiment", "/summarize"],
        accelerator="auto"
    )
    server.run(port=8000)

Authentication and Middleware

Protect your API with authentication.

import litserve as ls

class ProtectedAPI(ls.LitAPI):
    def setup(self, device):
        self.model = lambda x: x * 2

    def decoderequest(self, request):

        return request["value"]

    def predict(self, x):
        return self.model(x)

    def encoderesponse(self, output):
        return {"result": output}

    def authorize(self, authheader):

        """Validate API key from Authorization header."""
        if authheader != "Bearer my-secret-key":
            raise ls.HTTPException(statuscode=401, detail="Invalid API key")


if name == "main":
    api = ProtectedAPI()
    server = ls.LitServer(api)
    server.run(port=8000)

if name == "main":
    api = MyModelAPI()
    server = ls.LitServer(
        api,
        accelerator="gpu",
        devices="auto",        # Use all available GPUs
        workersperdevice=2,  # 2 workers per GPU

        timeout=30,            # 30 second timeout per request
    )
    server.run(port=8000)

server = ls.LitServer(
    api,
    accelerator="gpu",
    devices=[0, 1, 2, 3],    # Use GPUs 0-3
    workersperdevice=1,

)

from pydantic import BaseModel
from typing import List

class PredictRequest(BaseModel):
    features: List[float]
    modelversion: str = "v1"

class MyAPI(ls.LitAPI):
    def decoderequest(self, request: PredictRequest):

        return request.features

curl http://localhost:8000/health
Response: {"status": "ok"}

class RobustAPI(ls.LitAPI):
    def decoderequest(self, request):
        if "text" not in request:
            raise ls.HTTPException(
                statuscode=400,

                detail="Field 'text' is required"
            )
        if len(request["text"]) > 10000:
            raise ls.HTTPException(
                statuscode=400,
                detail="Text exceeds maximum length of 10000 characters"
            )
        return request["text"]

4. Logging and Debugging

import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(name)

class DebugAPI(ls.LitAPI):
    def predict(self, x):
        logger.info(f"Processing input shape: {x.shape}")
        result = self.model(x)
        logger.info(f"Prediction complete: {result[:5]}...")
        return result

5. Memory Optimization

class MemoryEfficientAPI(ls.LitAPI):
    def setup(self, device):
        self.model = loadmodel()

        if "cuda" in str(device):
            self.model = self.model.half()
        self.model.eval()

    def predict(self, x):
        with torch.nograd(), torch.cuda.amp.autocast():
            return self.model(x)

6. Docker Deployment

FROM python:3.11-slim

WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY server.py .

EXPOSE 8000
CMD ["python", "server.py"]

# requirements.txt
litserve
torch
transformers

Testing Your API

Unit Testing

import pytest
from server import IrisAPI

def testirisapi():
    api = IrisAPI()
    api.setup("cpu")

    # Test decode
    request = {"features": [5.1, 3.5, 1.4, 0.2]}
    x = api.decoderequest(request)

    assert x.shape == (1, 4)

    # Test predict
    result = api.predict(x)
    assert "classidx" in result
    assert "probabilities" in result

    # Test encode
    response = api.encoderesponse(result)

    assert "predictedclass" in response
    assert response["predictedclass"] in ["setosa", "versicolor", "virginica"]

# locustfile.py
from locust import HttpUser, task, between

class APIUser(HttpUser):
    waittime = between(0.1, 0.5)

    @task
    def predict(self):
        self.client.post("/predict", json={
            "features": [5.1, 3.5, 1.4, 0.2]
        })

pip install locust
locust -f locustfile.py --host=http://localhost:8000

Comparison with Other Frameworks

| Feature | LitServe | FastAPI | BentoML | TorchServe |

|---------|----------|---------|---------|------------|

| Setup Complexity | Low | Low | Medium | High |

| Auto Batching | Yes | No | Yes | Yes |

| Streaming | Yes | Yes | Yes | Yes |

| Multi-GPU | Yes | Manual | Yes | Yes |

| OpenAPI Docs | Yes | Yes | Yes | No |

| GPU Autoscaling | Yes | No | Yes | Limited |

Conclusion

LitServe simplifies the AI model deployment process by providing the right level of abstraction without sacrificing flexibility. By implementing just 4 core methods (setup, decoderequest, predict, encoderesponse), you can create production-ready APIs that support batching, streaming, multi-model serving, and GPU scaling.

Key takeaways:

• Start with a simple LitAPI implementation, then add features as needed
• Use batching to increase GPU throughput
• Leverage streaming for generative models
• Implement authentication for production APIs
• Perform load testing before deploying to production

LitServe is ideal for teams that want to focus on model development without spending time on complex serving infrastructure. With backing from Lightning AI, the framework continues to evolve with regular updates.