Spaces:
Running on Zero
Running on Zero
| """Watch forked ZeroGPU workers from the parent process (Linux /proc only). | |
| A worker that stalls before its first output leaves nothing in the logs, and a | |
| frozen process cannot report on itself. While a GPU call is pending, the parent | |
| logs each python child's state, per-thread kernel wait channels, CPU time and | |
| I/O. When a worker makes no progress between two looks, the parent sends it | |
| SIGUSR1 once; faulthandler, registered before the fork, then dumps every | |
| Python thread's stack from inside the worker. | |
| """ | |
| from __future__ import annotations | |
| from collections import Counter | |
| import os | |
| from pathlib import Path | |
| import signal | |
| PROC = Path("/proc") | |
| # No progress: under 5% of one CPU and under 1 MiB of reads/writes per interval. | |
| STALL_CPU_FRACTION = 0.05 | |
| STALL_IO_BYTES = 1 << 20 | |
| def _stat(text): | |
| """(comm, fields after comm); comm may itself contain spaces or parentheses.""" | |
| end = text.rindex(")") | |
| return text[text.index("(") + 1:end], text[end + 2:].split() | |
| def child_pids(parent=None, proc=PROC): | |
| parent = os.getpid() if parent is None else parent | |
| pids = [] | |
| for entry in proc.iterdir(): | |
| if not entry.name.isdigit(): | |
| continue | |
| try: | |
| comm, fields = _stat((entry / "stat").read_text()) | |
| except (OSError, ValueError): | |
| continue | |
| if int(fields[1]) == parent and "python" in comm: | |
| pids.append(int(entry.name)) | |
| return sorted(pids) | |
| def snapshot(pid, proc=PROC): | |
| base = proc / str(pid) | |
| _, fields = _stat((base / "stat").read_text()) | |
| try: | |
| io = dict(line.split(": ", 1) for line in (base / "io").read_text().splitlines()) | |
| io_bytes = int(io["rchar"]) + int(io["wchar"]) | |
| except (OSError, KeyError, ValueError): | |
| io_bytes = None | |
| threads = Counter() | |
| for task in (base / "task").iterdir(): | |
| try: | |
| state = _stat((task / "stat").read_text())[1][0] | |
| wchan = (task / "wchan").read_text().strip() or "-" | |
| except (OSError, ValueError, IndexError): | |
| continue | |
| threads[f"{state}:{wchan}"] += 1 | |
| return {"state": fields[0], "cpu_ticks": int(fields[11]) + int(fields[12]), | |
| "io_bytes": io_bytes, "threads": dict(threads)} | |
| def stalled(before, now, interval, ticks_per_s): | |
| if before is None: | |
| return False | |
| cpu_s = (now["cpu_ticks"] - before["cpu_ticks"]) / ticks_per_s | |
| io = None if None in (now["io_bytes"], before["io_bytes"]) else now["io_bytes"] - before["io_bytes"] | |
| return cpu_s < STALL_CPU_FRACTION * interval and (io is None or io < STALL_IO_BYTES) | |
| def watch(stop, interval, proc=PROC, kill=os.kill): | |
| """Until stop is set, log every python child each interval; dump a stalled one once.""" | |
| ticks_per_s = os.sysconf("SC_CLK_TCK") | |
| last, dumped = {}, set() | |
| while not stop.wait(interval): | |
| for pid in child_pids(proc=proc): | |
| try: | |
| now = snapshot(pid, proc) | |
| except (OSError, ValueError, IndexError): | |
| continue | |
| stuck = stalled(last.get(pid), now, interval, ticks_per_s) | |
| last[pid] = now | |
| threads = ",".join(f"{key}x{count}" for key, count in sorted(now["threads"].items())) | |
| io_mib = "?" if now["io_bytes"] is None else now["io_bytes"] >> 20 | |
| print(f"worker_probe pid={pid} state={now['state']} cpu_s={now['cpu_ticks'] / ticks_per_s:.0f} " | |
| f"io_mib={io_mib} stalled={stuck} threads={threads}", flush=True) | |
| if stuck and pid not in dumped: | |
| dumped.add(pid) | |
| try: | |
| kill(pid, signal.SIGUSR1) | |
| except OSError: | |
| pass | |