| """The flight recorder: one local JSON line per render, nothing phones home. |
| |
| Why it exists: layer (b) of the price meter must never predict minutes from |
| GPU specs. It predicts from the user's OWN history — s/step against work |
| units, keyed per device AND per model/quant so new hardware or a new quant |
| starts a fresh curve. |
| |
| Peaks come from IN-PROCESS torch counters, deliberately not nvidia-smi: |
| per-process readout lies on some cards (the Max-Q lesson, GPU1 here). |
| |
| Least-invasive hook: a pair of nodes the user drops at the front and back of |
| a graph. H3RecordStart resets the counters and passes its input through; |
| H3RecordStop reads them, times the gap, and appends the line. No monkey |
| patching of ComfyUI's executor, nothing that survives a failed run and |
| poisons the next one. |
| |
| The 'streamed' flag is read from ComfyUI's own partial-load log lines when a |
| log path is reachable, else null — never guessed. The window read is bounded |
| by the BYTE OFFSET captured at start(), so a partial-load line emitted by the |
| PREVIOUS run can never be attributed to this one. |
| |
| TODO(T2) — THE CRASH JOURNAL. This pair records successes. A run that OOMs, |
| is interrupted, or dies mid-VAE never reaches stop(), so it silently leaves |
| no record — and those are exactly the runs that locate the VRAM cliff, so |
| their absence CENSORS the calibration set toward optimism (only survivors |
| are sampled). The fix is not a bigger record, it is a different hook: |
| executor-level STAGE EVENTS (queued / model loaded / sampling / decoding / |
| finished|failed) journalled as they happen, so an aborted run leaves a |
| partial line with its stage and its last known peak. Wants a ComfyUI |
| execution callback rather than a node pair; deliberately not attempted |
| tonight, because a hook that survives a crash is exactly the kind of thing |
| that must not break the pack when it does not. |
| """ |
| import json |
| import os |
| import time |
|
|
| DEFAULT_PATH = os.path.join(os.path.expanduser("~"), ".cache", "mainodes", |
| "flight_recorder.jsonl") |
|
|
| _STATE = {} |
|
|
|
|
| def _torch(): |
| try: |
| import torch |
| return torch if torch.cuda.is_available() else None |
| except Exception: |
| return None |
|
|
|
|
| def _log_size(log_path): |
| try: |
| return os.path.getsize(log_path) if log_path else None |
| except OSError: |
| return None |
|
|
|
|
| def start(key="default", device_index=None, log_path=None): |
| t = _torch() |
| dev = None |
| if t is not None: |
| i = t.cuda.current_device() if device_index is None else int(device_index) |
| t.cuda.reset_peak_memory_stats(i) |
| dev = t.cuda.get_device_name(i) |
| _STATE[key] = {"t0": time.time(), "device": dev, "index": device_index, |
| |
| "log_offset": _log_size(log_path), "log_path": log_path} |
| return _STATE[key] |
|
|
|
|
| def stop(key="default", work_units=0, steps=1, pixels=0, frames=0, |
| model="unknown", path=None, log_path=None, extra=None): |
| """Append one record. Returns the record (also for the node's report).""" |
| st = _STATE.pop(key, None) or {"t0": time.time(), "device": None, |
| "index": None, "log_offset": None, |
| "log_path": None} |
| log_path = log_path or st.get("log_path") |
| wall = time.time() - st["t0"] |
| t = _torch() |
| peak = None |
| if t is not None: |
| i = (t.cuda.current_device() if st.get("index") is None |
| else int(st["index"])) |
| peak = int(t.cuda.max_memory_allocated(i)) |
| rec = { |
| "at": time.strftime("%Y-%m-%dT%H:%M:%S"), |
| "device": st.get("device"), |
| "model": str(model), |
| "work_units": int(work_units), |
| "steps": int(steps), |
| "pixels": int(pixels), |
| "frames": int(frames), |
| "wall_s": round(wall, 3), |
| "s_per_step": round(wall / max(1, int(steps)), 4), |
| "peak_bytes": peak, |
| "streamed": read_streamed_flag(log_path, since=st.get("log_offset")), |
| "log_offset": st.get("log_offset"), |
| } |
| if extra: |
| rec.update(dict(extra)) |
| p = os.path.abspath(path or DEFAULT_PATH) |
| os.makedirs(os.path.dirname(p), exist_ok=True) |
| with open(p, "a") as fh: |
| fh.write(json.dumps(rec) + "\n") |
| return rec |
|
|
|
|
| def read_streamed_flag(log_path, tail_bytes=200000, since=None): |
| """ComfyUI says so itself when it cannot fit the weights ('loaded |
| partially', 'lowvram'). None when we cannot see a log — an unknown is |
| reported as unknown, never as False. |
| |
| `since` is the byte offset captured at start(): reading from there is |
| what stops the previous run's partial-load line from being read as this |
| run's (the misattribution the crash-journal TODO is also about).""" |
| if not log_path or not os.path.isfile(log_path): |
| return None |
| try: |
| with open(log_path, "rb") as fh: |
| fh.seek(0, os.SEEK_END) |
| end = fh.tell() |
| lo = int(since) if since is not None else max(0, end - tail_bytes) |
| fh.seek(min(lo, end)) |
| tail = fh.read().decode(errors="replace").lower() |
| except OSError: |
| return None |
| return ("loaded partially" in tail or "lowvram" in tail |
| or "weight streaming" in tail) |
|
|
|
|
| def read(path=None, device=None, model=None): |
| out = [] |
| try: |
| with open(os.path.abspath(path or DEFAULT_PATH)) as fh: |
| for line in fh: |
| line = line.strip() |
| if not line: |
| continue |
| try: |
| r = json.loads(line) |
| except ValueError: |
| continue |
| if device and r.get("device") != device: |
| continue |
| if model and r.get("model") != model: |
| continue |
| out.append(r) |
| except OSError: |
| pass |
| return out |
|
|
