99 lines
3.2 KiB
Python
Executable File
99 lines
3.2 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""Generate the scenario list for the TP×DP matrix experiment.
|
||
|
||
Reads matrix.json and prints TSV lines:
|
||
|
||
mark input_len output_len concurrency num_prompts
|
||
|
||
mark is one of Y/P/N. The caller (run_bench.sh) decides how to treat each.
|
||
"""
|
||
import argparse
|
||
import json
|
||
import math
|
||
import os
|
||
from pathlib import Path
|
||
|
||
|
||
def sample_concurrency(low: int, high: int, target: int) -> list[int]:
|
||
"""Return only the low and high concurrency values in [low, high].
|
||
|
||
For the TP×DP matrix we only need the two endpoints of the concurrency
|
||
range (e.g. 1 and 128 for ISL=1024). The `target` argument is kept for
|
||
API compatibility but is ignored.
|
||
"""
|
||
assert 1 <= low <= high, f"invalid concurrency range: {low}-{high}"
|
||
if low == high:
|
||
return [low]
|
||
return [low, high]
|
||
|
||
|
||
def generate_scenarios(matrix_path: Path, mode: str, target_samples: int) -> list[dict]:
|
||
with open(matrix_path, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
|
||
matrix = data["matrix"]
|
||
concurrency_cfg = data["concurrency"]
|
||
|
||
scenarios = []
|
||
for isl_str in sorted(matrix.keys(), key=int):
|
||
dsl_map = matrix[isl_str]
|
||
low = concurrency_cfg[isl_str]["low"]
|
||
high = concurrency_cfg[isl_str]["high"]
|
||
concurrencies = sample_concurrency(low, high, target_samples)
|
||
|
||
for dsl_str in sorted(dsl_map.keys(), key=int):
|
||
mark = dsl_map[dsl_str]
|
||
if mode == "Y" and mark != "Y":
|
||
continue
|
||
if mode == "Y+P" and mark not in ("Y", "P"):
|
||
continue
|
||
# mode == "all" keeps everything, including N.
|
||
|
||
for conc in concurrencies:
|
||
scenarios.append(
|
||
{
|
||
"mark": mark,
|
||
"input_len": int(isl_str),
|
||
"output_len": int(dsl_str),
|
||
"concurrency": conc,
|
||
"num_prompts": conc * 5,
|
||
}
|
||
)
|
||
|
||
return scenarios
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--matrix", type=Path, default=Path("matrix.json"))
|
||
parser.add_argument("--mode", choices=["Y", "Y+P", "all"], default=None,
|
||
help="Scenario selection mode. Defaults to matrix.mode.")
|
||
parser.add_argument("--target-samples", type=int, default=0,
|
||
help="Target number of concurrency samples. 0 = heuristic (6-8).")
|
||
args = parser.parse_args()
|
||
|
||
with open(args.matrix, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
|
||
mode = args.mode if args.mode else data.get("mode", "Y+P")
|
||
|
||
target_samples = args.target_samples
|
||
if target_samples <= 0:
|
||
env_samples = os.getenv("CONCURRENCY_SAMPLES", "0")
|
||
try:
|
||
target_samples = int(env_samples)
|
||
except ValueError:
|
||
target_samples = 0
|
||
if target_samples <= 0:
|
||
target_samples = 7
|
||
|
||
scenarios = generate_scenarios(args.matrix, mode, target_samples)
|
||
|
||
print("mark\tinput_len\toutput_len\tconcurrency\tnum_prompts")
|
||
for s in scenarios:
|
||
print(f"{s['mark']}\t{s['input_len']}\t{s['output_len']}\t{s['concurrency']}\t{s['num_prompts']}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|