Metadata-Version: 2.4
Name: vrambatch
Version: 0.1.0
Summary: Run batched GPU work inside a VRAM budget, self-healing on CUDA out-of-memory
Author-email: Saurabh Sharma <hi.techbysaurabh@gmail.com>
License-Expression: MIT
Project-URL: Homepage, https://github.com/techbysaurabh/vrambatch
Project-URL: Repository, https://github.com/techbysaurabh/vrambatch
Project-URL: Issues, https://github.com/techbysaurabh/vrambatch/issues
Keywords: gpu,cuda,vram,oom,batching,inference,pytorch,memory
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Environment :: GPU :: NVIDIA CUDA
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Provides-Extra: torch
Requires-Dist: torch>=1.10; extra == "torch"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Dynamic: license-file

# vrambatch

Run batched GPU work inside a VRAM budget — and when a batch is still too big,
**halve it and retry instead of crashing**.

If you serve inference on a shared or memory-capped GPU, you have written this
code already: pick a batch size that fits, and pray no single request is
heavier than your estimate. `vrambatch` is that logic, extracted and tested —
size a pass from a memory budget, and self-heal on CUDA out-of-memory so a bad
estimate costs an extra pass, not a 500.

```bash
pip install vrambatch          # sizing only (pure Python)
pip install vrambatch[torch]   # + auto-sizing and OOM detection on GPU
```

## The problem it removes

A memory-capped process picks a batch size up front. Two things then go wrong:

- **Too small** and the GPU sits idle — most of your VRAM cap is unreachable
  and requests queue instead of fusing.
- **Too big** and one unusually heavy item OOMs the whole batch, failing every
  request in it.

The honest fix is to size for the *typical* case and recover from the rare
overflow. That needs a retry that can make progress — which is exactly what a
hand-rolled version usually lacks.

## Usage

```python
from vrambatch import run_in_passes

# process() takes a list of items, returns one result per item, in order.
def process(batch):
    feats = preprocess(batch)
    return model.generate(feats)        # your real GPU call

results = run_in_passes(
    items,
    process,
    gb_per_unit=0.25,        # measured VRAM per unit (see below)
    budget_gb=21,            # your cap; omit to read free VRAM automatically
    reserve_gb=3.0,          # model weights + context, subtracted first
    cost=lambda item: n_chunks(item),   # units per item (default 1)
)
```

- Items are grouped into passes whose total **cost** stays within the budget.
- Each pass is one `process()` call.
- On CUDA OOM the offending pass is **split in half and retried**, recursing to
  a single item. A lone item that still cannot fit raises `OOMError` — a real
  failure, reported as one. Non-OOM errors are never retried.

Pass `on_oom=lambda n: log(...)` to observe splits.

## Just the sizing, if that's all you need

```python
from vrambatch import plan_pass_size, available_vram_gb

n = plan_pass_size(budget_gb=21, gb_per_unit=0.25, reserve_gb=3.0, safety=0.8)
# -> units per pass

free = available_vram_gb()   # device-free + the allocator's reusable pool
```

## Measuring `gb_per_unit`

Do not guess it high "to be safe" — an over-large figure is what leaves your
cap unreachable. Measure it: run increasing batch sizes and fit a line to peak
memory.

```python
import torch
for n in (1, 2, 4, 8, 16):
    torch.cuda.reset_peak_memory_stats()
    process(items[:n])
    print(n, torch.cuda.max_memory_allocated() / 2**30)
# peak ≈ reserve_gb + gb_per_unit * n  → the slope is gb_per_unit
```

Because overflow is recoverable, size for the typical unit and let the
halve-and-retry cover the occasional heavy one. The `safety` factor (default
0.8) absorbs normal variance; the retry covers the rest.

## Notes

- **torch is optional.** `plan_pass_size` is pure Python. `available_vram_gb`
  and precise OOM typing use torch when present; OOM detection falls back to
  matching the error message so `RuntimeError("CUDA out of memory")` is caught
  either way.
- **Async servers:** `run_in_passes` is synchronous by design — call it from a
  thread (`await asyncio.to_thread(run_in_passes, ...)` or a single-worker
  executor) so it never blocks your event loop.
- **Fragmentation:** under a hard `set_per_process_memory_fraction` cap, also
  set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` before torch inits
  CUDA, or the allocator's cached pool can fragment and fail an allocation that
  is nominally under budget.

## License

MIT
