Skip to content

Commit 072d1b1

Browse files
authored
Merge pull request #23 from ShaneIsley/claude/rlm-testing-framework-byfLB
feat: add sample timeout to prevent indefinite hangs
2 parents 0a3739c + 72b288e commit 072d1b1

3 files changed

Lines changed: 34 additions & 2 deletions

File tree

‎benchmarks/cli.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -365,6 +365,12 @@ def main():
365365
help="Progress display mode: auto (uses tqdm if available), tqdm (progress bar), "
366366
"simple (periodic status), none (quiet). Default: auto",
367367
)
368+
run_parser.add_argument(
369+
"--sample-timeout",
370+
type=int,
371+
default=300,
372+
help="Timeout per sample in seconds (default: 300 = 5 minutes)",
373+
)
368374

369375
# Benchmark-specific options for run
370376
run_parser.add_argument("--context-length", type=int, default=100_000, help="NIAH context len")
@@ -470,6 +476,7 @@ def cmd_run(args: argparse.Namespace) -> int:
470476
log_dir=args.log_dir,
471477
max_workers=args.max_workers,
472478
progress=args.progress,
479+
sample_timeout=args.sample_timeout,
473480
)
474481

475482
for benchmark in benchmarks:

‎benchmarks/runner.py‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@
3838
("no such model", "Model not found. Check available models for your account."),
3939
# Rate limiting (without retry-after)
4040
("rate_limit", "Rate limited. Consider reducing --max-workers or adding delays."),
41+
# Timeout
42+
("timed out", "Sample timed out. Try --sample-timeout to increase, or reduce --context-length."),
4143
]
4244

4345

@@ -119,6 +121,7 @@ class RunnerConfig:
119121
max_workers: int = 1 # Number of parallel workers (1 = sequential)
120122
progress: str = "auto" # Progress display: "auto", "tqdm", "simple", "none"
121123
progress_callback: ProgressCallback | None = None # Custom progress callback
124+
sample_timeout: int = 300 # Timeout per sample in seconds (default: 5 minutes)
122125
backend_kwargs: dict[str, Any] = field(default_factory=dict)
123126
environment_kwargs: dict[str, Any] = field(default_factory=dict)
124127

@@ -157,6 +160,7 @@ def __init__(
157160
max_workers: int = 1,
158161
progress: str = "auto",
159162
progress_callback: ProgressCallback | None = None,
163+
sample_timeout: int = 300,
160164
**kwargs,
161165
):
162166
"""Initialize runner with configuration.
@@ -176,6 +180,7 @@ def __init__(
176180
- "none": No progress output
177181
progress_callback: Custom callback for progress updates.
178182
Signature: (completed, total, sample_result, stats) -> None
183+
sample_timeout: Timeout per sample in seconds (default: 300 = 5 minutes).
179184
**kwargs: Additional backend or environment kwargs.
180185
"""
181186
self.config = RunnerConfig(
@@ -188,6 +193,7 @@ def __init__(
188193
max_workers=max_workers,
189194
progress=progress,
190195
progress_callback=progress_callback,
196+
sample_timeout=sample_timeout,
191197
backend_kwargs={"model_name": model, **kwargs.get("backend_kwargs", {})},
192198
environment_kwargs=kwargs.get("environment_kwargs", {}),
193199
)
@@ -613,14 +619,24 @@ def _run_sample(
613619
inference_fn: Callable[[BenchmarkSample], tuple[str, dict[str, Any]]],
614620
benchmark: Benchmark,
615621
) -> SampleResult:
616-
"""Run a single sample and evaluate."""
622+
"""Run a single sample and evaluate with timeout."""
623+
from concurrent.futures import TimeoutError as FuturesTimeoutError
624+
617625
start_time = time.time()
618626
error = None
619627
prediction = ""
620628
metadata: dict[str, Any] = {}
629+
timeout = self.config.sample_timeout
621630

622631
try:
623-
prediction, metadata = inference_fn(sample)
632+
# Use a thread pool to enforce timeout
633+
with ThreadPoolExecutor(max_workers=1) as executor:
634+
future = executor.submit(inference_fn, sample)
635+
try:
636+
prediction, metadata = future.result(timeout=timeout)
637+
except FuturesTimeoutError:
638+
error = f"Sample timed out after {timeout}s. Try --sample-timeout to increase."
639+
prediction = ""
624640
except Exception as e:
625641
error = str(e)
626642
prediction = ""

‎tests/benchmarks/test_benchmarks.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -615,6 +615,15 @@ def test_classify_generic_error(self):
615615
is_fatal, suggestion = classify_error(error)
616616
assert not is_fatal
617617

618+
def test_classify_timeout_error(self):
619+
"""Test that timeout errors are classified as fatal with helpful suggestion."""
620+
from benchmarks.runner import classify_error
621+
622+
error = "Sample timed out after 300s"
623+
is_fatal, suggestion = classify_error(error)
624+
assert is_fatal
625+
assert "timeout" in suggestion.lower() or "context" in suggestion.lower()
626+
618627

619628
class TestBenchmarkIntegration:
620629
"""Integration tests for benchmark framework."""

0 commit comments

Comments
 (0)