Spaces:
Running on Zero
Running on Zero
| import hashlib | |
| import json | |
| import logging | |
| import os | |
| import time | |
| from contextlib import contextmanager | |
| from typing import Any, Optional | |
| logger = logging.getLogger(__name__) | |
| WORKER = {"pid": None, "calls": 0, "last_end": None} | |
| def session_tag(session_hash: Any) -> str: | |
| if not session_hash: | |
| return "" | |
| return hashlib.sha256(str(session_hash).encode("utf-8")).hexdigest()[:12] | |
| def current_session_tag() -> str: | |
| try: | |
| from gradio.context import LocalContext | |
| return session_tag(getattr(LocalContext.request.get(None), "session_hash", "")) | |
| except Exception: | |
| return "" | |
| def process_age_seconds() -> Optional[float]: | |
| try: | |
| with open("/proc/self/stat") as stat: | |
| start_ticks = int(stat.read().rsplit(")", 1)[1].split()[19]) | |
| with open("/proc/uptime") as uptime: | |
| uptime_seconds = float(uptime.read().split()[0]) | |
| return round(uptime_seconds - start_ticks / os.sysconf("SC_CLK_TCK"), 2) | |
| except Exception: | |
| return None | |
| def emit(kind: str, fields: dict) -> None: | |
| try: | |
| logger.info("%s %s", kind, json.dumps(fields, separators=(",", ":"))) | |
| except Exception: | |
| pass | |
| def labs_event(event: str, session: str, **extra) -> None: | |
| emit("LABS_EVENT", {"event": event, "session": session, "t": round(time.time(), 2), **extra}) | |
| def gpu_call(fn: str, session: str, texture_size: Any, **extra): | |
| started = time.time() | |
| record = {} | |
| try: | |
| pid = os.getpid() | |
| if WORKER["pid"] != pid: | |
| WORKER.update(pid=pid, calls=0, last_end=None) | |
| WORKER["calls"] += 1 | |
| last_end = WORKER["last_end"] | |
| record = { | |
| "fn": fn, "pid": pid, "calls_in_worker": WORKER["calls"], | |
| "since_prev_call_s": None if last_end is None else round(started - last_end, 2), | |
| "since_worker_start_s": process_age_seconds(), "texture_size": texture_size, "session": session, **extra, | |
| } | |
| except Exception: | |
| pass | |
| ok = False | |
| try: | |
| yield | |
| ok = True | |
| finally: | |
| try: | |
| WORKER["last_end"] = time.time() | |
| emit("GPU_CALL", {**record, "in_gpu_s": round(WORKER["last_end"] - started, 2), "ok": ok, | |
| "t": round(started, 2)}) | |
| except Exception: | |
| pass | |