Files
pollockjj-ComfyUI-MultiGPU/ci/extract_allocation.py
T
John Pollock 0b438f7a8b Refactor code for improved readability and performance
- Cleaned up unnecessary whitespace and comments in model_management_mgpu.py, nodes.py, wanvideo.py, and wrappers.py for better code clarity.
- Replaced list comprehensions with direct list conversions in nodes.py for efficiency.
- Updated memory logging format in model_management_mgpu.py to streamline data capture.
- Enhanced device management in wanvideo.py by ensuring consistent device setting and loading.
- Added linting configurations in pyproject.toml to enforce code quality standards.
- Removed unused imports and optimized existing ones across multiple files.
2026-03-06 05:59:12 -06:00

69 lines
2.1 KiB
Python

#!/usr/bin/env python3
"""Filter MultiGPU JSON logs for allocation summaries."""
import argparse
import json
import sys
from pathlib import Path
from typing import Iterable, Iterator, Dict, Any
def _write_stdout(message: str = "") -> None:
sys.stdout.write(f"{message}\n")
sys.stdout.flush()
def load_json_lines(path: Path) -> Iterator[Dict[str, Any]]:
with path.open("r", encoding="utf-8") as handle:
for line in handle:
line = line.strip()
if not line:
continue
try:
yield json.loads(line)
except json.JSONDecodeError:
continue
def is_allocation_event(entry: Dict[str, Any], keywords: Iterable[str]) -> bool:
message = entry.get("message", "")
return any(keyword in message for keyword in keywords)
def main() -> int:
parser = argparse.ArgumentParser(description="Extract allocation-related events from MultiGPU JSON logs")
parser.add_argument("logfile", type=Path, help="Path to JSONL log produced by MGPU_JSON_LOG_PATH")
parser.add_argument(
"--keywords",
nargs="*",
default=["Final Allocation String", "Total memory", "Virtual VRAM"],
help="Keywords that mark allocation events",
)
args = parser.parse_args()
entries = list(load_json_lines(args.logfile))
if not entries:
_write_stdout("No entries found in log file.")
return 0
matched = [entry for entry in entries if is_allocation_event(entry, args.keywords)]
if not matched:
_write_stdout("No allocation events matched provided keywords.")
return 0
for entry in matched:
timestamp = entry.get("timestamp", "unknown")
category = entry.get("event_category", "")
component = entry.get("component", "")
header_bits = [bit for bit in (timestamp, category, component) if bit]
header = " | ".join(header_bits) if header_bits else "allocation"
_write_stdout(f"## {header}")
_write_stdout(entry.get("message", ""))
_write_stdout()
return 0
if __name__ == "__main__":
raise SystemExit(main())