feat: parallelize rlm_query_batched subcalls for faster execution (#139)

* 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>

* move max_concurrent_subcalls as a general REPL property. add more tests

---------

Co-authored-by: dhananjai1729 <dhananjaikrishnakumarpm@Dhananjais-MacBook-Pro.local>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Alex Zhang <alex.lx.zhang@gmail.com>

Author: 4 people

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,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
+            env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls
             environment: BaseEnv = get_environment(self.environment_type, env_kwargs)
 
             if self.persistent:
@@ -765,6 +770,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,

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

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,9 +136,15 @@ 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)
+        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)
@@ -316,9 +323,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 +340,47 @@ 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