Add all options to plugin

Author: xzrderek

eval_protocol/pytest/evaluation_test.py

@@ -50,6 +50,7 @@
     generate_parameter_combinations,
     log_eval_status_and_rows,
     parse_ep_max_rows,
+    parse_ep_num_runs,
     rollout_processor_with_retry,
     sanitize_filename,
 )
@@ -456,7 +457,10 @@ def create_wrapper_with_signature() -> Callable:
 
             async def wrapper_body(**kwargs):
                 eval_metadata = None
-                all_results: List[List[EvaluationRow]] = [[] for _ in range(num_runs)]
+
+                # Apply environment override for num_runs if present
+                effective_num_runs = parse_ep_num_runs(num_runs)
+                all_results: List[List[EvaluationRow]] = [[] for _ in range(effective_num_runs)]
 
                 experiment_id = generate_id()
 
@@ -530,7 +534,7 @@ def _log_eval_error(
                         name=test_func.__name__,
                         description=test_func.__doc__,
                         status="running",
-                        num_runs=num_runs,
+                        num_runs=effective_num_runs,
                         aggregation_method=aggregation_method,
                         passed_threshold=threshold,
                         passed=None,
@@ -564,7 +568,7 @@ def _log_eval_error(
                         exception_handler_config=exception_handler_config,
                     )
 
-                    for i in range(num_runs):
+                    for i in range(effective_num_runs):
                         # Regenerate outputs each run by deep-copying the pristine dataset
                         # so model responses are not reused across runs.
                         run_id = generate_id()
@@ -693,7 +697,8 @@ async def _collect_result(config, lst):
                     # rollout_id is used to differentiate the result from different completion_params
                     if mode == "groupwise":
                         results_by_group = [
-                            [[] for _ in range(num_runs)] for _ in range(len(original_completion_params_list))
+                            [[] for _ in range(effective_num_runs)]
+                            for _ in range(len(original_completion_params_list))
                         ]
                         for i_run, result in enumerate(all_results):
                             for r in result:
@@ -708,7 +713,7 @@ async def _collect_result(config, lst):
                                 mode,
                                 original_completion_params_list[rollout_id],
                                 test_func.__name__,
-                                num_runs,
+                                effective_num_runs,
                             )
                     else:
                         postprocess(
@@ -719,7 +724,7 @@ async def _collect_result(config, lst):
                             mode,
                             completion_params,
                             test_func.__name__,
-                            num_runs,
+                            effective_num_runs,
                         )
 
                 except AssertionError:

eval_protocol/pytest/plugin.py

@@ -28,6 +28,12 @@ def pytest_addoption(parser) -> None:
             "Pass an integer (e.g., 2, 50) or 'all' for no limit."
         ),
     )
+    group.addoption(
+        "--ep-num-runs",
+        action="store",
+        default=None,
+        help=("Override the number of runs for evaluation_test. Pass an integer (e.g., 1, 5, 10)."),
+    )
     group.addoption(
         "--ep-print-summary",
         action="store_true",
@@ -92,6 +98,20 @@ def _normalize_max_rows(val: Optional[str]) -> Optional[str]:
         return None
 
 
+def _normalize_num_runs(val: Optional[str]) -> Optional[str]:
+    if val is None:
+        return None
+    s = val.strip()
+    # Validate int; if invalid, ignore and return None (no override)
+    try:
+        num = int(s)
+        if num <= 0:
+            return None  # num_runs must be positive
+        return str(num)
+    except ValueError:
+        return None
+
+
 def pytest_configure(config) -> None:
     # Quiet LiteLLM INFO spam early in pytest session unless user set a level
     try:
@@ -110,6 +130,11 @@ def pytest_configure(config) -> None:
     if norm is not None:
         os.environ["EP_MAX_DATASET_ROWS"] = norm
 
+    num_runs_val = config.getoption("--ep-num-runs")
+    norm_runs = _normalize_num_runs(num_runs_val)
+    if norm_runs is not None:
+        os.environ["EP_NUM_RUNS"] = norm_runs
+
     if config.getoption("--ep-print-summary"):
         os.environ["EP_PRINT_SUMMARY"] = "1"
 

eval_protocol/pytest/utils.py

@@ -139,17 +139,24 @@ def log_eval_status_and_rows(
 
 
 def parse_ep_max_rows(default_value: Optional[int]) -> Optional[int]:
-    """Read EP_MAX_DATASET_ROWS env override as int or None."""
+    """Read EP_MAX_DATASET_ROWS env override as int or None.
+
+    Assumes the environment variable was already validated by plugin.py.
+    """
     raw = os.getenv("EP_MAX_DATASET_ROWS")
     if raw is None:
         return default_value
-    s = raw.strip().lower()
-    if s == "none":
-        return None
-    try:
-        return int(s)
-    except ValueError:
-        return default_value
+    # plugin.py stores "None" as string for the "all" case
+    return None if raw.lower() == "none" else int(raw)
+
+
+def parse_ep_num_runs(default_value: int) -> int:
+    """Read EP_NUM_RUNS env override as int.
+
+    Assumes the environment variable was already validated by plugin.py.
+    """
+    raw = os.getenv("EP_NUM_RUNS")
+    return int(raw) if raw is not None else default_value
 
 
 def deep_update_dict(base: dict, override: dict) -> dict: