Download loopq_quantization/scripts/loopq/gpu_admission.py from JunYoungLee/ut-depth-probe-artifacts: direct link, hf CLI and curl.
- Browser
- Download file 1.02 kB
-
https://huggingface.co/JunYoungLee/ut-depth-probe-artifacts/resolve/main/loopq_quantization/scripts/loopq/gpu_admission.py
- Command line
-
hf download hf://JunYoungLee/ut-depth-probe-artifacts/loopq_quantization/scripts/loopq/gpu_admission.py
-
curl -L -o gpu_admission.py https://huggingface.co/JunYoungLee/ut-depth-probe-artifacts/resolve/main/loopq_quantization/scripts/loopq/gpu_admission.py
1.02 kB
| """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 | |