File size: 1,024 Bytes
9118991 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | """Admission policy shared by every LoopQ entry point (see loopQ/AGENTS.md)."""
from pathlib import Path
import os
import sys
ALLOWED_GPUS = ("0", "1")
def acquire_gpu(gpu: str):
"""Return a lock that the caller must retain until its GPU work exits."""
if gpu not in ALLOWED_GPUS:
raise ValueError(f"LoopQ permits physical GPUs {ALLOWED_GPUS}; GPU 2 is reserved")
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "scripts"))
from gpu_availability import try_acquire_gpu_lock, eligible_gpus
lock = try_acquire_gpu_lock(gpu)
if lock is None:
raise RuntimeError(f"GPU {gpu} admission lock is occupied; no job started")
try:
allowed, state = eligible_gpus(
[gpu], min_free_memory_mib=30000, shareable=frozenset())
if allowed != [gpu]:
raise RuntimeError(f"GPU {gpu} is not idle: {state.get(gpu)}")
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
return lock
except BaseException:
lock.close()
raise
|