Spaces:
Running on Zero
Running on Zero
| """音声やファイル名を含まない処理計測。測定失敗は本処理へ伝播させない。""" | |
| 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) | |
| 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 | |