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# =============================================================================
@@ -168,21 +136,14 @@ def __init__(
168136 custom_tools : dict [str , Any ] | None = None ,
169137 custom_sub_tools : dict [str , Any ] | None = None ,
170138 compaction : bool = False ,
171- max_concurrent_subcalls : int = 1 ,
172- subcall_semaphore : threading .Semaphore | None = None ,
139+ max_concurrent_subcalls : int = 4 ,
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