Skip to content

Commit 6b60ec9

Browse files
Ashutosh Srivastavaclaude
andcommitted
Slim parallel rlm_query_batched: remove global semaphore, fix _temp_cwd race, default=1
- 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) <noreply@anthropic.com> Made-with: Cursor
1 parent 460595c commit 6b60ec9

4 files changed

Lines changed: 30 additions & 362 deletions

File tree

‎rlm/core/rlm.py‎

Lines changed: 9 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import threading
21
import time
32
from collections.abc import Callable
43
from contextlib import contextmanager
@@ -104,10 +103,8 @@ def __init__(
104103
when root context reaches compaction_threshold_pct of the model's context limit.
105104
compaction_threshold_pct: When compaction is on, trigger summarization when root
106105
message token count reaches this fraction of the model context limit (default 0.85).
107-
max_concurrent_subcalls: Max children dispatched in parallel per rlm_query_batched() call.
108-
Default 1 = sequential (identical to previous behavior). Set >1 to enable parallel
109-
sub-agent execution. A global semaphore (default 16 slots) bounds total concurrent
110-
children across all depths to prevent thread/memory explosion.
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.
111108
on_subcall_start: Callback fired when a child RLM starts. Args: (depth, model, prompt_preview).
112109
on_subcall_complete: Callback fired when a child RLM completes. Args: (depth, model, duration, error_or_none).
113110
on_iteration_start: Callback fired when an iteration starts. Args: (depth, iteration_num).
@@ -163,8 +160,6 @@ def __init__(
163160
self._last_error: str | None = None
164161
self._best_partial_answer: str | None = None
165162
self._completion_start_time: float | None = None # Set when completion() starts
166-
# Lock for thread-safe cost/error tracking when children run in parallel
167-
self._state_lock = threading.Lock()
168163

169164
# Persistence support
170165
self.persistent = persistent
@@ -240,15 +235,15 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]):
240235
# For local environment with max_depth > 1, pass subcall callback for recursive RLM calls
241236
if self.environment_type == "local" and self.max_depth > 1:
242237
env_kwargs["subcall_fn"] = self._subcall
243-
env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls
244238
# Pass custom tools to the environment
245239
if self.custom_tools is not None:
246240
env_kwargs["custom_tools"] = self.custom_tools
247241
if self.custom_sub_tools is not None:
248242
env_kwargs["custom_sub_tools"] = self.custom_sub_tools
249243
if self.compaction and self.environment_type == "local":
250244
env_kwargs["compaction"] = True
251-
env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls
245+
if self.environment_type == "local":
246+
env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls
252247
environment: BaseEnv = get_environment(self.environment_type, env_kwargs)
253248

