106 lines
3.4 KiB
Python
Executable File
106 lines
3.4 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 a sorted list of concurrency values in [low, high].
|
||
|
||
If the range is small, return every integer. Otherwise sample roughly
|
||
`target` points linearly between low and high (inclusive).
|
||
"""
|
||
assert 1 <= low <= high, f"invalid concurrency range: {low}-{high}"
|
||
if high - low + 1 <= target:
|
||
return list(range(low, high + 1))
|
||
|
||
points = set()
|
||
points.add(low)
|
||
points.add(high)
|
||
step = (high - low) / (target - 1)
|
||
for i in range(1, target - 1):
|
||
v = low + round(step * i)
|
||
points.add(max(low, min(high, v)))
|
||
return sorted(points)
|
||
|
||
|
||
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()
|