moebiusT7/book-ocr-studio

1"""Bounded GPU admission; never unload or stop another application's models."""2import time3from scheduler import inventory4 5 6def wait_for_memory(gpu, minimum, cancelled=lambda: False, observe=lambda event: None,7                    timeout=30, sample=None, clock=time.monotonic, sleep=time.sleep):8    sample = sample or inventory9    deadline = clock() + timeout10    stable = 011    while True:12        if cancelled():13            raise RuntimeError('GPU admission cancelled')14        cards = sample()  # Query failures are errors, never treated as free VRAM.15        card = next((c for c in cards if c['index'] == gpu), None)16        if card is None:17            raise RuntimeError(f'GPU {gpu} is unavailable')18        free = card['memory.free']19        stable = stable + 1 if free >= minimum else 020        observe(dict(gpu=gpu, free_mib=free, required_mib=minimum, stable_samples=stable))21        if stable >= 2:22            return free23        if clock() >= deadline and free < minimum:24            raise RuntimeError(f'Out of VRAM: GPU {gpu} has {free}MiB free; {minimum}MiB required after waiting. Other applications were not stopped.')25        sleep(.5)26