File size: 2,982 Bytes
2f4202a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | #!/usr/bin/env python3
"""Trigger script for benign GGUFReader array parsing behavior.
Runs each parse in a child process so a no-progress parse can be timed out
without destabilizing the caller.
"""
from __future__ import annotations
import argparse
import json
import subprocess
import sys
import textwrap
import time
from pathlib import Path
PARSE_CODE = r"""
import json
import sys
import time
from importlib.metadata import version
import gguf
import numpy as np
import tracemalloc
tracemalloc.start()
start = time.perf_counter()
reader = gguf.GGUFReader(sys.argv[1], "r")
elapsed = time.perf_counter() - start
current, peak = tracemalloc.get_traced_memory()
array_field = None
for key, field in reader.fields.items():
if key.startswith("huntr.") and key.endswith("array"):
array_field = field
break
print(json.dumps({
"gguf": version("gguf"),
"numpy": np.__version__,
"elapsed_sec": elapsed,
"tracemalloc_peak_bytes": peak,
"field_count": len(reader.fields),
"array_key": array_field.name if array_field else None,
"array_part_count": len(array_field.parts) if array_field else None,
"array_data_indexes": len(array_field.data) if array_field else None,
}))
"""
def run_one(path: Path, timeout: float) -> dict[str, object]:
start = time.perf_counter()
try:
proc = subprocess.run(
[sys.executable, "-c", PARSE_CODE, str(path)],
text=True,
capture_output=True,
timeout=timeout,
check=False,
)
elapsed = time.perf_counter() - start
result: dict[str, object] = {
"path": str(path),
"timeout_sec": timeout,
"wall_sec": elapsed,
"returncode": proc.returncode,
"stdout": proc.stdout.strip(),
"stderr_tail": proc.stderr.strip()[-1000:],
}
if proc.stdout.strip().startswith("{"):
result["parsed"] = json.loads(proc.stdout)
return result
except subprocess.TimeoutExpired as exc:
return {
"path": str(path),
"timeout_sec": timeout,
"timed_out": True,
"wall_sec": time.perf_counter() - start,
"stdout": (exc.stdout or "").strip() if isinstance(exc.stdout, str) else "",
"stderr_tail": (exc.stderr or "").strip()[-1000:] if isinstance(exc.stderr, str) else "",
}
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("paths", nargs="+", type=Path)
parser.add_argument("--timeout", type=float, default=5.0)
parser.add_argument("--out", type=Path)
args = parser.parse_args()
results = [run_one(path, args.timeout) for path in args.paths]
rendered = json.dumps(results, indent=2)
print(rendered)
if args.out:
args.out.parent.mkdir(parents=True, exist_ok=True)
args.out.write_text(rendered + "\n", encoding="utf-8")
if __name__ == "__main__":
main()
|