mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-18 14:43:29 +08:00
yyj
98 lines
3.2 KiB
Python
98 lines
3.2 KiB
Python
"""Wait without allocating GPU memory; hold one host-wide advisory lock."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import fcntl
|
|
import json
|
|
import subprocess
|
|
import time
|
|
from pathlib import Path
|
|
|
|
|
|
class ResourceUnavailable(RuntimeError):
|
|
pass
|
|
|
|
|
|
def gpu_busy():
|
|
"""Fail closed on driver/query failures, and include non-compute VRAM users."""
|
|
apps = subprocess.check_output(
|
|
["nvidia-smi", "--query-compute-apps=pid", "--format=csv,noheader,nounits"],
|
|
text=True,
|
|
timeout=20,
|
|
).strip()
|
|
stats = subprocess.check_output(
|
|
[
|
|
"nvidia-smi",
|
|
"--query-gpu=memory.used,utilization.gpu",
|
|
"--format=csv,noheader,nounits",
|
|
],
|
|
text=True,
|
|
timeout=20,
|
|
).strip()
|
|
if not stats:
|
|
raise ResourceUnavailable("No GPUs found")
|
|
try:
|
|
rows = [
|
|
[int(value.strip()) for value in row.split(",")]
|
|
for row in stats.splitlines()
|
|
]
|
|
busy = bool(apps) or any(
|
|
memory > 256 or utilization > 5 for memory, utilization in rows
|
|
)
|
|
except (ValueError, TypeError) as exc:
|
|
raise ResourceUnavailable("Cannot determine GPU usage") from exc
|
|
return busy, {"compute_pids": apps, "gpu_memory_mib_utilization": rows}
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def reservation(
|
|
lock_path, log_path, timeout=21600, interval=15, idle_samples=3, probe=gpu_busy
|
|
):
|
|
"""No Actions concurrency group: it can replace older pending jobs.
|
|
|
|
A persistent lock inode must never be deleted after use. Manual workloads
|
|
should acquire the same lock to close the idle-check/start race.
|
|
"""
|
|
lock_path = Path(lock_path)
|
|
if not lock_path.parent.is_dir():
|
|
raise ResourceUnavailable(
|
|
"Administrator must provision the shared lock directory"
|
|
)
|
|
deadline = time.monotonic() + timeout
|
|
stable = 0
|
|
with lock_path.open("a") as lock, Path(log_path).open("a") as log:
|
|
owned = False
|
|
try:
|
|
while time.monotonic() < deadline:
|
|
if not owned:
|
|
try:
|
|
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
owned = True
|
|
except BlockingIOError:
|
|
pass
|
|
state = {
|
|
"status": "waiting_for_qj5090",
|
|
"reason": "another CI reservation",
|
|
}
|
|
if owned:
|
|
busy, evidence = probe()
|
|
stable = 0 if busy else stable + 1
|
|
state.update(evidence)
|
|
state["reason"] = (
|
|
"GPU workload present" if busy else "checking stable idle"
|
|
)
|
|
if stable >= idle_samples:
|
|
log.write(json.dumps({"status": "reserved"}) + "\n")
|
|
log.flush()
|
|
yield
|
|
return
|
|
print(json.dumps(state), flush=True)
|
|
log.write(json.dumps(state) + "\n")
|
|
log.flush()
|
|
time.sleep(interval)
|
|
raise ResourceUnavailable("Queue wait expired; no model test was run")
|
|
finally:
|
|
if owned:
|
|
fcntl.flock(lock, fcntl.LOCK_UN)
|