Metadata-Version: 2.4
Name: gpupilot
Version: 0.1.0
Summary: GPU job scheduling library
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: pynvml
Requires-Dist: torch
Dynamic: license-file

# GPUPilot

Run multiple PyTorch models across multiple GPUs without manually managing device placement, VRAM limits, or scheduling logic.

GPUPilot discovers your GPUs, measures actual VRAM usage when each model loads, and routes inference jobs to wherever there is room — evicting idle models under pressure and retrying once on OOM.

---

## Installation

```bash
pip install gpupilot
```

Requires Python 3.10+, PyTorch, and `nvidia-ml-py` (NVML bindings).

---

## Quick start

```python
import torch
import torch.nn as nn
from gpupilot import Scheduler

scheduler = Scheduler(headroom_gb=2.0)

@scheduler.model(name="my_model", memory_estimate_gb=0.5)
def load_my_model():
    return nn.Linear(1000, 1000)

scheduler.initialize()

result = scheduler.run("my_model", torch.randn(1, 1000))
print(result.shape)  # torch.Size([1, 1000])

scheduler.shutdown()
```

`initialize()` discovers GPUs via NVML and starts worker threads. `run()` blocks until the result is ready. Always call `shutdown()` when done.

---

## Real-world example: faster-whisper on multiple GPUs

[faster-whisper](https://github.com/SYSTRAN/faster-whisper) uses CTranslate2 under the hood, not `nn.Module`, so it manages its own device placement. The pattern below defers GPU assignment to the `to()` call that GPUPilot makes after loading:

```python
import torch
from gpupilot import Scheduler
from faster_whisper import WhisperModel


class WhisperWrapper:
    """Adapts faster-whisper to GPUPilot's loader protocol."""

    def __init__(self, model_size: str, compute_type: str = "float16"):
        self._model_size = model_size
        self._compute_type = compute_type
        self._model: WhisperModel | None = None

    def to(self, device: str) -> "WhisperWrapper":
        # GPUPilot calls .to("cuda:N") after loader_fn() returns.
        # ctranslate2 needs device="cuda" + device_index=N separately.
        backend, _, idx = device.partition(":")
        device_index = int(idx) if idx else 0
        self._model = WhisperModel(
            self._model_size,
            device=backend,
            device_index=device_index,
            compute_type=self._compute_type,
        )
        return self

    def __call__(self, audio_path: str):
        assert self._model is not None
        segments, info = self._model.transcribe(audio_path, beam_size=5)
        return list(segments), info


# ---- scheduler setup ----

scheduler = Scheduler(
    headroom_gb=2.0,    # keep 2 GB free per GPU as a safety buffer
    num_workers=4,      # inference threads (default: 4)
)

@scheduler.model(
    name="whisper-small",
    memory_estimate_gb=1.0,   # rough estimate; actual usage is measured on load
    idle_unload_seconds=300,  # unload from GPU after 5 min idle
)
def load_whisper_small():
    return WhisperWrapper("Systran/faster-whisper-small")


scheduler.initialize()

# Synchronous — blocks until the transcription completes
segments, info = scheduler.run("whisper-small", "interview.mp3")

for seg in segments:
    print(f"[{seg.start:.2f}s → {seg.end:.2f}s] {seg.text}")

scheduler.shutdown()
```

GPUPilot measures actual VRAM delta via NVML after the model loads and warns you if it deviates more than 20% from `memory_estimate_gb`. Update the estimate if you see that warning.

---

## Multiple models, async dispatch

Register as many models as you like. GPUPilot picks the best GPU for each one and evicts idle models when VRAM is tight.

```python
import torch
from gpupilot import Scheduler
from faster_whisper import WhisperModel


scheduler = Scheduler(headroom_gb=2.0, num_workers=8)


@scheduler.model(name="whisper-small", memory_estimate_gb=1.0, idle_unload_seconds=300)
def load_small():
    return WhisperWrapper("Systran/faster-whisper-small")


@scheduler.model(name="whisper-large", memory_estimate_gb=6.0, idle_unload_seconds=600)
def load_large():
    return WhisperWrapper("Systran/faster-whisper-large-v3", compute_type="float16")


scheduler.initialize()

# Fire off several jobs without waiting for each one
audio_files = ["clip1.mp3", "clip2.mp3", "clip3.mp3"]

futures = [
    scheduler.submit_async("whisper-small", path)
    for path in audio_files
]

# Collect results when ready
for path, future in zip(audio_files, futures):
    segments, info = future.result(timeout=120)
    print(f"{path}: {info.duration:.1f}s detected as {info.language}")

scheduler.shutdown()
```

`submit_async()` returns a `concurrent.futures.Future` immediately. Use `future.result(timeout=...)` to block on individual results.

---

## Priority jobs

Higher-priority jobs jump ahead in the queue:

```python
# Priority 10 runs before priority 0 (default)
urgent = scheduler.submit_async("whisper-small", "urgent.mp3", priority=10)
normal = scheduler.submit_async("whisper-small", "batch.mp3", priority=0)
```

---

## Metrics

```python
scheduler.initialize()
# ... run some jobs ...

m = scheduler.get_metrics()

for gpu in m.gpus:
    print(f"GPU {gpu.gpu_id} ({gpu.name}): "
          f"{gpu.used_vram_gb:.1f} / {gpu.total_vram_gb:.1f} GB used, "
          f"loaded: {gpu.loaded_models}")

for inst in m.instances:
    print(f"  {inst.model_name} on GPU {inst.gpu_id}: "
          f"status={inst.status}, active_jobs={inst.active_jobs}, "
          f"vram={inst.allocated_vram_gb:.2f} GB")

print(f"Queue depth: {m.queue_depth}")
```

---

## Configuration reference

### `Scheduler`

| Parameter | Default | Description |
|---|---|---|
| `headroom_gb` | `1.0` | VRAM to keep free per GPU as a safety buffer |
| `num_workers` | `4` | Inference worker threads |
| `poll_interval_s` | `1.0` | NVML monitor refresh interval |
| `reaper_interval_s` | `10.0` | How often the idle-reaper thread wakes to check for stale instances |

### `ModelSpec` / `@scheduler.model()`

| Parameter | Required | Description |
|---|---|---|
| `name` | yes | Unique model identifier |
| `memory_estimate_gb` | yes | Expected VRAM usage. Used for placement; actual usage is measured and may differ. |
| `idle_unload_seconds` | no (default `180`) | Unload from GPU after this many idle seconds. `None` = pinned (never reaped). |
| `max_concurrency` | no | Cap on simultaneous inference jobs for this model (not yet enforced — v1 roadmap). |

---

## How it works

1. **Discovery** — `initialize()` calls NVML to enumerate GPUs and their total VRAM.
2. **Placement** — when a job arrives for an unloaded model, the placement engine picks the GPU with the most safe free VRAM (`total − used − reserved − headroom`). `reserved_vram_gb` closes the race window between placement and NVML seeing the new allocation.
3. **VRAM measurement** — the loader wraps `loader_fn()` in a `measure_vram_delta` context that syncs CUDA and reads NVML before and after, storing the true delta on the instance.
4. **Eviction** — if no GPU has room, GPUPilot evicts idle instances LRU-first until enough VRAM is freed, then retries placement.
5. **OOM recovery** — if a `torch.cuda.OutOfMemoryError` escapes inference, GPUPilot unloads the instance and requeues the job once. A second OOM raises `OOMError` to the caller.
6. **Idle reaper** — a background thread unloads instances that have been idle longer than `idle_unload_seconds`. Models registered with `idle_unload_seconds=None` are pinned and never reaped.
