From 4dcd1dea0d5843447b3b253cbbb031376af3d70b Mon Sep 17 00:00:00 2001 From: dhananjai1729 Date: Thu, 19 Mar 2026 21:09:47 +0530 Subject: [PATCH 1/2] feat: parallelize rlm_query_batched subcalls using ThreadPoolExecutor When rlm_query_batched fans out multiple independent child RLM calls, they now run concurrently via a bounded thread pool instead of sequentially. This reduces wall time from N*T to ~T for N I/O-bound subcalls. Key changes: - Use ThreadPoolExecutor in _rlm_query_batched for 2+ prompts - Add max_concurrent_subcalls parameter (default 4) with semaphore-style concurrency control at each recursion level - Propagate setting from RLM -> LocalREPL -> child RLMs - Thread-safe completion collection with deterministic ordering - Skip thread pool for 0-1 prompts (no overhead) Closes #135 Co-Authored-By: Claude Opus 4.6 --- rlm/core/rlm.py | 8 +++ rlm/environments/local_repl.py | 55 +++++++++++++-- tests/test_rlm_query.py | 119 +++++++++++++++++++++++++++++++++ 3 files changed, 175 insertions(+), 7 deletions(-) diff --git a/rlm/core/rlm.py b/rlm/core/rlm.py index ce247ecdc..70c7ce056 100644 --- a/rlm/core/rlm.py +++ b/rlm/core/rlm.py @@ -70,6 +70,7 @@ def __init__( custom_sub_tools: dict[str, Any] | None = None, compaction: bool = False, compaction_threshold_pct: float = 0.85, + max_concurrent_subcalls: int = 4, on_subcall_start: Callable[[int, str, str], None] | None = None, on_subcall_complete: Callable[[int, str, float, str | None], None] | None = None, on_iteration_start: Callable[[int, int], None] | None = None, @@ -102,6 +103,8 @@ def __init__( when root context reaches compaction_threshold_pct of the model's context limit. compaction_threshold_pct: When compaction is on, trigger summarization when root message token count reaches this fraction of the model context limit (default 0.85). + max_concurrent_subcalls: Maximum number of parallel threads for rlm_query_batched subcalls. + Each child RLM runs in its own thread. Default 4. on_subcall_start: Callback fired when a child RLM starts. Args: (depth, model, prompt_preview). on_subcall_complete: Callback fired when a child RLM completes. Args: (depth, model, duration, error_or_none). on_iteration_start: Callback fired when an iteration starts. Args: (depth, iteration_num). @@ -132,6 +135,7 @@ def __init__( self.compaction = compaction self.compaction_threshold_pct = compaction_threshold_pct + self.max_concurrent_subcalls = max_concurrent_subcalls self.depth = depth self.max_depth = max_depth @@ -238,6 +242,8 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]): env_kwargs["custom_sub_tools"] = self.custom_sub_tools if self.compaction and self.environment_type == "local": env_kwargs["compaction"] = True + if self.environment_type == "local": + env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls environment: BaseEnv = get_environment(self.environment_type, env_kwargs) if self.persistent: @@ -765,6 +771,8 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion: # Propagate custom tools to children (sub_tools become the child's tools) custom_tools=self.custom_sub_tools, custom_sub_tools=self.custom_sub_tools, + # Propagate concurrency settings to children + max_concurrent_subcalls=self.max_concurrent_subcalls, # Propagate callbacks to children for nested tracking on_subcall_start=self.on_subcall_start, on_subcall_complete=self.on_subcall_complete, diff --git a/rlm/environments/local_repl.py b/rlm/environments/local_repl.py index afdde3b3d..ee9d46643 100644 --- a/rlm/environments/local_repl.py +++ b/rlm/environments/local_repl.py @@ -9,6 +9,7 @@ import time import uuid from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor, as_completed from contextlib import contextmanager from typing import Any @@ -135,12 +136,14 @@ def __init__( custom_tools: dict[str, Any] | None = None, custom_sub_tools: dict[str, Any] | None = None, compaction: bool = False, + max_concurrent_subcalls: int = 4, **kwargs, ): super().__init__(persistent=persistent, depth=depth, **kwargs) self.lm_handler_address = lm_handler_address self.subcall_fn = subcall_fn # Callback for recursive RLM calls (depth > 1 support) + self.max_concurrent_subcalls = max_concurrent_subcalls self.original_cwd = os.getcwd() self.temp_dir = tempfile.mkdtemp(prefix=f"repl_env_{uuid.uuid4()}_") self._lock = threading.Lock() @@ -316,9 +319,13 @@ def _rlm_query(self, prompt: str, model: str | None = None) -> str: return self._llm_query(prompt, model) def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> list[str]: - """Spawn recursive RLM sub-calls for multiple prompts. + """Spawn recursive RLM sub-calls for multiple prompts in parallel. + + Each prompt gets its own child RLM for deeper thinking. When multiple + prompts are provided, subcalls run concurrently using a thread pool + (bounded by max_concurrent_subcalls) since they are independent and + I/O-bound. Results are returned in the same order as input prompts. - Each prompt gets its own child RLM for deeper thinking. Falls back to llm_query_batched if no recursive capability is configured. Args: @@ -329,14 +336,48 @@ def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> li List of responses in the same order as input prompts. """ if self.subcall_fn is not None: - results = [] - for prompt in prompts: + # For 0 or 1 prompts, no need for thread pool overhead + if len(prompts) <= 1: + results = [] + for prompt in prompts: + try: + completion = self.subcall_fn(prompt, model) + self._pending_llm_calls.append(completion) + results.append(completion.response) + except Exception as e: + results.append(f"Error: RLM query failed - {e}") + return results + + # Parallel execution for multiple prompts + max_workers = min(self.max_concurrent_subcalls, len(prompts)) + # Pre-allocate result slots to preserve ordering + results: list[str] = [""] * len(prompts) + completions: list[tuple[int, RLMChatCompletion]] = [] + lock = threading.Lock() + + def _run_subcall(index: int, prompt: str) -> None: try: completion = self.subcall_fn(prompt, model) - self._pending_llm_calls.append(completion) - results.append(completion.response) + with lock: + completions.append((index, completion)) + results[index] = completion.response except Exception as e: - results.append(f"Error: RLM query failed - {e}") + results[index] = f"Error: RLM query failed - {e}" + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit(_run_subcall, i, prompt) + for i, prompt in enumerate(prompts) + ] + # Wait for all futures to complete; exceptions are captured inside _run_subcall + for future in as_completed(futures): + future.result() # Re-raises unexpected executor errors + + # Append completions in original prompt order for deterministic metadata + completions.sort(key=lambda x: x[0]) + for _, completion in completions: + self._pending_llm_calls.append(completion) + return results # Fall back to plain batched LM call if no recursive capability diff --git a/tests/test_rlm_query.py b/tests/test_rlm_query.py index 850a87fff..21790dec4 100644 --- a/tests/test_rlm_query.py +++ b/tests/test_rlm_query.py @@ -1,5 +1,7 @@ """Tests for rlm_query and rlm_query_batched in LocalREPL.""" +import threading +import time from unittest.mock import MagicMock from rlm.core.types import RLMChatCompletion, UsageSummary @@ -186,6 +188,123 @@ def test_llm_query_batched_ignores_subcall_fn(self): repl.cleanup() +class TestRlmQueryBatchedParallel: + """Tests for parallel execution of rlm_query_batched.""" + + def test_batched_runs_in_parallel(self): + """Multiple subcalls should execute concurrently, not sequentially.""" + call_times = {} + lock = threading.Lock() + + def slow_subcall(prompt, model): + start = time.monotonic() + time.sleep(0.2) # Simulate I/O-bound work + end = time.monotonic() + with lock: + call_times[prompt] = (start, end) + return _make_completion(f"answer for {prompt}") + + repl = LocalREPL(subcall_fn=slow_subcall, max_concurrent_subcalls=4) + start = time.monotonic() + result = repl.execute_code("answers = rlm_query_batched(['q1', 'q2', 'q3', 'q4'])") + wall_time = time.monotonic() - start + + assert result.stderr == "" + answers = repl.locals["answers"] + assert len(answers) == 4 + assert all("answer for" in a for a in answers) + + # If sequential, 4 × 0.2s = 0.8s minimum. Parallel should be ~0.2s. + # Use generous threshold to avoid flakiness. + assert wall_time < 0.6, f"Expected parallel execution but took {wall_time:.2f}s" + repl.cleanup() + + def test_batched_respects_max_concurrent(self): + """Thread pool should not exceed max_concurrent_subcalls.""" + max_concurrent = 2 + active_count = [] + active_lock = threading.Lock() + active = [0] + + def tracked_subcall(prompt, model): + with active_lock: + active[0] += 1 + active_count.append(active[0]) + time.sleep(0.1) + with active_lock: + active[0] -= 1 + return _make_completion(f"ok {prompt}") + + repl = LocalREPL(subcall_fn=tracked_subcall, max_concurrent_subcalls=max_concurrent) + result = repl.execute_code("answers = rlm_query_batched(['a', 'b', 'c', 'd'])") + assert result.stderr == "" + assert max(active_count) <= max_concurrent + repl.cleanup() + + def test_batched_preserves_order_parallel(self): + """Results must be in the same order as input prompts even with parallel execution.""" + import random + + def varying_delay_subcall(prompt, model): + time.sleep(random.uniform(0.01, 0.1)) + return _make_completion(f"result-{prompt}") + + repl = LocalREPL(subcall_fn=varying_delay_subcall, max_concurrent_subcalls=4) + result = repl.execute_code("answers = rlm_query_batched(['p0', 'p1', 'p2', 'p3'])") + assert result.stderr == "" + answers = repl.locals["answers"] + assert answers == ["result-p0", "result-p1", "result-p2", "result-p3"] + repl.cleanup() + + def test_batched_parallel_partial_failure(self): + """Failures in some threads should not affect others.""" + call_count = [0] + lock = threading.Lock() + + def sometimes_fail(prompt, model): + with lock: + call_count[0] += 1 + if prompt == "fail_me": + raise RuntimeError("intentional failure") + time.sleep(0.05) + return _make_completion(f"ok-{prompt}") + + repl = LocalREPL(subcall_fn=sometimes_fail, max_concurrent_subcalls=4) + result = repl.execute_code("answers = rlm_query_batched(['a', 'fail_me', 'c'])") + assert result.stderr == "" + answers = repl.locals["answers"] + assert answers[0] == "ok-a" + assert "Error" in answers[1] + assert "intentional failure" in answers[1] + assert answers[2] == "ok-c" + repl.cleanup() + + def test_batched_pending_calls_ordered(self): + """Pending LLM calls should be appended in prompt order for deterministic metadata.""" + def delayed_subcall(prompt, model): + # Reverse delay so later prompts finish first + delay = 0.1 if prompt == "first" else 0.01 + time.sleep(delay) + return _make_completion(f"resp-{prompt}") + + repl = LocalREPL(subcall_fn=delayed_subcall, max_concurrent_subcalls=4) + result = repl.execute_code("answers = rlm_query_batched(['first', 'second'])") + assert result.stderr == "" + assert len(result.rlm_calls) == 2 + assert result.rlm_calls[0].response == "resp-first" + assert result.rlm_calls[1].response == "resp-second" + repl.cleanup() + + def test_single_prompt_skips_thread_pool(self): + """Single prompt should not use ThreadPoolExecutor (no overhead).""" + subcall_fn = MagicMock(return_value=_make_completion("solo")) + repl = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=4) + result = repl.execute_code("answers = rlm_query_batched(['only'])") + assert repl.locals["answers"] == ["solo"] + subcall_fn.assert_called_once_with("only", None) + repl.cleanup() + + class TestRlmQueryScaffoldRestoration: """Test that rlm_query and rlm_query_batched are restored after overwrite.""" From 4090723dea07747912f6e5732113cd6b7cf970da Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Thu, 19 Mar 2026 14:31:40 -0700 Subject: [PATCH 2/2] move max_concurrent_subcalls as a general REPL property. add more tests --- rlm/core/rlm.py | 3 +- rlm/environments/base_env.py | 5 +- rlm/environments/local_repl.py | 11 ++- tests/test_rlm_query.py | 166 ++++++++++++++++++++++++++++++++- 4 files changed, 177 insertions(+), 8 deletions(-) diff --git a/rlm/core/rlm.py b/rlm/core/rlm.py index 70c7ce056..4a8e42203 100644 --- a/rlm/core/rlm.py +++ b/rlm/core/rlm.py @@ -242,8 +242,7 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]): env_kwargs["custom_sub_tools"] = self.custom_sub_tools if self.compaction and self.environment_type == "local": env_kwargs["compaction"] = True - if self.environment_type == "local": - env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls + env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls environment: BaseEnv = get_environment(self.environment_type, env_kwargs) if self.persistent: diff --git a/rlm/environments/base_env.py b/rlm/environments/base_env.py index afd6387f2..977a417f4 100644 --- a/rlm/environments/base_env.py +++ b/rlm/environments/base_env.py @@ -212,9 +212,12 @@ class BaseEnv(ABC): } """ - def __init__(self, persistent: bool = False, depth: int = 1, **kwargs): + def __init__( + self, persistent: bool = False, depth: int = 1, max_concurrent_subcalls: int = 4, **kwargs + ): self.persistent = persistent self.depth = depth + self.max_concurrent_subcalls = max_concurrent_subcalls self.kwargs = kwargs @abstractmethod diff --git a/rlm/environments/local_repl.py b/rlm/environments/local_repl.py index ee9d46643..fe33b3882 100644 --- a/rlm/environments/local_repl.py +++ b/rlm/environments/local_repl.py @@ -139,11 +139,15 @@ def __init__( max_concurrent_subcalls: int = 4, **kwargs, ): - super().__init__(persistent=persistent, depth=depth, **kwargs) + super().__init__( + persistent=persistent, + depth=depth, + max_concurrent_subcalls=max_concurrent_subcalls, + **kwargs, + ) self.lm_handler_address = lm_handler_address self.subcall_fn = subcall_fn # Callback for recursive RLM calls (depth > 1 support) - self.max_concurrent_subcalls = max_concurrent_subcalls self.original_cwd = os.getcwd() self.temp_dir = tempfile.mkdtemp(prefix=f"repl_env_{uuid.uuid4()}_") self._lock = threading.Lock() @@ -366,8 +370,7 @@ def _run_subcall(index: int, prompt: str) -> None: with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [ - executor.submit(_run_subcall, i, prompt) - for i, prompt in enumerate(prompts) + executor.submit(_run_subcall, i, prompt) for i, prompt in enumerate(prompts) ] # Wait for all futures to complete; exceptions are captured inside _run_subcall for future in as_completed(futures): diff --git a/tests/test_rlm_query.py b/tests/test_rlm_query.py index 21790dec4..a62cacdd4 100644 --- a/tests/test_rlm_query.py +++ b/tests/test_rlm_query.py @@ -281,6 +281,7 @@ def sometimes_fail(prompt, model): def test_batched_pending_calls_ordered(self): """Pending LLM calls should be appended in prompt order for deterministic metadata.""" + def delayed_subcall(prompt, model): # Reverse delay so later prompts finish first delay = 0.1 if prompt == "first" else 0.01 @@ -299,12 +300,175 @@ def test_single_prompt_skips_thread_pool(self): """Single prompt should not use ThreadPoolExecutor (no overhead).""" subcall_fn = MagicMock(return_value=_make_completion("solo")) repl = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=4) - result = repl.execute_code("answers = rlm_query_batched(['only'])") + repl.execute_code("answers = rlm_query_batched(['only'])") assert repl.locals["answers"] == ["solo"] subcall_fn.assert_called_once_with("only", None) repl.cleanup() +class TestMaxConcurrentSubcallsBounds: + """Tests for edge cases and boundary conditions of max_concurrent_subcalls.""" + + def test_max_concurrent_subcalls_default_value(self): + """Default max_concurrent_subcalls should be 4.""" + repl = LocalREPL() + assert repl.max_concurrent_subcalls == 4 + repl.cleanup() + + def test_max_concurrent_subcalls_custom_value(self): + """Should accept custom max_concurrent_subcalls values.""" + repl = LocalREPL(max_concurrent_subcalls=8) + assert repl.max_concurrent_subcalls == 8 + repl.cleanup() + + def test_max_concurrent_subcalls_one(self): + """max_concurrent_subcalls=1 should force sequential execution.""" + order = [] + lock = threading.Lock() + + def sequential_subcall(prompt, model): + with lock: + order.append(f"start-{prompt}") + time.sleep(0.05) + with lock: + order.append(f"end-{prompt}") + return _make_completion(f"done-{prompt}") + + repl = LocalREPL(subcall_fn=sequential_subcall, max_concurrent_subcalls=1) + result = repl.execute_code("answers = rlm_query_batched(['a', 'b', 'c'])") + assert result.stderr == "" + answers = repl.locals["answers"] + assert answers == ["done-a", "done-b", "done-c"] + # With max_workers=1, tasks run one at a time so each start + # should be followed by its end before the next start + for i in range(0, len(order) - 1, 2): + prompt_id = order[i].split("-", 1)[1] + assert order[i] == f"start-{prompt_id}" + assert order[i + 1] == f"end-{prompt_id}" + repl.cleanup() + + def test_max_concurrent_subcalls_larger_than_prompts(self): + """max_concurrent_subcalls larger than prompt count should work fine.""" + subcall_fn = MagicMock(return_value=_make_completion("ok")) + repl = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=100) + result = repl.execute_code("answers = rlm_query_batched(['a', 'b'])") + assert result.stderr == "" + assert repl.locals["answers"] == ["ok", "ok"] + assert subcall_fn.call_count == 2 + repl.cleanup() + + def test_max_concurrent_subcalls_exact_match(self): + """max_concurrent_subcalls == prompt count should allow full parallelism.""" + active_count = [] + active_lock = threading.Lock() + active = [0] + + def tracked_subcall(prompt, model): + with active_lock: + active[0] += 1 + active_count.append(active[0]) + time.sleep(0.1) + with active_lock: + active[0] -= 1 + return _make_completion(f"ok {prompt}") + + repl = LocalREPL(subcall_fn=tracked_subcall, max_concurrent_subcalls=3) + result = repl.execute_code("answers = rlm_query_batched(['a', 'b', 'c'])") + assert result.stderr == "" + # All 3 should be able to run at once + assert max(active_count) <= 3 + repl.cleanup() + + def test_batched_all_failures_parallel(self): + """All subcalls failing in parallel should return all error strings.""" + + def always_fail(prompt, model): + raise ValueError(f"fail-{prompt}") + + repl = LocalREPL(subcall_fn=always_fail, max_concurrent_subcalls=4) + result = repl.execute_code("answers = rlm_query_batched(['x', 'y', 'z'])") + assert result.stderr == "" + answers = repl.locals["answers"] + assert len(answers) == 3 + for i, prompt in enumerate(["x", "y", "z"]): + assert "Error" in answers[i] + assert f"fail-{prompt}" in answers[i] + repl.cleanup() + + def test_batched_large_batch_with_low_concurrency(self): + """Many prompts with low concurrency should still complete correctly.""" + max_concurrent = 2 + active_count = [] + active_lock = threading.Lock() + active = [0] + + def tracked_subcall(prompt, model): + with active_lock: + active[0] += 1 + active_count.append(active[0]) + time.sleep(0.02) + with active_lock: + active[0] -= 1 + return _make_completion(f"r-{prompt}") + + prompts = [f"p{i}" for i in range(10)] + repl = LocalREPL(subcall_fn=tracked_subcall, max_concurrent_subcalls=max_concurrent) + result = repl.execute_code(f"answers = rlm_query_batched({prompts!r})") + assert result.stderr == "" + answers = repl.locals["answers"] + assert len(answers) == 10 + assert answers == [f"r-p{i}" for i in range(10)] + # Concurrency cap should be respected + assert max(active_count) <= max_concurrent + repl.cleanup() + + def test_batched_empty_prompts_with_parallel(self): + """Empty prompt list should return empty list regardless of concurrency setting.""" + subcall_fn = MagicMock() + repl = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=4) + repl.execute_code("answers = rlm_query_batched([])") + assert repl.locals["answers"] == [] + subcall_fn.assert_not_called() + repl.cleanup() + + def test_pending_calls_exclude_failures(self): + """Failed subcalls should not appear in pending_llm_calls.""" + + def fail_second(prompt, model): + if prompt == "bad": + raise RuntimeError("boom") + return _make_completion(f"ok-{prompt}") + + repl = LocalREPL(subcall_fn=fail_second, max_concurrent_subcalls=4) + result = repl.execute_code("answers = rlm_query_batched(['good1', 'bad', 'good2'])") + assert result.stderr == "" + # Only 2 successful completions should be in rlm_calls + assert len(result.rlm_calls) == 2 + assert result.rlm_calls[0].response == "ok-good1" + assert result.rlm_calls[1].response == "ok-good2" + repl.cleanup() + + +class TestMaxConcurrentSubcallsOnBaseEnv: + """Tests that max_concurrent_subcalls is a property of BaseEnv, not just LocalREPL.""" + + def test_base_env_has_max_concurrent_subcalls(self): + """BaseEnv should accept and store max_concurrent_subcalls.""" + + # BaseEnv is abstract, so we test via LocalREPL which inherits from it + repl = LocalREPL(max_concurrent_subcalls=16) + assert repl.max_concurrent_subcalls == 16 + # Verify it's set via the BaseEnv.__init__ chain + assert hasattr(repl, "max_concurrent_subcalls") + repl.cleanup() + + def test_base_env_default_max_concurrent_subcalls(self): + """BaseEnv default should be 4.""" + repl = LocalREPL() + assert repl.max_concurrent_subcalls == 4 + repl.cleanup() + + class TestRlmQueryScaffoldRestoration: """Test that rlm_query and rlm_query_batched are restored after overwrite."""