bench: add --max-tokens and --max-concurrency overrides; plumb into runner

Author: benjibc

eval_protocol/benchmarks/registry.py

@@ -74,6 +74,7 @@ def _runner(
             max_rows: Optional[int | str] = None,
             num_runs: Optional[int] = None,
             input_params_override: Optional[Dict[str, Any]] = None,
+            max_concurrency: Optional[int] = None,
         ) -> Any:
             # Map convenience flags to EP_* env used by the pytest flow
             if print_summary:
@@ -122,6 +123,8 @@ def _deep_update(base: Dict[str, Any], over: Dict[str, Any]) -> Dict[str, Any]:
             max_dataset_rows = ep_config.get("max_dataset_rows")
             mcp_config_path = ep_config.get("mcp_config_path")
             max_concurrent_rollouts = ep_config.get("max_concurrent_rollouts")
+            if max_concurrency is not None:
+                max_concurrent_rollouts = int(max_concurrency)
             server_script_path = ep_config.get("server_script_path")
             steps = ep_config.get("steps")
             mode = ep_config.get("mode")

eval_protocol/benchmarks/run.py

@@ -38,6 +38,8 @@ def _parse_args() -> argparse.Namespace:
         help="Limit rows: integer or 'all' for no limit (maps to EP_MAX_DATASET_ROWS)",
     )
     parser.add_argument("--num-runs", type=int, help="Override num_runs if provided")
+    parser.add_argument("--max-tokens", type=int, help="Override max_tokens for generation requests")
+    parser.add_argument("--max-concurrency", type=int, help="Override max concurrent rollouts")
     # Allow overriding reasoning effort explicitly (low/medium/high). If omitted, suite default is used.
     # Already mapped by --reasoning-effort above.
     return parser.parse_args()
@@ -73,13 +75,20 @@ def main() -> int:
             max_rows = int(args.max_rows)
         except Exception:
             max_rows = str(args.max_rows)
+    # Build input params override if needed
+    ip_override = {}
+    if args.max_tokens is not None:
+        ip_override["max_tokens"] = int(args.max_tokens)
+
     _ = runner(
         model=args.model,
         print_summary=args.print_summary,
         out=args.out,
         reasoning_effort=args.reasoning_effort,
         max_rows=max_rows,
         num_runs=args.num_runs,
+        input_params_override=(ip_override or None),
+        max_concurrency=args.max_concurrency,
     )
     # Non-zero exit on failure gate is handled within the runner via assertions
     return 0