7 min read

LitServe: Building ML APIs with Batching & Caching

AI

ThinkTools Team

AI Research Lead

LitServe: Building ML APIs with Batching & Caching

Introduction

Deploying machine learning models as production‑ready services is a recurring challenge for data scientists and software engineers alike. The traditional approach of wrapping a model in a custom Flask or FastAPI application often leads to duplicated code, brittle deployment pipelines, and hidden performance bottlenecks. Over the past few years, a handful of lightweight serving frameworks have emerged to address these pain points, and LitServe has quickly become a favorite for teams that need to expose multiple models or endpoints without the overhead of a full‑blown inference platform.

LitServe distinguishes itself by offering a minimal API surface that still supports advanced features such as request batching, streaming responses, and in‑memory caching—all while keeping the runtime footprint small enough to run comfortably on a single laptop. This tutorial walks through a concrete implementation that showcases how to build a multi‑endpoint API capable of text generation, summarization, and classification, and how to tune each endpoint for optimal latency and throughput. By the end of the post you will understand how to structure your code, configure LitServe’s routing logic, and validate the system with realistic load tests.

The examples presented here are intentionally self‑contained; they rely solely on open‑source models and local resources, so you can experiment without paying for external API calls. The goal is to give you a practical, end‑to‑end recipe that you can adapt to your own models, whether they are transformer‑based language models, vision encoders, or custom tabular predictors.

Main Content

LitServe Overview

LitServe is built on top of PyTorch Lightning, which means that any model that can be wrapped in a LightningModule automatically becomes a candidate for serving. The framework exposes a simple decorator, @litserve.api, that turns a method into an HTTP endpoint. Behind the scenes, LitServe handles request parsing, response serialization, and concurrency management using Python’s asyncio event loop. One of the most compelling aspects of LitServe is its plug‑in architecture: you can drop in a caching layer, a batching scheduler, or a streaming response generator with minimal code changes.

Designing Multi‑Endpoint Architecture

When designing a multi‑endpoint API, the first decision is how to map incoming routes to model logic. In our example, we expose three distinct endpoints: /generate, /summarize, and /classify. Each endpoint is implemented as a separate method on a single LightningModule subclass, and the decorator automatically registers the route. Because LitServe supports path parameters and query strings, you can expose additional configuration options—such as temperature for generation or max length for summarization—directly in the URL.

A common pitfall is to bundle all logic into one large method, which makes maintenance difficult and obscures the performance characteristics of each task. By keeping the endpoints thin and delegating heavy lifting to helper functions, you preserve clarity and enable targeted optimizations. For instance, the generation endpoint can leverage a GPU‑accelerated transformer, while the classification endpoint can run on CPU without sacrificing latency.

Batching Strategies

Batching is a proven technique to amortize the cost of model inference, especially for transformer‑based architectures where the majority of the time is spent on matrix multiplications rather than data transfer. LitServe provides a Batcher class that collects incoming requests over a configurable time window or until a maximum batch size is reached. The batched inputs are then concatenated and fed into the model in a single forward pass.

In practice, you need to balance two competing objectives: latency and throughput. A batch size of 32 might deliver the highest throughput, but it also introduces a 10‑millisecond delay for the first request in the batch. LitServe’s Batcher exposes a max_latency parameter that allows you to cap the waiting time, ensuring that real‑time applications still meet their SLA. The example code demonstrates how to set max_batch_size=16 and max_latency=0.05 seconds for the generation endpoint, which yields a good trade‑off for interactive chat scenarios.

Streaming Responses

For long‑running inference tasks—such as generating a paragraph of text or translating a document—returning the entire result in a single HTTP response can be wasteful. LitServe supports streaming responses by yielding partial outputs from the model. The framework automatically converts these yields into Server‑Sent Events (SSE) that the client can consume incrementally.

Implementing streaming is straightforward: the endpoint method simply returns a generator that yields strings. Under the hood, LitServe wraps the generator in an async context and streams each chunk as it becomes available. In the tutorial, the /generate endpoint streams tokens as they are produced by a GPT‑style decoder, allowing the client to display the text in real time. This pattern is especially useful for building chatbots or real‑time translation services.

Caching Mechanisms

Repeated requests for the same input are common in production workloads. Without caching, each identical request would trigger a full forward pass, wasting compute resources. LitServe integrates seamlessly with functools.lru_cache and also offers a pluggable cache interface that can be backed by Redis, Memcached, or even a simple in‑memory dictionary.

The example demonstrates a simple LRU cache applied to the summarization endpoint. By decorating the summarization helper with @lru_cache(maxsize=128), the system stores the last 128 unique inputs and returns the cached output instantly for subsequent identical requests. For larger workloads, you can swap in a Redis cache by implementing the Cache interface and passing it to the LitServe constructor.

Local Inference and Performance

One of the key selling points of LitServe is that it can run entirely locally, eliminating the need for external API calls. This is particularly valuable for privacy‑sensitive applications or for teams that want to avoid vendor lock‑in. The tutorial shows how to load a Hugging Face transformer model from the local filesystem, configure the device mapping, and expose it through LitServe.

Performance tuning involves several layers: model quantization, mixed‑precision inference, and efficient batching. The post walks through quantizing a BERT model to 8‑bit weights using torch.quantization, which reduces memory usage by 75% and speeds up inference on CPUs. It also demonstrates how to enable PyTorch’s native torch.compile to accelerate the forward pass on recent hardware.

Testing and Validation

A robust serving pipeline must be validated against realistic traffic patterns. LitServe’s integration with pytest and httpx allows you to write end‑to‑end tests that simulate concurrent requests, measure latency distributions, and assert correctness of the outputs. The tutorial includes a test suite that sends 100 concurrent requests to each endpoint, verifies that the responses match expected patterns, and records the 95th percentile latency.

By automating these tests, you can catch regressions early and ensure that any changes to the model or the batching logic do not degrade performance. Continuous integration pipelines can run these tests on every commit, providing confidence that the API remains reliable as the codebase evolves.

Conclusion

Building a production‑grade machine learning API is no longer a matter of writing a few lines of Flask code. Modern frameworks like LitServe empower developers to expose sophisticated models with minimal friction while still offering fine‑grained control over performance characteristics. By combining request batching, streaming responses, and caching, you can create an API that delivers low latency for interactive use cases and high throughput for batch processing.

The tutorial presented a concrete, end‑to‑end example that covers the entire lifecycle—from model packaging to endpoint definition, from performance tuning to automated testing. The same principles apply regardless of the underlying model architecture, making LitServe a versatile tool in any data scientist’s toolbox.

Whether you are building a chatbot, a summarization service, or a custom classification engine, the patterns demonstrated here will help you deliver reliable, scalable, and maintainable machine learning services.

Call to Action

If you’re ready to move beyond prototype notebooks and into the realm of production‑ready inference, give LitServe a try. Start by cloning the example repository, loading your favorite Hugging Face model, and exposing it with a single decorator. Experiment with the batching and caching knobs, measure the impact on latency, and iterate until you hit your performance targets. Share your results on GitHub or Stack Overflow—your insights could help others accelerate their own ML deployments. Happy coding!

We value your privacy

We use cookies, including Google Analytics, to improve your experience on our site. By accepting, you agree to our use of these cookies. Learn more