Metadata-Version: 2.4
Name: smartbatch
Version: 0.1.0
Summary: High-throughput ML inference middleware with dynamic batching
Author-email: Veera Karthick <veerakarthick09@gmail.com>
License: MIT
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: fastapi>=0.68.0
Requires-Dist: uvicorn>=0.15.0
Requires-Dist: pydantic>=2.0.0
Requires-Dist: prometheus_client>=0.17.0
Requires-Dist: msgpack>=1.0.0
Requires-Dist: torch>=2.0.0
Requires-Dist: numpy>=1.21.0
Provides-Extra: test
Requires-Dist: pytest; extra == "test"
Requires-Dist: locust; extra == "test"
Requires-Dist: requests; extra == "test"
Dynamic: license-file

# SmartBatch: High-Throughput Async Inference Middleware

**SmartBatch** is a production-grade inference serving system designed to maximize GPU utilization and throughput for PyTorch models. It implements **Dynamic Batching** to group incoming requests on-the-fly, significantly reducing overhead compared to naive request-per-inference processing.

## 🚀 Key Features

*   **Dynamic Batching**: Automatically groups requests into batches (up to `MAX_BATCH_SIZE`) or processes them after a timeout (`MAX_WAIT_TIME`), striking the perfect balance between throughput and latency.
*   **Asynchronous API**: Built on `FastAPI` and `asyncio` to handle thousands of concurrent connections efficiently.
*   **Production Robustness**: Includes graceful shutdown, proper error handling, and thread-safe metrics.
*   **Real-World Load Testing**: Benchmarking suite included to simulate high-concurrency traffic with realistic payloads.
*   **Observability**: `/metrics` endpoint for real-time monitoring of latency, batch sizes, and throughput.

## 📊 Benchmark Results (Stress Test)

We conducted a rigorous 1-hour stress test comparing **SmartBatch** against a **Baseline** (no batching) implementation.

**Hardware**: Single Node (Simulated Production Environment)
**Load**: 200-1000 Concurrent Users
**Payload**: Real-world images (ResNet18 inputs)

| Metric | Baseline (Sequential) | SmartBatch (Batched) | Improvement |
| :--- | :--- | :--- | :--- |
| **Throughput (RPS)** | ~0.67 req/s | **~2.22 req/s** | **3.3x Higher** |
| **Median Latency (p50)** | > 1000s (Collapsed) | **~13s** (Stable) | **~99% Reduction** |
| **Tail Latency (p95)** | Unstable / Timeouts | Controlled by Batching | **Stabilized** |

---

## 🛠️ Installation

1.  **Clone the repository**:
    ```bash
    git clone https://github.com/yourusername/SmartBatch.git
    cd SmartBatch
    ```

2.  **Create a virtual environment**:
    ```bash
    python3.12 -m venv venv
    source venv/bin/activate
    ```

3.  **Install dependencies**:
    ```bash
    pip install .
    ```

## 🏃 Usage

### 1. The Decorator Pattern (Recommended)
Add `@batch` to any async function to automatically group requests.

```python
from smartbatch import batch
from typing import List

# 1. Define your batched function (List -> List)
@batch(max_batch_size=32, max_wait_time=0.01)
async def run_model(batch_inputs: List[float]) -> List[float]:
    # This runs ONLY when a batch is full or timeout matches
    return model.predict(batch_inputs)

# 2. Call it normally (Single Item -> Single Item)
# The decorator handles queueing and waiting!
result = await run_model(single_input) 
```

### 2. Binary Transport (MsgPack)
For high-performance clients, send binary packed data instead of JSON to reduce payload size.

**Header**: `Content-Type: application/msgpack`
**Body**: MsgPack encoded dict (e.g., `{"data": [...]}`) or raw list `[...]`.

```python
import msgpack, requests
payload = msgpack.packb([0.1, 0.2, 0.3])
requests.post("http://localhost:8000/predict", data=payload, headers={"Content-Type": "application/msgpack"})
```

### 3. Multi-Model Registry
To serve multiple models on dynamic routes (`/models/{name}/predict`), use `@register`:

```python
from smartbatch import batch, register

@register(name="yolo")
@batch(max_batch_size=8)
async def run_yolo(batch: List):
    return yolo_model(batch)

# Now available at: POST /models/yolo/predict
```

### 4. Input Schema Validation (Recommended)
Protect your workers by enforcing Pydantic schemas. Invalid requests (e.g. string instead of int) will raise an error *before* queueing.

```python
from pydantic import BaseModel

class ImageInput(BaseModel):
    data: List[float]
    threshold: float = 0.5

@batch(max_batch_size=32, input_schema=ImageInput)
async def safe_inference(batch: List[ImageInput]):
    # 'batch' contains valid Pydantic objects now!
    inputs = [item.data for item in batch]
    return model.predict(inputs)
```

### 5. Multi-GPU / Multi-Worker Support
Scale verticaly by running multiple worker loops. Use `worker_id` to select devices.

```python
# Models loaded on different GPUs
models = {
    0: load_model("cuda:0"),
    1: load_model("cuda:1")
}

@batch(max_batch_size=32, workers=2)
async def infer(batch, worker_id=0):
    # SmartBatch injects 'worker_id' (0 or 1) automatically
    model = models[worker_id]
    return model(batch)
```

## 🧪 Benchmarking

Reproduce the performance results yourself using the included benchmark suite.

**1. Run the Comparison Benchark**
This script runs both Baseline and SmartBatch scenarios sequentially and generates plots.
```bash
# usage: --users [NUM] --duration [TIME]
venv/bin/python scripts/production_benchmark.py --users 200 --duration 1h
```

**2. View Results**
The script will generate these files in your root directory:
*   `production_throughput.png`
*   `production_latency_p50.png`
*   `production_latency_p95.png`

## 📂 Project Structure

```
SmartBatch/
├── smartbatch/          # Core Package
│   ├── main.py          # Entry point & App lifecycle
│   ├── api.py           # FastAPI endpoints
│   ├── batching.py      # Batch logic & Queue management
│   ├── model.py         # PyTorch Model Wrapper
│   └── metrics.py       # Thread-safe metrics collection
├── scripts/             # Utilities
│   ├── production_benchmark.py # A/B Stress Test Script
│   └── study_params.py  # Hyperparameter optimization
├── tests/               # Testing
│   ├── locustfile.py    # Load generator
│   └── data/            # Test images (for realistic load)
└── README.md            # Documentation
```