254249
if self.persistent:
@@ -318,15 +313,6 @@ def completion(
318313
# Check timeout before each iteration
319314
self._check_timeout(i, time_start)
320315

321-
# Fire iteration start callback
322-
if self.on_iteration_start:
323-
try:
324-
self.on_iteration_start(self.depth, i + 1)
325-
except Exception:
326-
pass
327-
328-
iter_start = time.perf_counter()
329-
330316
# Compaction: check if context needs summarization
331317
if self.compaction and hasattr(environment, "append_compaction_entry"):
332318
current_tokens, threshold_tokens, max_tokens = self._get_compaction_status(
@@ -390,13 +376,6 @@ def completion(
390376
# Verbose output for this iteration
391377
self.verbose.print_iteration(iteration, i + 1)
392378

393-
# Fire iteration complete callback
394-
if self.on_iteration_complete:
395-
try:
396-
self.on_iteration_complete(self.depth, i + 1, time.perf_counter() - iter_start)
397-
except Exception:
398-
pass
399-
400379
if final_answer is not None:
401380
time_end = time.perf_counter()
402381
usage = lm_handler.get_usage_summary()
@@ -728,11 +707,9 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion:
728707
)
729708

730709
# Calculate remaining budget for child (if budget tracking enabled)
731-
# Read _cumulative_cost under lock for thread-safe snapshot
732710
remaining_budget = None
733711
if self.max_budget is not None:
734-
with self._state_lock:
735-
remaining_budget = self.max_budget - self._cumulative_cost
712+
remaining_budget = self.max_budget - self._cumulative_cost
736713
if remaining_budget <= 0:
737714
return RLMChatCompletion(
738715
root_model=resolved_model,
@@ -799,20 +776,16 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion:
799776
# Propagate callbacks to children for nested tracking
800777
on_subcall_start=self.on_subcall_start,
801778
on_subcall_complete=self.on_subcall_complete,
802-
on_iteration_start=self.on_iteration_start,
803-
on_iteration_complete=self.on_iteration_complete,
804779
)
805780
try:
806781
result = child.completion(prompt, root_prompt=None)
807-
# Thread-safe: track child's cost in parent's cumulative cost
782+
# Track child's cost in parent's cumulative cost
808783
if result.usage_summary and result.usage_summary.total_cost:
809-
with self._state_lock:
810-
self._cumulative_cost += result.usage_summary.total_cost
784+
self._cumulative_cost += result.usage_summary.total_cost
811785
return result
812786
except BudgetExceededError as e:
813-
# Thread-safe: propagate child's spending to parent
814-
with self._state_lock:
815-
self._cumulative_cost += e.spent
787+
# Propagate child's spending to parent
788+
self._cumulative_cost += e.spent
816789
error_msg = f"Budget exceeded - {e}"
817790
return RLMChatCompletion(
818791
root_model=resolved_model,

‎rlm/environments/base_env.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,7 @@ class BaseEnv(ABC):
213213
"""
214214

215215
def __init__(
216-
self, persistent: bool = False, depth: int = 1, max_concurrent_subcalls: int = 4, **kwargs
216+
self, persistent: bool = False, depth: int = 1, max_concurrent_subcalls: int = 1, **kwargs
217217
):
218218
self.persistent = persistent
219219
self.depth = depth

‎rlm/environments/local_repl.py‎

Lines changed: 16 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -22,38 +22,6 @@
2222
validate_custom_tools,
2323
)
2424

25-
# =============================================================================
26-
# Global Subcall Semaphore
27-
# =============================================================================
28-
# Bounds the total number of concurrent child RLMs across ALL depths in
29-
# the process. Each depth level creates its own ThreadPoolExecutor (avoiding
30-
# pool-within-pool deadlock), but every child must acquire a slot from this
31-
# semaphore before doing real work. When all slots are taken, new children
32-
# wait until a running child finishes.
33-
#
34-
# Default: 16. Override via set_global_max_subcalls() or pass a custom
35-
# semaphore to LocalREPL(subcall_semaphore=...).
36-
37-
_DEFAULT_MAX_GLOBAL_SUBCALLS = 16
38-
_global_subcall_semaphore = threading.Semaphore(_DEFAULT_MAX_GLOBAL_SUBCALLS)
39-
40-
41-
42-
def set_global_max_subcalls(n: int) -> None:
43-
"""Replace the global subcall semaphore with a new one of size n.
44-
45-
Call this BEFORE creating any RLM instances. Not safe to call while
46-
subcalls are in flight.
47-
"""
48-
global _global_subcall_semaphore
49-
_global_subcall_semaphore = threading.Semaphore(n)
50-
51-
52-
def get_global_subcall_semaphore() -> threading.Semaphore:
53-
"""Return the current global subcall semaphore."""
54-
return _global_subcall_semaphore
55-
56-
5725
# =============================================================================
5826
# Safe Builtins
5927
# =============================================================================
@@ -169,20 +137,13 @@ def __init__(
169137
custom_sub_tools: dict[str, Any] | None = None,
170138
compaction: bool = False,
171139
max_concurrent_subcalls: int = 1,
172-
subcall_semaphore: threading.Semaphore | None = None,
173140
**kwargs,
174141
):
175-
super().__init__(
176-
persistent=persistent,
177-
depth=depth,
178-
max_concurrent_subcalls=max_concurrent_subcalls,
179-
**kwargs,
180-
)
142+
super().__init__(persistent=persistent, depth=depth, **kwargs)
181143

182144
self.lm_handler_address = lm_handler_address
183145
self.subcall_fn = subcall_fn # Callback for recursive RLM calls (depth > 1 support)
184146
self.max_concurrent_subcalls = max_concurrent_subcalls
185-
self._subcall_semaphore = subcall_semaphore or _global_subcall_semaphore
186147
self.original_cwd = os.getcwd()
187148
self.temp_dir = tempfile.mkdtemp(prefix=f"repl_env_{uuid.uuid4()}_")
188149
self._lock = threading.Lock()
@@ -360,19 +321,11 @@ def _rlm_query(self, prompt: str, model: str | None = None) -> str:
360321

361322
def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> list[str]:
362323
"""Spawn recursive RLM sub-calls for multiple prompts in parallel.
363-
Each prompt gets its own child RLM for deeper thinking.
364324
365-
Execution modes (controlled by max_concurrent_subcalls):
366-
- 1 (default): Sequential execution, identical to previous behavior.
367-
Zero overhead — no locks, no thread pool, no semaphore.
368-
- >1: Parallel execution via ThreadPoolExecutor. Each child RLM
369-
runs in its own thread with its own LMHandler and LocalREPL.
370-
A global semaphore bounds total concurrent children across all
371-
depths to prevent thread/memory explosion with deep recursion.
372-
373-
No deadlock risk: each depth level creates its own short-lived
374-
ThreadPoolExecutor (no pool-within-pool). The global semaphore does
375-
not own threads — workers acquire a slot, do work, release.
325+
Each prompt gets its own child RLM for deeper thinking. When multiple
326+
prompts are provided and max_concurrent_subcalls > 1, subcalls run
327+
concurrently using a thread pool. Results are returned in the same
328+
order as input prompts.
376329
377330
Falls back to llm_query_batched if no recursive capability is configured.
378331
@@ -396,63 +349,26 @@ def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> li
396349
for prompt in prompts:
397350
try:
398351
completion = self.subcall_fn(prompt, model)
399-
with lock:
400-
completions.append((index, completion))
401-
results[index] = completion.response
352+
self._pending_llm_calls.append(completion)
353+
results.append(completion.response)
402354
except Exception as e:
403-
results[index] = f"Error: RLM query failed - {e}"
404-
405-
with ThreadPoolExecutor(max_workers=max_workers) as executor:
406-
futures = [
407-
executor.submit(_run_subcall, i, prompt) for i, prompt in enumerate(prompts)
408-
]
409-
# Wait for all futures to complete; exceptions are captured inside _run_subcall
410-
for future in as_completed(futures):
411-
future.result() # Re-raises unexpected executor errors
412-
413-
# Append completions in original prompt order for deterministic metadata
414-
completions.sort(key=lambda x: x[0])
415-
for _, completion in completions:
416-
self._pending_llm_calls.append(completion)
417-
355+
results.append(f"Error: RLM query failed - {e}")
418356
return results
419357

420358
# Fall back to plain batched LM call if no recursive capability
421359
return self._llm_query_batched(prompts, model)
422360

423361
def _rlm_query_batched_parallel(self, prompts: list[str], model: str | None = None) -> list[str]:
424-
"""Execute multiple RLM sub-calls in parallel using ThreadPoolExecutor.
425-
426-
Each child RLM gets its own thread. Results maintain the original
427-
prompt order regardless of which child finishes first.
428-
429-
Concurrency control (two layers):
430-
1. Local ThreadPoolExecutor: bounds children dispatched by THIS REPL
431-
(max_concurrent_subcalls). Short-lived — created and destroyed
432-
per batch call.
433-
2. Global semaphore (_subcall_semaphore): bounds total concurrent
434-
children across ALL depths in the process. A worker acquires a
435-
slot before calling subcall_fn and releases it after, so the
436-
semaphore never blocks the pool itself — only the actual work.
437-
438-
Thread safety:
439-
- Each child RLM is fully isolated (own LMHandler, own LocalREPL).
440-
- _pending_llm_calls is protected by _calls_lock for concurrent appends.
441-
- subcall_fn (RLM._subcall) is called from worker threads — the parent
442-
RLM's _state_lock protects cumulative_cost and error tracking.
443-
"""
362+
"""Execute multiple RLM sub-calls in parallel using ThreadPoolExecutor."""
444363
results: list[str | None] = [None] * len(prompts)
445-
sem = self._subcall_semaphore
364+
completions: list[tuple[int, RLMChatCompletion]] = []
446365

447366
def _execute_one(idx: int, prompt: str) -> tuple[int, str, RLMChatCompletion | None]:
448-
sem.acquire()
449367
try:
450368
completion = self.subcall_fn(prompt, model)
451369
return (idx, completion.response, completion)
452370
except Exception as e:
453371
return (idx, f"Error: RLM query failed - {e}", None)
454-
finally:
455-
sem.release()
456372

457373
max_workers = min(self.max_concurrent_subcalls, len(prompts))
458374
with ThreadPoolExecutor(max_workers=max_workers) as executor:
@@ -464,8 +380,12 @@ def _execute_one(idx: int, prompt: str) -> tuple[int, str, RLMChatCompletion | N
464380
idx, response, completion = future.result()
465381
results[idx] = response
466382
if completion is not None:
467-
with self._calls_lock:
468-
self._pending_llm_calls.append(completion)
383+
completions.append((idx, completion))
384+
385+
# Append in prompt order for deterministic metadata
386+
completions.sort(key=lambda x: x[0])
387+
for _, completion in completions:
388+
self._pending_llm_calls.append(completion)
469389

470390
return results
471391

0 commit comments

Comments
 (0)