Spaces:
Running on Zero
Running on Zero
File size: 5,396 Bytes
aab35ad | 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 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | """音声やファイル名を含まない処理計測。測定失敗は本処理へ伝播させない。"""
import functools
import inspect
import json
import math
import re
import subprocess
import threading
import time
from pathlib import Path
_lock = threading.Lock()
def _path(root, jid):
if not re.fullmatch(r"[a-f0-9]{32}", str(jid)):
raise ValueError("invalid_job_id")
return Path(root) / jid / "gpu_usage.json"
def _update(root, jid, update):
try:
with _lock:
target = _path(root, jid)
try:
data = json.loads(target.read_text(encoding="utf-8"))
except (OSError, ValueError):
data = {}
update(data)
temporary = target.with_suffix(".usage.tmp")
temporary.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")
temporary.replace(target)
except Exception:
print("[GPU Usage] measurement_write_failed", flush=True)
def begin(root, jid, duration=None):
def change(data):
data.clear()
data.update(schema_version=1, request_started_at=int(time.time()*1000),
processing_started_at=int(time.time()*1000),
audio_duration_seconds=duration if isinstance(duration, (int, float)) and math.isfinite(duration) and duration > 0 else None,
processing_ms=None, queue_wait_ms=None, gpu_seconds=None,
gpu_seconds_is_estimated=False, gpu_function_ms=None,
device_utilization_pct=None, gpu_sample_count=0, status="processing")
_update(root, jid, change)
def read(root, jid):
try:
with _lock:
return json.loads(_path(root, jid).read_text(encoding="utf-8"))
except Exception:
return None
def finish(root, jid, succeeded):
def change(data):
end = int(time.time()*1000)
data.update(request_finished_at=end, processing_finished_at=end,
status="completed" if succeeded is True else "failed" if succeeded is False else "unknown",
error_code=None if succeeded is True else "processing_failed" if succeeded is False else "status_unavailable")
start = data.get("processing_started_at")
data["processing_ms"] = max(0, end-start) if start else None
_update(root, jid, change)
def sample_devices():
"""機器全体の瞬時使用率。プロセス単位のGPU時間ではない。"""
result = subprocess.run([
"nvidia-smi", "--query-gpu=index,utilization.gpu,memory.used,memory.total",
"--format=csv,noheader,nounits"], capture_output=True, text=True, timeout=1)
if result.returncode:
return []
devices = []
for line in result.stdout.splitlines():
values = [value.strip() for value in line.split(",")]
if len(values) != 4:
continue
try:
index, utilization, used, total = map(float, values)
if not all(math.isfinite(value) for value in [index, utilization, used, total]) or not 0 <= utilization <= 100 or total <= 0:
continue
devices.append({"index": int(index), "utilization_pct": utilization,
"memory_used_mib": used, "memory_total_mib": total})
except ValueError:
continue
return devices
def record_gpu_function(function):
"""spaces.GPUの内側に置く。関数の経過時間と機器サンプルを取得する。"""
signature = inspect.signature(function)
@functools.wraps(function)
def wrapped(*args, **kwargs):
bound = signature.bind(*args, **kwargs)
jid = bound.arguments.get("jid")
root = function.__globals__.get("TEMP_DIR")
if root is None or not re.fullmatch(r"[a-f0-9]{32}", str(jid)):
return function(*args, **kwargs)
started = time.perf_counter()
stop = threading.Event()
samples = []
def sample():
while not stop.is_set():
try:
devices = sample_devices()
# 複数GPUは処理対象機器を特定できないため平均へ混ぜない。
if len(devices) == 1:
samples.append(devices[0]["utilization_pct"])
except Exception:
return
stop.wait(2)
thread = threading.Thread(target=sample, daemon=True)
try:
thread.start()
except Exception:
pass
try:
return function(*args, **kwargs)
finally:
elapsed = round((time.perf_counter()-started)*1000)
stop.set()
def change(data):
data["gpu_function_ms"] = (data.get("gpu_function_ms") or 0) + elapsed
# 壁時計時間をGPU稼働時間に変換しない。
data["gpu_seconds"] = None
captured = list(samples)
count = len(captured)
old_count = data.get("gpu_sample_count", 0)
old_value = data.get("device_utilization_pct") or 0
data["gpu_sample_count"] = old_count + count
data["device_utilization_pct"] = (old_value*old_count+sum(captured))/(old_count+count) if old_count+count else None
_update(root, jid, change)
return wrapped
|