47 lines
1.2 KiB
Python
47 lines
1.2 KiB
Python
"""Shared GPU lease and process-group lifecycle helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import threading
|
|
from contextlib import suppress
|
|
|
|
|
|
class ResourceBusyError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class GpuLease:
|
|
def __init__(self):
|
|
self.lock = threading.RLock()
|
|
self.owner: str | None = None
|
|
|
|
def acquire(self, owner: str) -> None:
|
|
with self.lock:
|
|
if self.owner is not None and self.owner != owner:
|
|
raise ResourceBusyError(f"计算资源正由 {self.owner} 使用")
|
|
self.owner = owner
|
|
|
|
def release(self, owner: str) -> None:
|
|
with self.lock:
|
|
if self.owner == owner:
|
|
self.owner = None
|
|
|
|
def public(self) -> str | None:
|
|
with self.lock:
|
|
return self.owner
|
|
|
|
|
|
def terminate_process(process: subprocess.Popen[str], grace_seconds: float = 5.0) -> None:
|
|
if process.poll() is not None:
|
|
return
|
|
with suppress(ProcessLookupError):
|
|
os.killpg(process.pid, signal.SIGTERM)
|
|
try:
|
|
process.wait(timeout=grace_seconds)
|
|
except subprocess.TimeoutExpired:
|
|
with suppress(ProcessLookupError):
|
|
os.killpg(process.pid, signal.SIGKILL)
|