Skip to content

Commit 4dcd1de

Browse files
dhananjai1729claude
authored andcommitted
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 <noreply@anthropic.com>
1 parent 2e5764b commit 4dcd1de

3 files changed

Lines changed: 175 additions & 7 deletions

File tree

‎rlm/core/rlm.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ def __init__(
7070
custom_sub_tools: dict[str, Any] | None = None,
7171
compaction: bool = False,
7272
compaction_threshold_pct: float = 0.85,
73+
max_concurrent_subcalls: int = 4,
7374
on_subcall_start: Callable[[int, str, str], None] | None = None,
7475
on_subcall_complete: Callable[[int, str, float, str | None], None] | None = None,
7576
on_iteration_start: Callable[[int, int], None] | None = None,
@@ -102,6 +103,8 @@ def __init__(
102103
when root context reaches compaction_threshold_pct of the model's context limit.
103104
compaction_threshold_pct: When compaction is on, trigger summarization when root
104105
message token count reaches this fraction of the model context limit (default 0.85).
106+
max_concurrent_subcalls: Maximum number of parallel threads for rlm_query_batched subcalls.
107+
Each child RLM runs in its own thread. Default 4.
105108
on_subcall_start: Callback fired when a child RLM starts. Args: (depth, model, prompt_preview).
106109
on_subcall_complete: Callback fired when a child RLM completes. Args: (depth, model, duration, error_or_none).
107110
on_iteration_start: Callback fired when an iteration starts. Args: (depth, iteration_num).
@@ -132,6 +135,7 @@ def __init__(
132135

133136
self.compaction = compaction
134137
self.compaction_threshold_pct = compaction_threshold_pct
138+
self.max_concurrent_subcalls = max_concurrent_subcalls
135139

136140
self.depth = depth
137141
self.max_depth = max_depth
@@ -238,6 +242,8 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]):
238242
env_kwargs["custom_sub_tools"] = self.custom_sub_tools
239243
if self.compaction and self.environment_type == "local":
240244
env_kwargs["compaction"] = True
245+
if self.environment_type == "local":
246+
env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls
241247
environment: BaseEnv = get_environment(self.environment_type, env_kwargs)
242248

243249
if self.persistent:
@@ -765,6 +771,8 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion:
765771
# Propagate custom tools to children (sub_tools become the child's tools)
766772
custom_tools=self.custom_sub_tools,
767773
custom_sub_tools=self.custom_sub_tools,
774+
# Propagate concurrency settings to children
775+
max_concurrent_subcalls=self.max_concurrent_subcalls,
768776
# Propagate callbacks to children for nested tracking
769777
on_subcall_start=self.on_subcall_start,
770778
on_subcall_complete=self.on_subcall_complete,

‎rlm/environments/local_repl.py‎

Lines changed: 48 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import time
1010
import uuid
1111
from collections.abc import Callable
12+
from concurrent.futures import ThreadPoolExecutor, as_completed
1213
from contextlib import contextmanager
1314
from typing import Any
1415

@@ -135,12 +136,14 @@ def __init__(
135136
custom_tools: dict[str, Any] | None = None,
136137
custom_sub_tools: dict[str, Any] | None = None,
137138
compaction: bool = False,
139+
max_concurrent_subcalls: int = 4,
138140
**kwargs,
139141
):
140142
super().__init__(persistent=persistent, depth=depth, **kwargs)
141143

142144
self.lm_handler_address = lm_handler_address
143145
self.subcall_fn = subcall_fn # Callback for recursive RLM calls (depth > 1 support)
146+
self.max_concurrent_subcalls = max_concurrent_subcalls
144147
self.original_cwd = os.getcwd()
145148
self.temp_dir = tempfile.mkdtemp(prefix=f"repl_env_{uuid.uuid4()}_")
146149
self._lock = threading.Lock()
@@ -316,9 +319,13 @@ def _rlm_query(self, prompt: str, model: str | None = None) -> str:
316319
return self._llm_query(prompt, model)
317320

318321
def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> list[str]:
319-
"""Spawn recursive RLM sub-calls for multiple prompts.
322+
"""Spawn recursive RLM sub-calls for multiple prompts in parallel.
323+
324+
Each prompt gets its own child RLM for deeper thinking. When multiple
325+
prompts are provided, subcalls run concurrently using a thread pool
326+
(bounded by max_concurrent_subcalls) since they are independent and
327+
I/O-bound. Results are returned in the same order as input prompts.
320328
321-
Each prompt gets its own child RLM for deeper thinking.
322329
Falls back to llm_query_batched if no recursive capability is configured.
323330
324331
Args:
@@ -329,14 +336,48 @@ def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> li
329336
List of responses in the same order as input prompts.
330337
"""
331338
if self.subcall_fn is not None:
332-
results = []
333-
for prompt in prompts:
339+
# For 0 or 1 prompts, no need for thread pool overhead
340+
if len(prompts) <= 1:
341+
results = []
342+
for prompt in prompts:
343+
try:
344+
completion = self.subcall_fn(prompt, model)
345+
self._pending_llm_calls.append(completion)
346+
results.append(completion.response)
347+
except Exception as e:
348+
results.append(f"Error: RLM query failed - {e}")
349+
return results
350+
351+
# Parallel execution for multiple prompts
352+
max_workers = min(self.max_concurrent_subcalls, len(prompts))
353+
# Pre-allocate result slots to preserve ordering
354+
results: list[str] = [""] * len(prompts)
355+
completions: list[tuple[int, RLMChatCompletion]] = []
356+
lock = threading.Lock()
357+
358+
def _run_subcall(index: int, prompt: str) -> None:
334359
try:
335360
completion = self.subcall_fn(prompt, model)
336-
self._pending_llm_calls.append(completion)
337-
results.append(completion.response)
361+
with lock:
362+
completions.append((index, completion))
363+
results[index] = completion.response
338364
except Exception as e:
339-
results.append(f"Error: RLM query failed - {e}")
365+
results[index] = f"Error: RLM query failed - {e}"
366+
367+
with ThreadPoolExecutor(max_workers=max_workers) as executor:
368+
futures = [
369+
executor.submit(_run_subcall, i, prompt)
370+
for i, prompt in enumerate(prompts)
371+
]
372+
# Wait for all futures to complete; exceptions are captured inside _run_subcall
373+
for future in as_completed(futures):
374+
future.result() # Re-raises unexpected executor errors
375+
376+
# Append completions in original prompt order for deterministic metadata
377+
completions.sort(key=lambda x: x[0])
378+
for _, completion in completions:
379+
self._pending_llm_calls.append(completion)
380+
340381
return results
341382

342383
# Fall back to plain batched LM call if no recursive capability

‎tests/test_rlm_query.py‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
"""Tests for rlm_query and rlm_query_batched in LocalREPL."""
22

3+
import threading
4+
import time
35
from unittest.mock import MagicMock
46

57
from rlm.core.types import RLMChatCompletion, UsageSummary
@@ -186,6 +188,123 @@ def test_llm_query_batched_ignores_subcall_fn(self):
186188
repl.cleanup()
187189

188190

191+
class TestRlmQueryBatchedParallel:
192+
"""Tests for parallel execution of rlm_query_batched."""
193+
194+
def test_batched_runs_in_parallel(self):
195+
"""Multiple subcalls should execute concurrently, not sequentially."""
196+
call_times = {}
197+
lock = threading.Lock()
198+
199+
def slow_subcall(prompt, model):
200+
start = time.monotonic()
201+
time.sleep(0.2) # Simulate I/O-bound work
202+
end = time.monotonic()
203+
with lock:
204+
call_times[prompt] = (start, end)
205+
return _make_completion(f"answer for {prompt}")
206+
207+
repl = LocalREPL(subcall_fn=slow_subcall, max_concurrent_subcalls=4)
208+
start = time.monotonic()
209+
result = repl.execute_code("answers = rlm_query_batched(['q1', 'q2', 'q3', 'q4'])")
210+
wall_time = time.monotonic() - start
211+
212+
assert result.stderr == ""
213+
answers = repl.locals["answers"]
214+
assert len(answers) == 4
215+
assert all("answer for" in a for a in answers)
216+
217+
# If sequential, 4 × 0.2s = 0.8s minimum. Parallel should be ~0.2s.
218+
# Use generous threshold to avoid flakiness.
219+
assert wall_time < 0.6, f"Expected parallel execution but took {wall_time:.2f}s"
220+
repl.cleanup()
221+
222+
def test_batched_respects_max_concurrent(self):
223+
"""Thread pool should not exceed max_concurrent_subcalls."""
224+
max_concurrent = 2
225+
active_count = []
226+
active_lock = threading.Lock()
227+
active = [0]
228+
229+
def tracked_subcall(prompt, model):
230+
with active_lock:
231+
active[0] += 1
232+
active_count.append(active[0])
233+
time.sleep(0.1)
234+
with active_lock:
235+
active[0] -= 1
236+
return _make_completion(f"ok {prompt}")
237+
238+
repl = LocalREPL(subcall_fn=tracked_subcall, max_concurrent_subcalls=max_concurrent)
239+
result = repl.execute_code("answers = rlm_query_batched(['a', 'b', 'c', 'd'])")
240+
assert result.stderr == ""
241+
assert max(active_count) <= max_concurrent
242+
repl.cleanup()
243+
244+
def test_batched_preserves_order_parallel(self):
245+
"""Results must be in the same order as input prompts even with parallel execution."""
246+
import random
247+
248+
def varying_delay_subcall(prompt, model):
249+
time.sleep(random.uniform(0.01, 0.1))
250+
return _make_completion(f"result-{prompt}")
251+
252+
repl = LocalREPL(subcall_fn=varying_delay_subcall, max_concurrent_subcalls=4)
253+
result = repl.execute_code("answers = rlm_query_batched(['p0', 'p1', 'p2', 'p3'])")
254+
assert result.stderr == ""
255+
answers = repl.locals["answers"]
256+
assert answers == ["result-p0", "result-p1", "result-p2", "result-p3"]
257+
repl.cleanup()
258+
259+
def test_batched_parallel_partial_failure(self):
260+
"""Failures in some threads should not affect others."""
261+
call_count = [0]
262+
lock = threading.Lock()
263+
264+
def sometimes_fail(prompt, model):
265+
with lock:
266+
call_count[0] += 1
267+
if prompt == "fail_me":
268+
raise RuntimeError("intentional failure")
269+
time.sleep(0.05)
270+
return _make_completion(f"ok-{prompt}")
271+
272+
repl = LocalREPL(subcall_fn=sometimes_fail, max_concurrent_subcalls=4)
273+
result = repl.execute_code("answers = rlm_query_batched(['a', 'fail_me', 'c'])")
274+
assert result.stderr == ""
275+
answers = repl.locals["answers"]
276+
assert answers[0] == "ok-a"
277+
assert "Error" in answers[1]
278+
assert "intentional failure" in answers[1]
279+
assert answers[2] == "ok-c"
280+
repl.cleanup()
281+
282+
def test_batched_pending_calls_ordered(self):
283+
"""Pending LLM calls should be appended in prompt order for deterministic metadata."""
284+
def delayed_subcall(prompt, model):
285+
# Reverse delay so later prompts finish first
286+
delay = 0.1 if prompt == "first" else 0.01
287+
time.sleep(delay)
288+
return _make_completion(f"resp-{prompt}")
289+
290+
repl = LocalREPL(subcall_fn=delayed_subcall, max_concurrent_subcalls=4)
291+
result = repl.execute_code("answers = rlm_query_batched(['first', 'second'])")
292+
assert result.stderr == ""
293+
assert len(result.rlm_calls) == 2
294+
assert result.rlm_calls[0].response == "resp-first"
295+
assert result.rlm_calls[1].response == "resp-second"
296+
repl.cleanup()
297+
298+
def test_single_prompt_skips_thread_pool(self):
299+
"""Single prompt should not use ThreadPoolExecutor (no overhead)."""
300+
subcall_fn = MagicMock(return_value=_make_completion("solo"))
301+
repl = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=4)
302+
result = repl.execute_code("answers = rlm_query_batched(['only'])")
303+
assert repl.locals["answers"] == ["solo"]
304+
subcall_fn.assert_called_once_with("only", None)
305+
repl.cleanup()
306+
307+
189308
class TestRlmQueryScaffoldRestoration:
190309
"""Test that rlm_query and rlm_query_batched are restored after overwrite."""
191310

0 commit comments

Comments
 (0)