#!/usr/bin/env python3
"""
Hold a fixed amount of VRAM, so contention can be reproduced without ComfyUI.

    ballast.py 4096          # pin ~4 GiB until killed

Prints `READY <mib>` on stdout once the memory is actually resident, then waits
for SIGTERM. The point is a *reproducible* co-tenant: a real image-generation
job's footprint swings by gigabytes between steps, which makes it useless as a
controlled variable. A ballast process holds one number flat, so the only thing
changing between runs is how much of the card the other tenant took.

Allocation is done in 256 MiB chunks and then touched, because PyTorch's caching
allocator will happily hand back a tensor whose pages the driver has not yet
committed — an untouched allocation can read as far less resident than asked for.
"""

import signal
import sys
import time

import torch

CHUNK_MIB = 256


def main() -> None:
    want = int(sys.argv[1])
    if not torch.cuda.is_available():
        sys.exit("no CUDA device")

    held = []
    got = 0
    while got < want:
        n = min(CHUNK_MIB, want - got)
        try:
            t = torch.empty(n * 1024 * 1024, dtype=torch.uint8, device="cuda")
        except torch.cuda.OutOfMemoryError:
            break
        t.fill_(1)          # commit the pages
        held.append(t)
        got += n
    torch.cuda.synchronize()

    running = True

    def stop(*_):
        nonlocal running
        running = False

    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)

    print(f"READY {got}", flush=True)
    while running:
        time.sleep(0.2)


if __name__ == "__main__":
    main()