63 lines
1.9 KiB
Python
63 lines
1.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Filter MultiGPU JSON logs for allocation summaries."""
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Iterable, Iterator, Dict, Any
|
|
|
|
|
|
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:
|
|
print("No entries found in log file.")
|
|
return 0
|
|
|
|
matched = [entry for entry in entries if is_allocation_event(entry, args.keywords)]
|
|
if not matched:
|
|
print("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"
|
|
print(f"## {header}")
|
|
print(entry.get("message", ""))
|
|
print()
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|