Lightning-AI / Lightning-AI/LitServe
Decouple API interface from batching configuration
- Dominant language
- Python
- Stars
- 3.9k
- Forks
- 304
- Avg merge
- 3d 13h
- Merged PRs (30d)
- 6
Description
## 🚀 Feature
Implement consistent `predict` method interface independent of batching configuration in LitServe.
### Motivation
Example code I'm using ⤵️
```python
import random
import litserve as ls
import os
import time
class LitServeBatchingDemoAPI(ls.LitAPI):
def setup(self, device):
print(f"Loading models in process {os.getpid()}")
def decode_request(self, request):
return request["inputs"]
def predict(self, batch):
print("Received batch of size", len(batch), batch)
results = [random.random() for _ in batch]
time.sleep(1.5)
return results
def encode_response(self, output):
return {"output": output}
if __name__ == "__main__":
print(f"Starting server in process {os.getpid()}")
server = ls.LitServer(LitServeBatchingDemoAPI(),
workers_per_device=1,
)
server.run(port=8000)
```
Currently, enabling batching in LitServe (by setting `max_batch_size` and `batch_timeout`) changes the expected implementation of the `predict` method in `LitAPI` subclasses. This creates several issues:
1. The same API implementation behaves differently based on server configuration parameters
2. Developers need to maintain different implementations or add conditional logic based on whether batching is enabled
3. It violates the principle of separation of concerns - server configuration parameters should not affect the API contract
For example, with batching disabled, sending:
```python
requests.post("http://127.0.0.1:8000/predict", json={"inputs": "my-input"})
```
```python
def predict(self, batch):
# batch is a single input, len(batch) will return length of the string
print("Received batch of size", len(batch), batch)
```
With batching enabled (`max_batch_size>=2`):
```python
def predict(self, batch):
# batch is a list of inputs
print("Received batch of size", len(batch), batch)
```
This inconsistency makes it harder to maintain and test APIs, especially when batching configuration might change between development and production environments.
### Pitch
LitServe should provide a consistent interface for the `predict` method regardless of batching configuration.
### Additional context
Similar serving frameworks like TorchServe and RayServe have similar approaches to batching, but this doesn't mean LitServe can't improve upon their design. A consistent API contract would make LitServe more intuitive and easier to use correctly.
More context in [recent Discord discussion.](https://discord.com/channels/1077906959069626439/1286717103247724675/1328493010849108010)
The proposed change:
- Maintains backward compatibility when batching is enabled
- Simplifies API implementation by providing a consistent interface
- Follows the principle of least surprise
- Makes testing easier as there's only one behavior to test
- Reduces potential bugs from incorrect handling of single vs batched inputs
Contributor guide
Research direction
Start by tracing LitAPI.predict through the LitServer paths controlled by max_batch_size and batch_timeout. Verify how single requests and batched requests are passed to predict, then define tests showing one consistent interface while preserving existing batched behavior; no file or test path is identified in the issue.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- api, backend-api-design
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100