Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
This commit is contained in:
DarkLight1337 2025-10-17 08:27:03 +00:00
parent c809ff9e72
commit 481e1e247a

View File

@ -8,6 +8,7 @@ import shlex
import signal import signal
import subprocess import subprocess
from datetime import datetime from datetime import datetime
from pathlib import Path
import pandas as pd import pandas as pd
@ -52,12 +53,12 @@ def _override_args(cmd: list[str], params: dict[str, object]):
return cmd return cmd
def _get_path_one_comb(output_dir: str, params: dict[str, object]): def _get_path_one_comb(output_dir: Path, params: dict[str, object]):
return os.path.join(output_dir, "_".join(f"{k}={v}" for k, v in params.items())) return output_dir / "_".join(f"{k}={v}" for k, v in params.items())
def _get_path_one_run(result_dir: str, run_number: int): def _get_path_one_run(result_dir: Path, run_number: int):
return os.path.join(result_dir, f"run={run_number}.json") return result_dir / f"run={run_number}.json"
def benchmark_one_run( def benchmark_one_run(
@ -65,7 +66,7 @@ def benchmark_one_run(
bench_cmd: list[str], bench_cmd: list[str],
serve_comb: dict[str, object], serve_comb: dict[str, object],
run_number: int, run_number: int,
result_dir: str, result_dir: Path,
dry_run: bool, dry_run: bool,
): ):
result_path = _get_path_one_run(result_dir, run_number) result_path = _get_path_one_run(result_dir, run_number)
@ -77,7 +78,7 @@ def benchmark_one_run(
"--result-dir", "--result-dir",
result_dir, result_dir,
"--result-filename", "--result-filename",
os.path.basename(result_path), result_path.name,
] ]
print("=" * 60) print("=" * 60)
@ -102,7 +103,7 @@ def benchmark_one_run(
# Kill entire process group # Kill entire process group
os.killpg(os.getpgid(server_process.pid), signal.SIGKILL) os.killpg(os.getpgid(server_process.pid), signal.SIGKILL)
with open(result_path, "rb") as f: with result_path.open("rb") as f:
run_data = json.load(f) run_data = json.load(f)
run_data["run_number"] = run_number run_data["run_number"] = run_number
@ -115,13 +116,13 @@ def benchmark_one_comb(
serve_cmd: list[str], serve_cmd: list[str],
bench_cmd: list[str], bench_cmd: list[str],
serve_comb: dict[str, object], serve_comb: dict[str, object],
output_dir: str, output_dir: Path,
num_runs: int, num_runs: int,
dry_run: bool, dry_run: bool,
): ):
result_dir = _get_path_one_comb(output_dir, serve_comb) result_dir = _get_path_one_comb(output_dir, serve_comb)
if not dry_run: if not dry_run:
os.makedirs(result_dir, exist_ok=True) result_dir.mkdir(parents=True, exist_ok=True)
comb_data = [ comb_data = [
benchmark_one_run( benchmark_one_run(
@ -138,22 +139,22 @@ def benchmark_one_comb(
if dry_run: if dry_run:
return None return None
with open(os.path.join(result_dir, "summary.json"), "w") as f: with (result_dir / "summary.json").open("w") as f:
json.dump(comb_data, f) json.dump(comb_data, f)
return pd.DataFrame.from_records(comb_data) return pd.DataFrame.from_records(comb_data) # type: ignore[arg-type]
def benchmark_all( def benchmark_all(
serve_cmd: list[str], serve_cmd: list[str],
bench_cmd: list[str], bench_cmd: list[str],
serve_params: list[dict[str, object]], serve_params: list[dict[str, object]],
output_dir: str, output_dir: Path,
num_runs: int, num_runs: int,
dry_run: bool, dry_run: bool,
): ):
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
output_dir = os.path.join(output_dir, timestamp) output_dir = output_dir / timestamp
result_dfs = [ result_dfs = [
benchmark_one_comb( benchmark_one_comb(
@ -171,7 +172,7 @@ def benchmark_all(
return None return None
combined_df = pd.concat(result_dfs) combined_df = pd.concat(result_dfs)
combined_df.to_csv(os.path.join(output_dir, "summary.csv")) combined_df.to_csv(output_dir / "summary.csv")
return combined_df return combined_df
@ -234,7 +235,7 @@ def main():
serve_cmd=serve_cmd, serve_cmd=serve_cmd,
bench_cmd=bench_cmd, bench_cmd=bench_cmd,
serve_params=serve_params, serve_params=serve_params,
output_dir=args.output_dir, output_dir=Path(args.output_dir),
num_runs=args.num_runs, num_runs=args.num_runs,
dry_run=args.dry_run, dry_run=args.dry_run,
) )