From a485c0d55bfe4aeabaf217c9b6c6cb7c4681c123 Mon Sep 17 00:00:00 2001 From: Ashutosh Srivastava Date: Wed, 11 Mar 2026 11:59:10 -0700 Subject: [PATCH 1/6] Make rlm subcalls parallel with thread safety and semaphore --- rlm/core/rlm.py | 44 ++++++- rlm/environments/local_repl.py | 132 ++++++++++++++++++- tests/test_rlm_query.py | 227 +++++++++++++++++++++++++++++++++ 3 files changed, 395 insertions(+), 8 deletions(-) diff --git a/rlm/core/rlm.py b/rlm/core/rlm.py index ce247ecdc..fc1e233dc 100644 --- a/rlm/core/rlm.py +++ b/rlm/core/rlm.py @@ -1,3 +1,4 @@ +import threading import time from collections.abc import Callable from contextlib import contextmanager @@ -70,6 +71,7 @@ def __init__( custom_sub_tools: dict[str, Any] | None = None, compaction: bool = False, compaction_threshold_pct: float = 0.85, + max_concurrent_subcalls: int = 1, 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 +104,10 @@ 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: Max children dispatched in parallel per rlm_query_batched() call. + Default 1 = sequential (identical to previous behavior). Set >1 to enable parallel + sub-agent execution. A global semaphore (default 16 slots) bounds total concurrent + children across all depths to prevent thread/memory explosion. 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 +138,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 @@ -156,6 +163,8 @@ def __init__( self._last_error: str | None = None self._best_partial_answer: str | None = None self._completion_start_time: float | None = None # Set when completion() starts + # Lock for thread-safe cost/error tracking when children run in parallel + self._state_lock = threading.Lock() # Persistence support self.persistent = persistent @@ -231,6 +240,7 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]): # For local environment with max_depth > 1, pass subcall callback for recursive RLM calls if self.environment_type == "local" and self.max_depth > 1: env_kwargs["subcall_fn"] = self._subcall + env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls # Pass custom tools to the environment if self.custom_tools is not None: env_kwargs["custom_tools"] = self.custom_tools @@ -307,6 +317,15 @@ def completion( # Check timeout before each iteration self._check_timeout(i, time_start) + # Fire iteration start callback + if self.on_iteration_start: + try: + self.on_iteration_start(self.depth, i + 1) + except Exception: + pass + + iter_start = time.perf_counter() + # Compaction: check if context needs summarization if self.compaction and hasattr(environment, "append_compaction_entry"): current_tokens, threshold_tokens, max_tokens = self._get_compaction_status( @@ -370,6 +389,13 @@ def completion( # Verbose output for this iteration self.verbose.print_iteration(iteration, i + 1) + # Fire iteration complete callback + if self.on_iteration_complete: + try: + self.on_iteration_complete(self.depth, i + 1, time.perf_counter() - iter_start) + except Exception: + pass + if final_answer is not None: time_end = time.perf_counter() usage = lm_handler.get_usage_summary() @@ -701,9 +727,11 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion: ) # Calculate remaining budget for child (if budget tracking enabled) + # Read _cumulative_cost under lock for thread-safe snapshot remaining_budget = None if self.max_budget is not None: - remaining_budget = self.max_budget - self._cumulative_cost + with self._state_lock: + remaining_budget = self.max_budget - self._cumulative_cost if remaining_budget <= 0: return RLMChatCompletion( root_model=resolved_model, @@ -765,19 +793,25 @@ 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, + on_iteration_start=self.on_iteration_start, + on_iteration_complete=self.on_iteration_complete, ) try: result = child.completion(prompt, root_prompt=None) - # Track child's cost in parent's cumulative cost + # Thread-safe: track child's cost in parent's cumulative cost if result.usage_summary and result.usage_summary.total_cost: - self._cumulative_cost += result.usage_summary.total_cost + with self._state_lock: + self._cumulative_cost += result.usage_summary.total_cost return result except BudgetExceededError as e: - # Propagate child's spending to parent - self._cumulative_cost += e.spent + # Thread-safe: propagate child's spending to parent + with self._state_lock: + self._cumulative_cost += e.spent error_msg = f"Budget exceeded - {e}" return RLMChatCompletion( root_model=resolved_model, diff --git a/rlm/environments/local_repl.py b/rlm/environments/local_repl.py index afdde3b3d..8aa8d7b3f 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 @@ -21,6 +22,38 @@ validate_custom_tools, ) +# ============================================================================= +# Global Subcall Semaphore +# ============================================================================= +# Bounds the total number of concurrent child RLMs across ALL depths in +# the process. Each depth level creates its own ThreadPoolExecutor (avoiding +# pool-within-pool deadlock), but every child must acquire a slot from this +# semaphore before doing real work. When all slots are taken, new children +# wait until a running child finishes. +# +# Default: 16. Override via set_global_max_subcalls() or pass a custom +# semaphore to LocalREPL(subcall_semaphore=...). + +_DEFAULT_MAX_GLOBAL_SUBCALLS = 16 +_global_subcall_semaphore = threading.Semaphore(_DEFAULT_MAX_GLOBAL_SUBCALLS) + + + +def set_global_max_subcalls(n: int) -> None: + """Replace the global subcall semaphore with a new one of size n. + + Call this BEFORE creating any RLM instances. Not safe to call while + subcalls are in flight. + """ + global _global_subcall_semaphore + _global_subcall_semaphore = threading.Semaphore(n) + + +def get_global_subcall_semaphore() -> threading.Semaphore: + """Return the current global subcall semaphore.""" + return _global_subcall_semaphore + + # ============================================================================= # Safe Builtins # ============================================================================= @@ -135,15 +168,29 @@ def __init__( custom_tools: dict[str, Any] | None = None, custom_sub_tools: dict[str, Any] | None = None, compaction: bool = False, + max_concurrent_subcalls: int = 1, + subcall_semaphore: threading.Semaphore | None = None, **kwargs, ): + """ + Args: + max_concurrent_subcalls: Max children to dispatch in parallel from + THIS REPL instance (controls local ThreadPoolExecutor size). + Default 1 = sequential (identical to previous behavior). + subcall_semaphore: Global semaphore bounding total concurrent child + RLMs across ALL depths. If None, uses the module-level default + (16 slots). Pass a custom one for testing or tighter control. + """ 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._subcall_semaphore = subcall_semaphore or _global_subcall_semaphore self.original_cwd = os.getcwd() self.temp_dir = tempfile.mkdtemp(prefix=f"repl_env_{uuid.uuid4()}_") self._lock = threading.Lock() + self._calls_lock = threading.Lock() # Protects _pending_llm_calls during parallel subcalls self._context_count: int = 0 self._history_count: int = 0 self.compaction = compaction @@ -319,6 +366,19 @@ def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> li """Spawn recursive RLM sub-calls for multiple prompts. Each prompt gets its own child RLM for deeper thinking. + + Execution modes (controlled by max_concurrent_subcalls): + - 1 (default): Sequential execution, identical to previous behavior. + Zero overhead — no locks, no thread pool, no semaphore. + - >1: Parallel execution via ThreadPoolExecutor. Each child RLM + runs in its own thread with its own LMHandler and LocalREPL. + A global semaphore bounds total concurrent children across all + depths to prevent thread/memory explosion with deep recursion. + + No deadlock risk: each depth level creates its own short-lived + ThreadPoolExecutor (no pool-within-pool). The global semaphore does + not own threads — workers acquire a slot, do work, release. + Falls back to llm_query_batched if no recursive capability is configured. Args: @@ -329,6 +389,14 @@ 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: + if not prompts: + return [] + + # Parallel execution: multiple prompts + concurrency enabled + if self.max_concurrent_subcalls > 1 and len(prompts) > 1: + return self._rlm_query_batched_parallel(prompts, model) + + # Sequential execution (default): zero overhead, identical to previous behavior results = [] for prompt in prompts: try: @@ -342,6 +410,55 @@ def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> li # Fall back to plain batched LM call if no recursive capability return self._llm_query_batched(prompts, model) + def _rlm_query_batched_parallel(self, prompts: list[str], model: str | None = None) -> list[str]: + """Execute multiple RLM sub-calls in parallel using ThreadPoolExecutor. + + Each child RLM gets its own thread. Results maintain the original + prompt order regardless of which child finishes first. + + Concurrency control (two layers): + 1. Local ThreadPoolExecutor: bounds children dispatched by THIS REPL + (max_concurrent_subcalls). Short-lived — created and destroyed + per batch call. + 2. Global semaphore (_subcall_semaphore): bounds total concurrent + children across ALL depths in the process. A worker acquires a + slot before calling subcall_fn and releases it after, so the + semaphore never blocks the pool itself — only the actual work. + + Thread safety: + - Each child RLM is fully isolated (own LMHandler, own LocalREPL). + - _pending_llm_calls is protected by _calls_lock for concurrent appends. + - subcall_fn (RLM._subcall) is called from worker threads — the parent + RLM's _state_lock protects cumulative_cost and error tracking. + """ + results: list[str | None] = [None] * len(prompts) + sem = self._subcall_semaphore + + def _execute_one(idx: int, prompt: str) -> tuple[int, str, RLMChatCompletion | None]: + sem.acquire() + try: + completion = self.subcall_fn(prompt, model) + return (idx, completion.response, completion) + except Exception as e: + return (idx, f"Error: RLM query failed - {e}", None) + finally: + sem.release() + + max_workers = min(self.max_concurrent_subcalls, len(prompts)) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = { + executor.submit(_execute_one, idx, prompt): idx + for idx, prompt in enumerate(prompts) + } + for future in as_completed(futures): + idx, response, completion = future.result() + results[idx] = response + if completion is not None: + with self._calls_lock: + self._pending_llm_calls.append(completion) + + return results + def load_context(self, context_payload: dict | list | str): """Load context into the environment as context_0 (and 'context' alias).""" self.add_context(context_payload, 0) @@ -449,13 +566,22 @@ def _capture_output(self): @contextmanager def _temp_cwd(self): - """Temporarily change to temp directory for execution.""" - old_cwd = os.getcwd() + """Temporarily change to temp directory for execution. + + Always restores to self.original_cwd (set at init time) rather than + os.getcwd(), which avoids [Errno 2] when parallel threads delete + each other's temp dirs. No lock needed: the worst case is a brief + cwd mismatch during concurrent exec(), which only affects relative + file I/O in model code (rare — context is passed via variables). + """ try: os.chdir(self.temp_dir) yield finally: - os.chdir(old_cwd) + try: + os.chdir(self.original_cwd) + except OSError: + pass def _restore_scaffold(self) -> None: """Restore scaffold names after execution so overwrites (e.g. context = 'x') don't persist.""" diff --git a/tests/test_rlm_query.py b/tests/test_rlm_query.py index 850a87fff..a385c58ee 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 @@ -208,3 +210,228 @@ def test_rlm_query_batched_restored_after_overwrite(self): repl.execute_code("answers = rlm_query_batched(['q1'])") assert repl.locals["answers"] == ["real"] repl.cleanup() + + +# ============================================================================= +# Parallel execution tests (max_concurrent_subcalls > 1) +# ============================================================================= + + +def _make_slow_subcall_fn(delay: float = 0.3): + """Create a thread-safe subcall_fn that sleeps to simulate real LLM latency. + + Returns the function. The function records the prompt in its response + so result ordering can be verified. + """ + def subcall_fn(prompt: str, model: str | None = None) -> RLMChatCompletion: + time.sleep(delay) + return _make_completion(f"answer_for_{prompt}") + + return subcall_fn + + +class TestRlmQueryBatchedParallel: + """Tests for parallel rlm_query_batched when max_concurrent_subcalls > 1.""" + + def test_parallel_faster_than_sequential(self): + """Parallel execution should be measurably faster than an actual sequential run.""" + delay = 0.3 + prompts = ["q0", "q1", "q2", "q3", "q4"] + subcall_fn = _make_slow_subcall_fn(delay) + + # Actual sequential run (max_concurrent=1) + seq_repl = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=1) + seq_start = time.perf_counter() + seq_repl.execute_code(f"answers = rlm_query_batched({prompts})") + seq_elapsed = time.perf_counter() - seq_start + seq_answers = seq_repl.locals["answers"] + seq_repl.cleanup() + + # Actual parallel run (max_concurrent=5) + par_repl = LocalREPL( + subcall_fn=subcall_fn, + max_concurrent_subcalls=5, + subcall_semaphore=threading.Semaphore(16), + ) + par_start = time.perf_counter() + par_repl.execute_code(f"answers = rlm_query_batched({prompts})") + par_elapsed = time.perf_counter() - par_start + par_answers = par_repl.locals["answers"] + par_repl.cleanup() + + # Both should produce the same results + assert len(seq_answers) == len(par_answers) == 5 + assert seq_answers == par_answers + + # Parallel should be at least 2x faster than sequential + assert par_elapsed < seq_elapsed * 0.6, ( + f"Parallel ({par_elapsed:.2f}s) not fast enough vs sequential ({seq_elapsed:.2f}s)" + ) + + def test_parallel_results_maintain_order(self): + """Results should match input prompt order, not completion order.""" + # Give each prompt a different delay so they finish out of order + delays = {"p0": 0.25, "p1": 0.05, "p2": 0.15, "p3": 0.01, "p4": 0.10} + + def subcall_fn(prompt: str, model: str | None = None) -> RLMChatCompletion: + time.sleep(delays.get(prompt, 0.05)) + return _make_completion(f"answer_for_{prompt}") + + repl = LocalREPL( + subcall_fn=subcall_fn, + max_concurrent_subcalls=5, + subcall_semaphore=threading.Semaphore(16), + ) + result = repl.execute_code( + "answers = rlm_query_batched(['p0', 'p1', 'p2', 'p3', 'p4'])" + ) + assert result.stderr == "" + answers = repl.locals["answers"] + for i in range(5): + assert answers[i] == f"answer_for_p{i}", f"answers[{i}] = {answers[i]}" + repl.cleanup() + + def test_parallel_partial_failure(self): + """One failed subcall should not block or corrupt other results.""" + def subcall_fn(prompt: str, model: str | None = None) -> RLMChatCompletion: + time.sleep(0.05) + if prompt == "c": + raise RuntimeError("child exploded") + return _make_completion(f"ok_{prompt}") + + repl = LocalREPL( + subcall_fn=subcall_fn, + max_concurrent_subcalls=4, + subcall_semaphore=threading.Semaphore(16), + ) + result = repl.execute_code( + "answers = rlm_query_batched(['a', 'b', 'c', 'd'])" + ) + assert result.stderr == "" + answers = repl.locals["answers"] + assert len(answers) == 4 + assert answers[0] == "ok_a" + assert answers[1] == "ok_b" + assert "Error" in answers[2] and "child exploded" in answers[2] + assert answers[3] == "ok_d" + repl.cleanup() + + def test_parallel_tracks_pending_calls(self): + """All successful completions should be tracked in rlm_calls.""" + def subcall_fn(prompt: str, model: str | None = None) -> RLMChatCompletion: + time.sleep(0.02) + return _make_completion(f"resp_{prompt}") + + repl = LocalREPL( + subcall_fn=subcall_fn, + max_concurrent_subcalls=4, + subcall_semaphore=threading.Semaphore(16), + ) + result = repl.execute_code( + "rlm_query_batched(['a', 'b', 'c', 'd'])" + ) + assert len(result.rlm_calls) == 4 + repl.cleanup() + + def test_sequential_fallback_when_max_concurrent_1(self): + """max_concurrent_subcalls=1 should execute sequentially (default behavior).""" + completions = [_make_completion(f"seq_{i}") for i in range(3)] + subcall_fn = MagicMock(side_effect=completions) + + repl = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=1) + result = repl.execute_code("answers = rlm_query_batched(['a', 'b', 'c'])") + + assert result.stderr == "" + assert repl.locals["answers"] == ["seq_0", "seq_1", "seq_2"] + assert subcall_fn.call_count == 3 + # Verify calls were in order (sequential) + calls = [c[0][0] for c in subcall_fn.call_args_list] + assert calls == ["a", "b", "c"] + repl.cleanup() + + def test_single_prompt_skips_parallel_path(self): + """A single prompt should use sequential path even with max_concurrent > 1.""" + subcall_fn = MagicMock(return_value=_make_completion("single")) + repl = LocalREPL( + subcall_fn=subcall_fn, + max_concurrent_subcalls=8, + subcall_semaphore=threading.Semaphore(16), + ) + result = repl.execute_code("answers = rlm_query_batched(['only'])") + assert repl.locals["answers"] == ["single"] + subcall_fn.assert_called_once_with("only", None) + repl.cleanup() + + +class TestGlobalSemaphoreBounding: + """Tests that the global semaphore correctly bounds total concurrency.""" + + def test_semaphore_limits_concurrent_children(self): + """With semaphore(2), at most 2 children should run simultaneously.""" + max_seen = 0 + current = 0 + lock = threading.Lock() + + def subcall_fn(prompt: str, model: str | None = None) -> RLMChatCompletion: + nonlocal max_seen, current + with lock: + current += 1 + if current > max_seen: + max_seen = current + time.sleep(0.15) + with lock: + current -= 1 + return _make_completion(f"ok_{prompt}") + + sem = threading.Semaphore(2) + repl = LocalREPL( + subcall_fn=subcall_fn, + max_concurrent_subcalls=5, # pool allows 5 + subcall_semaphore=sem, # but semaphore only allows 2 + ) + result = repl.execute_code( + "answers = rlm_query_batched(['a', 'b', 'c', 'd', 'e'])" + ) + assert result.stderr == "" + assert len(repl.locals["answers"]) == 5 + assert all(a.startswith("ok_") for a in repl.locals["answers"]) + # Semaphore should have limited concurrent execution to 2 + assert max_seen <= 2, f"max concurrent was {max_seen}, expected <= 2" + repl.cleanup() + + def test_semaphore_shared_across_repls(self): + """Two REPLs sharing the same semaphore should respect the combined limit.""" + max_seen = 0 + current = 0 + lock = threading.Lock() + + def subcall_fn(prompt: str, model: str | None = None) -> RLMChatCompletion: + nonlocal max_seen, current + with lock: + current += 1 + if current > max_seen: + max_seen = current + time.sleep(0.15) + with lock: + current -= 1 + return _make_completion("ok") + + shared_sem = threading.Semaphore(3) + repl1 = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=4, subcall_semaphore=shared_sem) + repl2 = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=4, subcall_semaphore=shared_sem) + + # Run both in parallel threads + def run_repl(repl): + repl.execute_code("rlm_query_batched(['x', 'y', 'z'])") + + t1 = threading.Thread(target=run_repl, args=(repl1,)) + t2 = threading.Thread(target=run_repl, args=(repl2,)) + t1.start() + t2.start() + t1.join(timeout=10) + t2.join(timeout=10) + + # 6 total tasks, but semaphore(3) means at most 3 concurrent + assert max_seen <= 3, f"max concurrent was {max_seen}, expected <= 3" + repl1.cleanup() + repl2.cleanup() From bc52f441d48f5ee802507eccdd082fcc15921e10 Mon Sep 17 00:00:00 2001 From: Ashutosh Srivastava Date: Tue, 17 Mar 2026 22:51:07 -0700 Subject: [PATCH 2/6] Fix lint failure --- tests/test_rlm_query.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_rlm_query.py b/tests/test_rlm_query.py index a385c58ee..46bfeb434 100644 --- a/tests/test_rlm_query.py +++ b/tests/test_rlm_query.py @@ -357,7 +357,7 @@ def test_single_prompt_skips_parallel_path(self): max_concurrent_subcalls=8, subcall_semaphore=threading.Semaphore(16), ) - result = repl.execute_code("answers = rlm_query_batched(['only'])") + _ = repl.execute_code("answers = rlm_query_batched(['only'])") assert repl.locals["answers"] == ["single"] subcall_fn.assert_called_once_with("only", None) repl.cleanup() From dbb701fccb3dd448ab4a7a41ae3089ca2481620b Mon Sep 17 00:00:00 2001 From: Ashutosh Srivastava Date: Mon, 6 Apr 2026 11:49:06 -0700 Subject: [PATCH 3/6] Slim parallel rlm_query_batched: remove global semaphore, fix _temp_cwd race, default=1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove global subcall semaphore and related API (set_global_max_subcalls, etc.) per maintainer feedback — concurrency bounded by ThreadPoolExecutor(max_workers) only - Fix _temp_cwd race condition: restore to self.original_cwd instead of os.getcwd() which could point to a deleted temp dir under concurrent execution ([Errno 2]) - Change default max_concurrent_subcalls from 4 to 1 (sequential, backward-compatible) - Scope max_concurrent_subcalls env_kwargs pass to local environment only - Append pending_llm_calls in prompt order for deterministic metadata - Remove duplicate semaphore-dependent tests, fix default assertions (4→1) Co-Authored-By: Claude Opus 4.6 (1M context) Made-with: Cursor --- rlm/core/rlm.py | 47 ++----- rlm/environments/local_repl.py | 114 +++-------------- tests/test_rlm_query.py | 225 --------------------------------- 3 files changed, 27 insertions(+), 359 deletions(-) diff --git a/rlm/core/rlm.py b/rlm/core/rlm.py index 1a4881426..70c7ce056 100644 --- a/rlm/core/rlm.py +++ b/rlm/core/rlm.py @@ -1,4 +1,3 @@ -import threading import time from collections.abc import Callable from contextlib import contextmanager @@ -71,7 +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 = 1, + 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, @@ -104,10 +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: Max children dispatched in parallel per rlm_query_batched() call. - Default 1 = sequential (identical to previous behavior). Set >1 to enable parallel - sub-agent execution. A global semaphore (default 16 slots) bounds total concurrent - children across all depths to prevent thread/memory explosion. + 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). @@ -163,8 +160,6 @@ def __init__( self._last_error: str | None = None self._best_partial_answer: str | None = None self._completion_start_time: float | None = None # Set when completion() starts - # Lock for thread-safe cost/error tracking when children run in parallel - self._state_lock = threading.Lock() # Persistence support self.persistent = persistent @@ -240,7 +235,6 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]): # For local environment with max_depth > 1, pass subcall callback for recursive RLM calls if self.environment_type == "local" and self.max_depth > 1: env_kwargs["subcall_fn"] = self._subcall - env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls # Pass custom tools to the environment if self.custom_tools is not None: env_kwargs["custom_tools"] = self.custom_tools @@ -248,7 +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 - env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls + 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: @@ -318,15 +313,6 @@ def completion( # Check timeout before each iteration self._check_timeout(i, time_start) - # Fire iteration start callback - if self.on_iteration_start: - try: - self.on_iteration_start(self.depth, i + 1) - except Exception: - pass - - iter_start = time.perf_counter() - # Compaction: check if context needs summarization if self.compaction and hasattr(environment, "append_compaction_entry"): current_tokens, threshold_tokens, max_tokens = self._get_compaction_status( @@ -390,13 +376,6 @@ def completion( # Verbose output for this iteration self.verbose.print_iteration(iteration, i + 1) - # Fire iteration complete callback - if self.on_iteration_complete: - try: - self.on_iteration_complete(self.depth, i + 1, time.perf_counter() - iter_start) - except Exception: - pass - if final_answer is not None: time_end = time.perf_counter() usage = lm_handler.get_usage_summary() @@ -728,11 +707,9 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion: ) # Calculate remaining budget for child (if budget tracking enabled) - # Read _cumulative_cost under lock for thread-safe snapshot remaining_budget = None if self.max_budget is not None: - with self._state_lock: - remaining_budget = self.max_budget - self._cumulative_cost + remaining_budget = self.max_budget - self._cumulative_cost if remaining_budget <= 0: return RLMChatCompletion( root_model=resolved_model, @@ -799,20 +776,16 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion: # Propagate callbacks to children for nested tracking on_subcall_start=self.on_subcall_start, on_subcall_complete=self.on_subcall_complete, - on_iteration_start=self.on_iteration_start, - on_iteration_complete=self.on_iteration_complete, ) try: result = child.completion(prompt, root_prompt=None) - # Thread-safe: track child's cost in parent's cumulative cost + # Track child's cost in parent's cumulative cost if result.usage_summary and result.usage_summary.total_cost: - with self._state_lock: - self._cumulative_cost += result.usage_summary.total_cost + self._cumulative_cost += result.usage_summary.total_cost return result except BudgetExceededError as e: - # Thread-safe: propagate child's spending to parent - with self._state_lock: - self._cumulative_cost += e.spent + # Propagate child's spending to parent + self._cumulative_cost += e.spent error_msg = f"Budget exceeded - {e}" return RLMChatCompletion( root_model=resolved_model, diff --git a/rlm/environments/local_repl.py b/rlm/environments/local_repl.py index a711e4d14..6300c94a7 100644 --- a/rlm/environments/local_repl.py +++ b/rlm/environments/local_repl.py @@ -22,38 +22,6 @@ validate_custom_tools, ) -# ============================================================================= -# Global Subcall Semaphore -# ============================================================================= -# Bounds the total number of concurrent child RLMs across ALL depths in -# the process. Each depth level creates its own ThreadPoolExecutor (avoiding -# pool-within-pool deadlock), but every child must acquire a slot from this -# semaphore before doing real work. When all slots are taken, new children -# wait until a running child finishes. -# -# Default: 16. Override via set_global_max_subcalls() or pass a custom -# semaphore to LocalREPL(subcall_semaphore=...). - -_DEFAULT_MAX_GLOBAL_SUBCALLS = 16 -_global_subcall_semaphore = threading.Semaphore(_DEFAULT_MAX_GLOBAL_SUBCALLS) - - - -def set_global_max_subcalls(n: int) -> None: - """Replace the global subcall semaphore with a new one of size n. - - Call this BEFORE creating any RLM instances. Not safe to call while - subcalls are in flight. - """ - global _global_subcall_semaphore - _global_subcall_semaphore = threading.Semaphore(n) - - -def get_global_subcall_semaphore() -> threading.Semaphore: - """Return the current global subcall semaphore.""" - return _global_subcall_semaphore - - # ============================================================================= # Safe Builtins # ============================================================================= @@ -168,21 +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 = 1, - subcall_semaphore: threading.Semaphore | None = None, + max_concurrent_subcalls: int = 4, **kwargs, ): - super().__init__( - persistent=persistent, - depth=depth, - max_concurrent_subcalls=max_concurrent_subcalls, - **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._subcall_semaphore = subcall_semaphore or _global_subcall_semaphore self.original_cwd = os.getcwd() self.temp_dir = tempfile.mkdtemp(prefix=f"repl_env_{uuid.uuid4()}_") self._lock = threading.Lock() @@ -360,19 +321,11 @@ def _rlm_query(self, prompt: str, model: str | None = None) -> str: def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> list[str]: """Spawn recursive RLM sub-calls for multiple prompts in parallel. - Each prompt gets its own child RLM for deeper thinking. - Execution modes (controlled by max_concurrent_subcalls): - - 1 (default): Sequential execution, identical to previous behavior. - Zero overhead — no locks, no thread pool, no semaphore. - - >1: Parallel execution via ThreadPoolExecutor. Each child RLM - runs in its own thread with its own LMHandler and LocalREPL. - A global semaphore bounds total concurrent children across all - depths to prevent thread/memory explosion with deep recursion. - - No deadlock risk: each depth level creates its own short-lived - ThreadPoolExecutor (no pool-within-pool). The global semaphore does - not own threads — workers acquire a slot, do work, release. + Each prompt gets its own child RLM for deeper thinking. When multiple + prompts are provided and max_concurrent_subcalls > 1, subcalls run + concurrently using a thread pool. Results are returned in the same + order as input prompts. Falls back to llm_query_batched if no recursive capability is configured. @@ -396,63 +349,26 @@ def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> li for prompt in prompts: try: completion = self.subcall_fn(prompt, model) - with lock: - completions.append((index, completion)) - results[index] = completion.response + self._pending_llm_calls.append(completion) + results.append(completion.response) except Exception as 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) - + results.append(f"Error: RLM query failed - {e}") return results # Fall back to plain batched LM call if no recursive capability return self._llm_query_batched(prompts, model) def _rlm_query_batched_parallel(self, prompts: list[str], model: str | None = None) -> list[str]: - """Execute multiple RLM sub-calls in parallel using ThreadPoolExecutor. - - Each child RLM gets its own thread. Results maintain the original - prompt order regardless of which child finishes first. - - Concurrency control (two layers): - 1. Local ThreadPoolExecutor: bounds children dispatched by THIS REPL - (max_concurrent_subcalls). Short-lived — created and destroyed - per batch call. - 2. Global semaphore (_subcall_semaphore): bounds total concurrent - children across ALL depths in the process. A worker acquires a - slot before calling subcall_fn and releases it after, so the - semaphore never blocks the pool itself — only the actual work. - - Thread safety: - - Each child RLM is fully isolated (own LMHandler, own LocalREPL). - - _pending_llm_calls is protected by _calls_lock for concurrent appends. - - subcall_fn (RLM._subcall) is called from worker threads — the parent - RLM's _state_lock protects cumulative_cost and error tracking. - """ + """Execute multiple RLM sub-calls in parallel using ThreadPoolExecutor.""" results: list[str | None] = [None] * len(prompts) - sem = self._subcall_semaphore + completions: list[tuple[int, RLMChatCompletion]] = [] def _execute_one(idx: int, prompt: str) -> tuple[int, str, RLMChatCompletion | None]: - sem.acquire() try: completion = self.subcall_fn(prompt, model) return (idx, completion.response, completion) except Exception as e: return (idx, f"Error: RLM query failed - {e}", None) - finally: - sem.release() max_workers = min(self.max_concurrent_subcalls, len(prompts)) with ThreadPoolExecutor(max_workers=max_workers) as executor: @@ -464,8 +380,12 @@ def _execute_one(idx: int, prompt: str) -> tuple[int, str, RLMChatCompletion | N idx, response, completion = future.result() results[idx] = response if completion is not None: - with self._calls_lock: - self._pending_llm_calls.append(completion) + completions.append((idx, completion)) + + # Append in prompt order for deterministic metadata + completions.sort(key=lambda x: x[0]) + for _, completion in completions: + self._pending_llm_calls.append(completion) return results diff --git a/tests/test_rlm_query.py b/tests/test_rlm_query.py index dfd87d9a3..a62cacdd4 100644 --- a/tests/test_rlm_query.py +++ b/tests/test_rlm_query.py @@ -491,228 +491,3 @@ def test_rlm_query_batched_restored_after_overwrite(self): repl.execute_code("answers = rlm_query_batched(['q1'])") assert repl.locals["answers"] == ["real"] repl.cleanup() - - -# ============================================================================= -# Parallel execution tests (max_concurrent_subcalls > 1) -# ============================================================================= - - -def _make_slow_subcall_fn(delay: float = 0.3): - """Create a thread-safe subcall_fn that sleeps to simulate real LLM latency. - - Returns the function. The function records the prompt in its response - so result ordering can be verified. - """ - def subcall_fn(prompt: str, model: str | None = None) -> RLMChatCompletion: - time.sleep(delay) - return _make_completion(f"answer_for_{prompt}") - - return subcall_fn - - -class TestRlmQueryBatchedParallel: - """Tests for parallel rlm_query_batched when max_concurrent_subcalls > 1.""" - - def test_parallel_faster_than_sequential(self): - """Parallel execution should be measurably faster than an actual sequential run.""" - delay = 0.3 - prompts = ["q0", "q1", "q2", "q3", "q4"] - subcall_fn = _make_slow_subcall_fn(delay) - - # Actual sequential run (max_concurrent=1) - seq_repl = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=1) - seq_start = time.perf_counter() - seq_repl.execute_code(f"answers = rlm_query_batched({prompts})") - seq_elapsed = time.perf_counter() - seq_start - seq_answers = seq_repl.locals["answers"] - seq_repl.cleanup() - - # Actual parallel run (max_concurrent=5) - par_repl = LocalREPL( - subcall_fn=subcall_fn, - max_concurrent_subcalls=5, - subcall_semaphore=threading.Semaphore(16), - ) - par_start = time.perf_counter() - par_repl.execute_code(f"answers = rlm_query_batched({prompts})") - par_elapsed = time.perf_counter() - par_start - par_answers = par_repl.locals["answers"] - par_repl.cleanup() - - # Both should produce the same results - assert len(seq_answers) == len(par_answers) == 5 - assert seq_answers == par_answers - - # Parallel should be at least 2x faster than sequential - assert par_elapsed < seq_elapsed * 0.6, ( - f"Parallel ({par_elapsed:.2f}s) not fast enough vs sequential ({seq_elapsed:.2f}s)" - ) - - def test_parallel_results_maintain_order(self): - """Results should match input prompt order, not completion order.""" - # Give each prompt a different delay so they finish out of order - delays = {"p0": 0.25, "p1": 0.05, "p2": 0.15, "p3": 0.01, "p4": 0.10} - - def subcall_fn(prompt: str, model: str | None = None) -> RLMChatCompletion: - time.sleep(delays.get(prompt, 0.05)) - return _make_completion(f"answer_for_{prompt}") - - repl = LocalREPL( - subcall_fn=subcall_fn, - max_concurrent_subcalls=5, - subcall_semaphore=threading.Semaphore(16), - ) - result = repl.execute_code( - "answers = rlm_query_batched(['p0', 'p1', 'p2', 'p3', 'p4'])" - ) - assert result.stderr == "" - answers = repl.locals["answers"] - for i in range(5): - assert answers[i] == f"answer_for_p{i}", f"answers[{i}] = {answers[i]}" - repl.cleanup() - - def test_parallel_partial_failure(self): - """One failed subcall should not block or corrupt other results.""" - def subcall_fn(prompt: str, model: str | None = None) -> RLMChatCompletion: - time.sleep(0.05) - if prompt == "c": - raise RuntimeError("child exploded") - return _make_completion(f"ok_{prompt}") - - repl = LocalREPL( - subcall_fn=subcall_fn, - max_concurrent_subcalls=4, - subcall_semaphore=threading.Semaphore(16), - ) - result = repl.execute_code( - "answers = rlm_query_batched(['a', 'b', 'c', 'd'])" - ) - assert result.stderr == "" - answers = repl.locals["answers"] - assert len(answers) == 4 - assert answers[0] == "ok_a" - assert answers[1] == "ok_b" - assert "Error" in answers[2] and "child exploded" in answers[2] - assert answers[3] == "ok_d" - repl.cleanup() - - def test_parallel_tracks_pending_calls(self): - """All successful completions should be tracked in rlm_calls.""" - def subcall_fn(prompt: str, model: str | None = None) -> RLMChatCompletion: - time.sleep(0.02) - return _make_completion(f"resp_{prompt}") - - repl = LocalREPL( - subcall_fn=subcall_fn, - max_concurrent_subcalls=4, - subcall_semaphore=threading.Semaphore(16), - ) - result = repl.execute_code( - "rlm_query_batched(['a', 'b', 'c', 'd'])" - ) - assert len(result.rlm_calls) == 4 - repl.cleanup() - - def test_sequential_fallback_when_max_concurrent_1(self): - """max_concurrent_subcalls=1 should execute sequentially (default behavior).""" - completions = [_make_completion(f"seq_{i}") for i in range(3)] - subcall_fn = MagicMock(side_effect=completions) - - repl = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=1) - result = repl.execute_code("answers = rlm_query_batched(['a', 'b', 'c'])") - - assert result.stderr == "" - assert repl.locals["answers"] == ["seq_0", "seq_1", "seq_2"] - assert subcall_fn.call_count == 3 - # Verify calls were in order (sequential) - calls = [c[0][0] for c in subcall_fn.call_args_list] - assert calls == ["a", "b", "c"] - repl.cleanup() - - def test_single_prompt_skips_parallel_path(self): - """A single prompt should use sequential path even with max_concurrent > 1.""" - subcall_fn = MagicMock(return_value=_make_completion("single")) - repl = LocalREPL( - subcall_fn=subcall_fn, - max_concurrent_subcalls=8, - subcall_semaphore=threading.Semaphore(16), - ) - _ = repl.execute_code("answers = rlm_query_batched(['only'])") - assert repl.locals["answers"] == ["single"] - subcall_fn.assert_called_once_with("only", None) - repl.cleanup() - - -class TestGlobalSemaphoreBounding: - """Tests that the global semaphore correctly bounds total concurrency.""" - - def test_semaphore_limits_concurrent_children(self): - """With semaphore(2), at most 2 children should run simultaneously.""" - max_seen = 0 - current = 0 - lock = threading.Lock() - - def subcall_fn(prompt: str, model: str | None = None) -> RLMChatCompletion: - nonlocal max_seen, current - with lock: - current += 1 - if current > max_seen: - max_seen = current - time.sleep(0.15) - with lock: - current -= 1 - return _make_completion(f"ok_{prompt}") - - sem = threading.Semaphore(2) - repl = LocalREPL( - subcall_fn=subcall_fn, - max_concurrent_subcalls=5, # pool allows 5 - subcall_semaphore=sem, # but semaphore only allows 2 - ) - result = repl.execute_code( - "answers = rlm_query_batched(['a', 'b', 'c', 'd', 'e'])" - ) - assert result.stderr == "" - assert len(repl.locals["answers"]) == 5 - assert all(a.startswith("ok_") for a in repl.locals["answers"]) - # Semaphore should have limited concurrent execution to 2 - assert max_seen <= 2, f"max concurrent was {max_seen}, expected <= 2" - repl.cleanup() - - def test_semaphore_shared_across_repls(self): - """Two REPLs sharing the same semaphore should respect the combined limit.""" - max_seen = 0 - current = 0 - lock = threading.Lock() - - def subcall_fn(prompt: str, model: str | None = None) -> RLMChatCompletion: - nonlocal max_seen, current - with lock: - current += 1 - if current > max_seen: - max_seen = current - time.sleep(0.15) - with lock: - current -= 1 - return _make_completion("ok") - - shared_sem = threading.Semaphore(3) - repl1 = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=4, subcall_semaphore=shared_sem) - repl2 = LocalREPL(subcall_fn=subcall_fn, max_concurrent_subcalls=4, subcall_semaphore=shared_sem) - - # Run both in parallel threads - def run_repl(repl): - repl.execute_code("rlm_query_batched(['x', 'y', 'z'])") - - t1 = threading.Thread(target=run_repl, args=(repl1,)) - t2 = threading.Thread(target=run_repl, args=(repl2,)) - t1.start() - t2.start() - t1.join(timeout=10) - t2.join(timeout=10) - - # 6 total tasks, but semaphore(3) means at most 3 concurrent - assert max_seen <= 3, f"max concurrent was {max_seen}, expected <= 3" - repl1.cleanup() - repl2.cleanup() From 7dc89bbd1e3f1126d6227d52a77a50e2e6aa45df Mon Sep 17 00:00:00 2001 From: Ashutosh Srivastava Date: Mon, 20 Apr 2026 22:45:40 -0700 Subject: [PATCH 4/6] Fix lint --- rlm/environments/local_repl.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/rlm/environments/local_repl.py b/rlm/environments/local_repl.py index 6300c94a7..ad1d2e58c 100644 --- a/rlm/environments/local_repl.py +++ b/rlm/environments/local_repl.py @@ -358,7 +358,9 @@ def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> li # Fall back to plain batched LM call if no recursive capability return self._llm_query_batched(prompts, model) - def _rlm_query_batched_parallel(self, prompts: list[str], model: str | None = None) -> list[str]: + def _rlm_query_batched_parallel( + self, prompts: list[str], model: str | None = None + ) -> list[str]: """Execute multiple RLM sub-calls in parallel using ThreadPoolExecutor.""" results: list[str | None] = [None] * len(prompts) completions: list[tuple[int, RLMChatCompletion]] = [] From 7a58640cfad2230edb110c09a6f437da594210a4 Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Sat, 25 Apr 2026 23:24:16 -0400 Subject: [PATCH 5/6] remove thread based batch RLM, move to subprocesses --- pyproject.toml | 1 + rlm/core/rlm.py | 132 ++++++++++++++++ rlm/environments/local_repl.py | 275 ++++++++++++++++++++++++++++----- tests/test_rlm_query.py | 230 +-------------------------- uv.lock | 4 +- 5 files changed, 373 insertions(+), 269 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 489080359..caf2ab433 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ classifiers = [ dependencies = [ "anthropic>=0.75.0", + "dill>=0.3.7", "google-genai>=1.56.0", "openai>=2.14.0", "portkey-ai>=2.1.0", diff --git a/rlm/core/rlm.py b/rlm/core/rlm.py index 70c7ce056..7aefdfc65 100644 --- a/rlm/core/rlm.py +++ b/rlm/core/rlm.py @@ -235,6 +235,9 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]): # For local environment with max_depth > 1, pass subcall callback for recursive RLM calls if self.environment_type == "local" and self.max_depth > 1: env_kwargs["subcall_fn"] = self._subcall + # Preamble/finalize used by the subprocess-based parallel batched path + env_kwargs["subcall_preamble_fn"] = self._subcall_preamble + env_kwargs["subcall_finalize_fn"] = self._subcall_finalize # Pass custom tools to the environment if self.custom_tools is not None: env_kwargs["custom_tools"] = self.custom_tools @@ -648,6 +651,135 @@ def _fallback_answer(self, message: str | dict[str, Any]) -> str: response = client.completion(message) return response + def _fire_subcall_start(self, next_depth: int, resolved_model: str, prompt: str) -> None: + if self.on_subcall_start: + try: + prompt_preview = prompt[:80] if len(prompt) > 80 else prompt + self.on_subcall_start(next_depth, str(resolved_model), prompt_preview) + except Exception: + pass + + def _subcall_preamble( + self, prompt: str, model: str | None = None + ) -> tuple[str, Any, float, int, str]: + """ + Plan a sub-call without running it; used by subprocess-based batched sub-calls. + + Returns one of: + ("short_circuit", completion, start_time, next_depth, resolved_model) + ("plain_lm", (backend, backend_kwargs, root_model), start_time, next_depth, resolved_model) + ("rlm", config_dict, start_time, next_depth, resolved_model) + + Fires on_subcall_start for non-short-circuit cases. + """ + next_depth = self.depth + 1 + + if model is not None: + child_backend_kwargs = (self.backend_kwargs or {}).copy() + child_backend_kwargs["model_name"] = model + else: + child_backend_kwargs = self.backend_kwargs + resolved_model = model or (child_backend_kwargs or {}).get("model_name", "unknown") + + start_time = time.perf_counter() + + # Max-depth boundary: plain LM call at the leaf + if next_depth >= self.max_depth: + if self.other_backends and self.other_backend_kwargs: + backend = self.other_backends[0] + backend_kwargs = self.other_backend_kwargs[0] + else: + backend = self.backend + backend_kwargs = child_backend_kwargs or {} + root_model = model or backend_kwargs.get("model_name", "unknown") + self._fire_subcall_start(next_depth, root_model, prompt) + return ( + "plain_lm", + (backend, backend_kwargs, root_model), + start_time, + next_depth, + root_model, + ) + + # Budget short-circuit + remaining_budget = None + if self.max_budget is not None: + remaining_budget = self.max_budget - self._cumulative_cost + if remaining_budget <= 0: + completion = RLMChatCompletion( + root_model=resolved_model, + prompt=prompt, + response=( + "Error: Budget exhausted " + f"(spent ${self._cumulative_cost:.6f} of ${self.max_budget:.6f})" + ), + usage_summary=UsageSummary(model_usage_summaries={}), + execution_time=0.0, + ) + return ("short_circuit", completion, start_time, next_depth, resolved_model) + + # Timeout short-circuit + remaining_timeout = None + if self.max_timeout is not None and self._completion_start_time is not None: + elapsed = time.perf_counter() - self._completion_start_time + remaining_timeout = self.max_timeout - elapsed + if remaining_timeout <= 0: + completion = RLMChatCompletion( + root_model=resolved_model, + prompt=prompt, + response=f"Error: Timeout exhausted ({elapsed:.1f}s of {self.max_timeout:.1f}s)", + usage_summary=UsageSummary(model_usage_summaries={}), + execution_time=0.0, + ) + return ("short_circuit", completion, start_time, next_depth, resolved_model) + + self._fire_subcall_start(next_depth, resolved_model, prompt) + + # Serializable config for the child RLM. subcall_fn is NOT included; the subprocess's + # child will wire its own in completion(). + config = { + "backend": self.backend, + "backend_kwargs": child_backend_kwargs, + "environment": self.environment_type, + "environment_kwargs": self.environment_kwargs, + "depth": next_depth, + "max_depth": self.max_depth, + "max_iterations": self.max_iterations, + "max_budget": remaining_budget, + "max_timeout": remaining_timeout, + "max_tokens": self.max_tokens, + "max_errors": self.max_errors, + "custom_system_prompt": self.system_prompt, + "other_backends": self.other_backends, + "other_backend_kwargs": self.other_backend_kwargs, + "logger": None, + "verbose": False, + "custom_tools": self.custom_sub_tools, + "custom_sub_tools": self.custom_sub_tools, + "max_concurrent_subcalls": self.max_concurrent_subcalls, + } + return ("rlm", config, start_time, next_depth, resolved_model) + + def _subcall_finalize( + self, + completion: RLMChatCompletion, + start_time: float, + next_depth: int, + resolved_model: str, + error_msg: str | None, + cost_already_tracked: bool = False, + ) -> None: + """Track cumulative cost and fire on_subcall_complete after a sub-call resolves.""" + if not cost_already_tracked and error_msg is None: + if completion.usage_summary and completion.usage_summary.total_cost: + self._cumulative_cost += completion.usage_summary.total_cost + if self.on_subcall_complete: + try: + duration = time.perf_counter() - start_time + self.on_subcall_complete(next_depth, str(resolved_model), duration, error_msg) + except Exception: + pass + def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion: """ Handle a subcall from the environment, potentially spawning a child RLM. diff --git a/rlm/environments/local_repl.py b/rlm/environments/local_repl.py index ad1d2e58c..3b28ddbee 100644 --- a/rlm/environments/local_repl.py +++ b/rlm/environments/local_repl.py @@ -1,6 +1,7 @@ import copy import io import json +import multiprocessing as mp import os import shutil import sys @@ -9,12 +10,14 @@ import time import uuid from collections.abc import Callable -from concurrent.futures import ThreadPoolExecutor, as_completed +from concurrent.futures import ProcessPoolExecutor, as_completed from contextlib import contextmanager from typing import Any +import dill + from rlm.core.comms_utils import LMRequest, send_lm_request, send_lm_request_batched -from rlm.core.types import REPLResult, RLMChatCompletion +from rlm.core.types import REPLResult, RLMChatCompletion, UsageSummary from rlm.environments.base_env import ( RESERVED_TOOL_NAMES, NonIsolatedEnv, @@ -22,6 +25,115 @@ validate_custom_tools, ) +# ============================================================================= +# Global Subprocess Budget +# ============================================================================= +# Bounds the number of concurrent subprocess sub-calls actively doing work +# across the entire Python process, at every depth. A worker holding a slot +# is one that's currently inside child.completion(). When a worker itself +# needs to spawn grandchildren (rlm_query_batched inside its REPL), it +# releases its slot before awaiting them and reacquires after — this keeps +# the cap enforceable without deadlocking on recursion. + +_DEFAULT_MAX_SUBPROCESSES = 16 +_global_subprocess_semaphore: Any = None + +# Per-process flag: set True inside a subprocess worker so _rlm_query_batched_parallel +# knows to release/reacquire the slot when orchestrating grandchildren. +_IN_SUBPROCESS_WORKER: bool = False + + +def set_max_subprocesses(n: int) -> None: + """Set the global cap on concurrent subprocess-based RLM sub-calls. + + Must be called before any subprocess-based sub-call runs. Not safe to + change mid-run. + """ + global _global_subprocess_semaphore + _global_subprocess_semaphore = mp.get_context("spawn").BoundedSemaphore(n) + + +def _get_subprocess_semaphore(): + global _global_subprocess_semaphore + if _global_subprocess_semaphore is None: + _global_subprocess_semaphore = mp.get_context("spawn").BoundedSemaphore( + _DEFAULT_MAX_SUBPROCESSES + ) + return _global_subprocess_semaphore + + +def _init_subprocess_worker(semaphore) -> None: + """ProcessPoolExecutor initializer: install the global semaphore in this worker. + + BoundedSemaphores can't be passed through submit() arguments under the + spawn start method — only through initargs (or fork inheritance). We set + the module-level _global_subprocess_semaphore here so workers and any + grandchild pools they spawn can share the same kernel semaphore. + """ + global _global_subprocess_semaphore + _global_subprocess_semaphore = semaphore + + +def _subprocess_worker_entry(job_bytes: bytes) -> bytes: + """Top-level entry for subprocess-based sub-calls. + + Unpacks a dill-pickled (kind, data, prompt) tuple and either constructs a + child RLM and runs its completion, or makes a plain LM call at max-depth. + Acquires the global semaphore (set by _init_subprocess_worker) while doing + work and sets _IN_SUBPROCESS_WORKER so recursive rlm_query_batched inside + this worker can release/reacquire the slot. + """ + global _IN_SUBPROCESS_WORKER + semaphore = _global_subprocess_semaphore + kind, data, prompt = dill.loads(job_bytes) + + semaphore.acquire() + _IN_SUBPROCESS_WORKER = True + try: + if kind == "rlm": + from rlm.core.rlm import RLM + + child = RLM(**data) + try: + result = child.completion(prompt, root_prompt=None) + finally: + child.close() + return dill.dumps(result) + + if kind == "plain_lm": + from rlm.clients import get_client + + backend, backend_kwargs, root_model = data + client = get_client(backend, backend_kwargs) + start = time.perf_counter() + try: + response = client.completion(prompt) + end = time.perf_counter() + model_usage = client.get_last_usage() + result = RLMChatCompletion( + root_model=root_model, + prompt=prompt, + response=response, + usage_summary=UsageSummary(model_usage_summaries={root_model: model_usage}), + execution_time=end - start, + ) + except Exception as e: + end = time.perf_counter() + result = RLMChatCompletion( + root_model=root_model, + prompt=prompt, + response=f"Error: LM query failed at max depth - {e}", + usage_summary=UsageSummary(model_usage_summaries={}), + execution_time=end - start, + ) + return dill.dumps(result) + + raise ValueError(f"Unknown subprocess job kind: {kind!r}") + finally: + _IN_SUBPROCESS_WORKER = False + semaphore.release() + + # ============================================================================= # Safe Builtins # ============================================================================= @@ -133,21 +245,28 @@ def __init__( persistent: bool = False, depth: int = 1, subcall_fn: Callable[[str, str | None], RLMChatCompletion] | None = None, + subcall_preamble_fn: Callable[..., tuple] | None = None, + subcall_finalize_fn: Callable[..., None] | None = None, 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) - self.max_concurrent_subcalls = max_concurrent_subcalls + self.subcall_preamble_fn = subcall_preamble_fn + self.subcall_finalize_fn = subcall_finalize_fn self.original_cwd = os.getcwd() self.temp_dir = tempfile.mkdtemp(prefix=f"repl_env_{uuid.uuid4()}_") self._lock = threading.Lock() - self._calls_lock = threading.Lock() # Protects _pending_llm_calls during parallel subcalls self._context_count: int = 0 self._history_count: int = 0 self.compaction = compaction @@ -361,35 +480,115 @@ def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> li def _rlm_query_batched_parallel( self, prompts: list[str], model: str | None = None ) -> list[str]: - """Execute multiple RLM sub-calls in parallel using ThreadPoolExecutor.""" - results: list[str | None] = [None] * len(prompts) - completions: list[tuple[int, RLMChatCompletion]] = [] + """ + Spawn recursive RLM sub-calls for multiple prompts in parallel subprocesses. - def _execute_one(idx: int, prompt: str) -> tuple[int, str, RLMChatCompletion | None]: - try: - completion = self.subcall_fn(prompt, model) - return (idx, completion.response, completion) - except Exception as e: - return (idx, f"Error: RLM query failed - {e}", None) - - max_workers = min(self.max_concurrent_subcalls, len(prompts)) - with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = { - executor.submit(_execute_one, idx, prompt): idx - for idx, prompt in enumerate(prompts) - } - for future in as_completed(futures): - idx, response, completion = future.result() - results[idx] = response - if completion is not None: - completions.append((idx, completion)) - - # Append in prompt order for deterministic metadata - completions.sort(key=lambda x: x[0]) - for _, completion in completions: - self._pending_llm_calls.append(completion) - - return results + Each sub-call runs in its own subprocess so it has its own cwd (via + _temp_cwd) and file-system state without racing with siblings. Global + concurrency is bounded by a module-level BoundedSemaphore; see + set_max_subprocesses. + """ + if self.subcall_preamble_fn is None or self.subcall_finalize_fn is None: + # Missing wiring — fall back to sequential sub-calls via subcall_fn. + # The subprocess path requires preamble/finalize callbacks that RLM + # installs at completion-time. + 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 + + semaphore = _get_subprocess_semaphore() + global _IN_SUBPROCESS_WORKER + is_reentrant = _IN_SUBPROCESS_WORKER + + # If we're inside a subprocess worker, release our slot while orchestrating + # grandchildren to prevent deadlock; reclaim on the way out. + if is_reentrant: + semaphore.release() + + try: + # Plan each sub-call in the parent process (synchronous, cheap) + preambles = [self.subcall_preamble_fn(prompt, model) for prompt in prompts] + # preamble = (kind, data, start_time, next_depth, resolved_model) + + # Collect jobs that need a subprocess (rlm or plain_lm) + pool_jobs: list[tuple[int, tuple]] = [] + for idx, preamble in enumerate(preambles): + if preamble[0] != "short_circuit": + pool_jobs.append((idx, preamble)) + + results_by_idx: dict[int, tuple[RLMChatCompletion, str | None]] = {} + + if pool_jobs: + max_workers = min(self.max_concurrent_subcalls, len(pool_jobs)) + ctx = mp.get_context("spawn") + with ProcessPoolExecutor( + max_workers=max_workers, + mp_context=ctx, + initializer=_init_subprocess_worker, + initargs=(semaphore,), + ) as executor: + fut_to_idx: dict[Any, int] = {} + for idx, preamble in pool_jobs: + kind, data, _, _, resolved_model = preamble + prompt = prompts[idx] + job_bytes = dill.dumps((kind, data, prompt)) + fut = executor.submit(_subprocess_worker_entry, job_bytes) + fut_to_idx[fut] = idx + + for fut in as_completed(fut_to_idx): + idx = fut_to_idx[fut] + preamble = preambles[idx] + resolved_model = preamble[4] + try: + completion = dill.loads(fut.result()) + error_msg = None + except Exception as e: + error_msg = f"Subprocess sub-call failed - {e}" + completion = RLMChatCompletion( + root_model=resolved_model, + prompt=prompts[idx], + response=f"Error: {error_msg}", + usage_summary=UsageSummary(model_usage_summaries={}), + execution_time=0.0, + ) + results_by_idx[idx] = (completion, error_msg) + + # Build response list and fire finalize callbacks in prompt order + responses: list[str | None] = [None] * len(prompts) + ordered_completions: list[RLMChatCompletion] = [] + for idx, preamble in enumerate(preambles): + kind, data, start_time, next_depth, resolved_model = preamble + if kind == "short_circuit": + completion = data # completion stored in slot 1 for short-circuit + error_msg = None + cost_already_tracked = True # short-circuit completions carry no cost + else: + completion, error_msg = results_by_idx[idx] + cost_already_tracked = False + self.subcall_finalize_fn( + completion, + start_time, + next_depth, + resolved_model, + error_msg, + cost_already_tracked, + ) + responses[idx] = completion.response + ordered_completions.append(completion) + + for completion in ordered_completions: + self._pending_llm_calls.append(completion) + + return responses # type: ignore[return-value] + finally: + if is_reentrant: + semaphore.acquire() def load_context(self, context_payload: dict | list | str): """Load context into the environment as context_0 (and 'context' alias).""" @@ -498,13 +697,11 @@ def _capture_output(self): @contextmanager def _temp_cwd(self): - """Temporarily change to temp directory for execution. + """chdir into temp_dir for the duration of exec(). - Always restores to self.original_cwd (set at init time) rather than - os.getcwd(), which avoids [Errno 2] when parallel threads delete - each other's temp dirs. No lock needed: the worst case is a brief - cwd mismatch during concurrent exec(), which only affects relative - file I/O in model code (rare — context is passed via variables). + Safe because this process is single-threaded for its own execute_code: + parallel RLM sub-calls run in subprocesses, not threads, so no other + thread can race on os.chdir. """ try: os.chdir(self.temp_dir) diff --git a/tests/test_rlm_query.py b/tests/test_rlm_query.py index a62cacdd4..ffcfe2fa0 100644 --- a/tests/test_rlm_query.py +++ b/tests/test_rlm_query.py @@ -1,7 +1,5 @@ """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 @@ -188,126 +186,8 @@ 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) - 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.""" + """Tests for max_concurrent_subcalls config and edge cases.""" def test_max_concurrent_subcalls_default_value(self): """Default max_concurrent_subcalls should be 4.""" @@ -321,32 +201,6 @@ def test_max_concurrent_subcalls_custom_value(self): 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")) @@ -357,71 +211,6 @@ def test_max_concurrent_subcalls_larger_than_prompts(self): 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() @@ -431,23 +220,6 @@ def test_batched_empty_prompts_with_parallel(self): 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.""" diff --git a/uv.lock b/uv.lock index 6abd703d0..c96204fd5 100644 --- a/uv.lock +++ b/uv.lock @@ -1760,10 +1760,11 @@ wheels = [ [[package]] name = "rlms" -version = "0.1.0" +version = "0.1.1" source = { editable = "." } dependencies = [ { name = "anthropic" }, + { name = "dill" }, { name = "google-genai" }, { name = "openai" }, { name = "portkey-ai" }, @@ -1807,6 +1808,7 @@ test = [ requires-dist = [ { name = "anthropic", specifier = ">=0.75.0" }, { name = "daytona", marker = "extra == 'daytona'", specifier = ">=0.128.1" }, + { name = "dill", specifier = ">=0.3.7" }, { name = "dill", marker = "extra == 'daytona'", specifier = ">=0.3.7" }, { name = "dill", marker = "extra == 'e2b'", specifier = ">=0.3.7" }, { name = "dill", marker = "extra == 'modal'", specifier = ">=0.3.7" }, From 9a2ee81b192298a555aa41409972e6060a26fc30 Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Sat, 25 Apr 2026 23:24:30 -0400 Subject: [PATCH 6/6] add extra tests --- tests/test_subprocess_subcalls.py | 218 ++++++++++++++++++++++++++++++ 1 file changed, 218 insertions(+) create mode 100644 tests/test_subprocess_subcalls.py diff --git a/tests/test_subprocess_subcalls.py b/tests/test_subprocess_subcalls.py new file mode 100644 index 000000000..946c46781 --- /dev/null +++ b/tests/test_subprocess_subcalls.py @@ -0,0 +1,218 @@ +"""Tests for subprocess-backed parallel rlm_query_batched. + +These tests target the building blocks (global semaphore, subprocess worker) +without standing up a full RLM + LLM backend, since spawning subprocesses +that go through the real client machinery is expensive and fragile. +""" + +import multiprocessing as mp +import os +import time +from concurrent.futures import ProcessPoolExecutor, as_completed + +import dill +import pytest + +from rlm.environments import local_repl as lr + +# --------------------------------------------------------------------------- +# Helpers run in subprocesses — must be top-level for spawn to pickle them. +# --------------------------------------------------------------------------- + + +_TEST_SEMAPHORE = None + + +def _install_test_semaphore(sem) -> None: + """ProcessPoolExecutor initializer — semaphores can't cross spawn via submit().""" + global _TEST_SEMAPHORE + _TEST_SEMAPHORE = sem + + +def _slot_holder(hold_seconds: float) -> tuple[float, float]: + """Acquire the test semaphore, sleep, release. Returns (start_ts, end_ts).""" + _TEST_SEMAPHORE.acquire() + try: + start = time.monotonic() + time.sleep(hold_seconds) + return (start, time.monotonic()) + finally: + _TEST_SEMAPHORE.release() + + +def _isolated_writer(out_dir: str, payload: str) -> tuple[str, str]: + """chdir into out_dir, write a file by relative path via three different + APIs (open, os.open, pathlib), and return (cwd_seen, file_we_wrote).""" + import pathlib + + os.chdir(out_dir) + cwd_seen = os.getcwd() + + # Plain open with a relative path + with open("via_open.txt", "w") as f: + f.write(payload) + + # os.open with a relative path + fd = os.open("via_os_open.txt", os.O_CREAT | os.O_WRONLY) + try: + os.write(fd, payload.encode()) + finally: + os.close(fd) + + # pathlib relative path + pathlib.Path("via_pathlib.txt").write_text(payload) + + return (cwd_seen, payload) + + +# --------------------------------------------------------------------------- +# Global semaphore enforcement +# --------------------------------------------------------------------------- + + +class TestGlobalSubprocessSemaphore: + def test_semaphore_caps_concurrent_workers(self): + """With N=2 slots and 4 workers each holding for 0.5s, total wall + time must be at least ~1s (two waves of two).""" + ctx = mp.get_context("spawn") + sem = ctx.BoundedSemaphore(2) + + with ProcessPoolExecutor( + max_workers=4, + mp_context=ctx, + initializer=_install_test_semaphore, + initargs=(sem,), + ) as pool: + t0 = time.monotonic() + futures = [pool.submit(_slot_holder, 0.5) for _ in range(4)] + intervals = [f.result() for f in as_completed(futures)] + wall = time.monotonic() - t0 + + # Two waves of two — ≥1s total. Use generous bound to stay non-flaky. + assert wall >= 0.9, f"semaphore did not cap concurrency: wall={wall:.2f}s" + + # Sweep over (start, +1) and (end, -1) events to confirm the live + # count of holders never exceeds 2. + events = [] + for s, e in intervals: + events.append((s, 1)) + events.append((e, -1)) + events.sort() + live = 0 + peak = 0 + for _, delta in events: + live += delta + peak = max(peak, live) + assert peak <= 2, f"more than 2 workers held the semaphore concurrently: peak={peak}" + + def test_set_max_subprocesses_replaces_global_semaphore(self): + """set_max_subprocesses(n) installs a fresh semaphore of size n.""" + original = lr._global_subprocess_semaphore + try: + lr.set_max_subprocesses(3) + sem = lr._get_subprocess_semaphore() + # BoundedSemaphore exposes the value via _value (CPython detail); + # just confirm three acquires don't block and a fourth would. + assert sem.acquire(timeout=0.1) + assert sem.acquire(timeout=0.1) + assert sem.acquire(timeout=0.1) + assert not sem.acquire(timeout=0.1) + # Cleanup: release what we acquired. + sem.release() + sem.release() + sem.release() + finally: + lr._global_subprocess_semaphore = original + + def test_get_subprocess_semaphore_lazily_creates(self): + """First call to _get_subprocess_semaphore creates with default size.""" + original = lr._global_subprocess_semaphore + try: + lr._global_subprocess_semaphore = None + sem = lr._get_subprocess_semaphore() + assert sem is not None + assert lr._global_subprocess_semaphore is sem + finally: + lr._global_subprocess_semaphore = original + + +# --------------------------------------------------------------------------- +# Cwd isolation across parallel subprocess workers +# --------------------------------------------------------------------------- + + +class TestSubprocessCwdIsolation: + def test_parallel_workers_each_have_their_own_cwd(self, tmp_path): + """Four subprocesses each chdir into their own dir and write files + via open/os.open/pathlib. None of the writes should cross-contaminate. + """ + dirs = [] + for i in range(4): + d = tmp_path / f"worker_{i}" + d.mkdir() + dirs.append(str(d)) + + ctx = mp.get_context("spawn") + with ProcessPoolExecutor(max_workers=4, mp_context=ctx) as pool: + futures = {pool.submit(_isolated_writer, dirs[i], f"payload-{i}"): i for i in range(4)} + for fut in as_completed(futures): + i = futures[fut] + cwd, payload = fut.result() + assert cwd == dirs[i] + assert payload == f"payload-{i}" + + # Verify each file landed in the correct dir + for i in range(4): + for fname in ("via_open.txt", "via_os_open.txt", "via_pathlib.txt"): + p = os.path.join(dirs[i], fname) + assert os.path.exists(p), f"missing {p}" + with open(p) as f: + assert f.read() == f"payload-{i}", f"wrong contents in {p}" + + # Spot-check: no cross-leak. dir 0 should not have dir 1's content. + with open(os.path.join(dirs[0], "via_open.txt")) as f: + assert f.read() == "payload-0" + + def test_parent_cwd_unchanged(self, tmp_path): + """Subprocesses' chdir doesn't affect the parent's cwd.""" + before = os.getcwd() + d = tmp_path / "iso" + d.mkdir() + ctx = mp.get_context("spawn") + with ProcessPoolExecutor(max_workers=1, mp_context=ctx) as pool: + pool.submit(_isolated_writer, str(d), "x").result() + assert os.getcwd() == before + + +# --------------------------------------------------------------------------- +# _subprocess_worker_entry — exercise the dill envelope and semaphore +# --------------------------------------------------------------------------- + + +def _no_op_dill_job(semaphore) -> bytes: + """Pretend to be a worker_entry that just returns a marker.""" + semaphore.acquire() + try: + return dill.dumps("ok") + finally: + semaphore.release() + + +class TestDillEnvelope: + def test_dill_serializes_lambda_in_job_payload(self): + """A lambda inside a custom_tools-like dict survives dill round-trip + but would fail under stdlib pickle. Proves dill actually buys us + something for the subprocess path.""" + import pickle + + payload = { + "kind": "rlm", + "custom_sub_tools": {"square": lambda x: x * x}, + } + # stdlib pickle should fail on the lambda + with pytest.raises((pickle.PicklingError, AttributeError)): + pickle.dumps(payload) + + # dill should handle it + roundtrip = dill.loads(dill.dumps(payload)) + assert roundtrip["custom_sub_tools"]["square"](5) == 25