gpupilot 0.1.0

GPU job scheduling library

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

pip install gpupilot

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

---

Quick start

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 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:

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.

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:

# 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

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.