From 7e9de9c60076a752b3111ae5492e3d3f77ec90ba Mon Sep 17 00:00:00 2001 From: papertager <2567587994@qq.com> Date: Tue, 9 Jun 2026 21:20:35 +0800 Subject: [PATCH 1/2] Add benchmark log capture runner --- tools/run_benchmark_with_logs.py | 65 +++++++++++++++++++++++ unit_test/test_run_benchmark_with_logs.py | 20 +++++++ 2 files changed, 85 insertions(+) create mode 100644 tools/run_benchmark_with_logs.py create mode 100644 unit_test/test_run_benchmark_with_logs.py diff --git a/tools/run_benchmark_with_logs.py b/tools/run_benchmark_with_logs.py new file mode 100644 index 0000000..9e00288 --- /dev/null +++ b/tools/run_benchmark_with_logs.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python3 +"""Run an mcoplib benchmark command and capture reproducible logs.""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path + + +def build_command(root: Path, op: str, extra_args: list[str]) -> list[str]: + return [ + sys.executable, + str(root / "benchmark" / "mcoplib_mxbenchmark_ops.py"), + "--op", + op, + *extra_args, + ] + + +def run_with_logs(root: Path, op: str, log_root: Path, extra_args: list[str]) -> int: + run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + f"_{op}" + run_dir = log_root / run_id + run_dir.mkdir(parents=True, exist_ok=True) + + command = build_command(root, op, extra_args) + metadata = { + "generated_at": datetime.now(timezone.utc).isoformat(), + "root": str(root), + "op": op, + "command": command, + "environment": { + "MACA_PATH": os.environ.get("MACA_PATH"), + "CUDA_HOME": os.environ.get("CUDA_HOME"), + "LD_LIBRARY_PATH": os.environ.get("LD_LIBRARY_PATH"), + }, + } + (run_dir / "metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") + + with (run_dir / "stdout.log").open("w", encoding="utf-8") as out, ( + run_dir / "stderr.log" + ).open("w", encoding="utf-8") as err: + proc = subprocess.run(command, cwd=root / "benchmark", stdout=out, stderr=err, text=True) + + (run_dir / "exit_code.txt").write_text(str(proc.returncode) + "\n", encoding="utf-8") + print(f"Benchmark logs written to: {run_dir}") + return proc.returncode + + +def main() -> int: + parser = argparse.ArgumentParser(description="Run mcoplib benchmark with structured logs.") + parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1]) + parser.add_argument("--op", required=True) + parser.add_argument("--log-root", type=Path, default=Path("benchmark_logs")) + args, extra = parser.parse_known_args() + + return run_with_logs(args.root.resolve(), args.op, args.log_root.resolve(), extra) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/unit_test/test_run_benchmark_with_logs.py b/unit_test/test_run_benchmark_with_logs.py new file mode 100644 index 0000000..c9baba0 --- /dev/null +++ b/unit_test/test_run_benchmark_with_logs.py @@ -0,0 +1,20 @@ +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from tools.run_benchmark_with_logs import build_command + + +class RunBenchmarkWithLogsTest(unittest.TestCase): + def test_build_command_forwards_extra_args(self): + root = Path("/repo") + command = build_command(root, "rms_norm", ["--generate", "--csv", "out.csv"]) + + self.assertEqual(Path(command[1]), root / "benchmark" / "mcoplib_mxbenchmark_ops.py") + self.assertEqual(command[2:], ["--op", "rms_norm", "--generate", "--csv", "out.csv"]) + + +if __name__ == "__main__": + unittest.main() From 02c2eb2da39a7e217aadc720472d4854f5d2beb5 Mon Sep 17 00:00:00 2001 From: papertager <2567587994@qq.com> Date: Thu, 11 Jun 2026 00:39:38 +0800 Subject: [PATCH 2/2] Harden benchmark log capture inputs --- tools/run_benchmark_with_logs.py | 23 +++++++++++++++++++++-- unit_test/test_run_benchmark_with_logs.py | 11 ++++++++++- 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/tools/run_benchmark_with_logs.py b/tools/run_benchmark_with_logs.py index 9e00288..7a58dcc 100644 --- a/tools/run_benchmark_with_logs.py +++ b/tools/run_benchmark_with_logs.py @@ -23,6 +23,19 @@ def build_command(root: Path, op: str, extra_args: list[str]) -> list[str]: def run_with_logs(root: Path, op: str, log_root: Path, extra_args: list[str]) -> int: + if not op or not all(ch.isalnum() or ch == "_" for ch in op): + print( + f"Error: Invalid operator name '{op}'. Only alphanumeric characters and underscores are allowed.", + file=sys.stderr, + ) + return 1 + + benchmark_dir = root / "benchmark" + script_path = benchmark_dir / "mcoplib_mxbenchmark_ops.py" + if not script_path.is_file(): + print(f"Error: Benchmark script not found at '{script_path}'", file=sys.stderr) + return 1 + run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + f"_{op}" run_dir = log_root / run_id run_dir.mkdir(parents=True, exist_ok=True) @@ -37,14 +50,20 @@ def run_with_logs(root: Path, op: str, log_root: Path, extra_args: list[str]) -> "MACA_PATH": os.environ.get("MACA_PATH"), "CUDA_HOME": os.environ.get("CUDA_HOME"), "LD_LIBRARY_PATH": os.environ.get("LD_LIBRARY_PATH"), + "PYTHONPATH": os.environ.get("PYTHONPATH"), + "CUDA_VISIBLE_DEVICES": os.environ.get("CUDA_VISIBLE_DEVICES"), + "MACA_VISIBLE_DEVICES": os.environ.get("MACA_VISIBLE_DEVICES"), }, } - (run_dir / "metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") + (run_dir / "metadata.json").write_text( + json.dumps(metadata, indent=2, ensure_ascii=False), + encoding="utf-8", + ) with (run_dir / "stdout.log").open("w", encoding="utf-8") as out, ( run_dir / "stderr.log" ).open("w", encoding="utf-8") as err: - proc = subprocess.run(command, cwd=root / "benchmark", stdout=out, stderr=err, text=True) + proc = subprocess.run(command, cwd=benchmark_dir, stdout=out, stderr=err, text=True) (run_dir / "exit_code.txt").write_text(str(proc.returncode) + "\n", encoding="utf-8") print(f"Benchmark logs written to: {run_dir}") diff --git a/unit_test/test_run_benchmark_with_logs.py b/unit_test/test_run_benchmark_with_logs.py index c9baba0..5e58bad 100644 --- a/unit_test/test_run_benchmark_with_logs.py +++ b/unit_test/test_run_benchmark_with_logs.py @@ -1,10 +1,11 @@ import sys +import tempfile import unittest from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1])) -from tools.run_benchmark_with_logs import build_command +from tools.run_benchmark_with_logs import build_command, run_with_logs class RunBenchmarkWithLogsTest(unittest.TestCase): @@ -15,6 +16,14 @@ def test_build_command_forwards_extra_args(self): self.assertEqual(Path(command[1]), root / "benchmark" / "mcoplib_mxbenchmark_ops.py") self.assertEqual(command[2:], ["--op", "rms_norm", "--generate", "--csv", "out.csv"]) + def test_rejects_operator_names_with_path_traversal(self): + with tempfile.TemporaryDirectory() as tmpdir: + root = Path(tmpdir) + + code = run_with_logs(root, "../../etc", root / "logs", []) + + self.assertEqual(code, 1) + if __name__ == "__main__": unittest.main()