diff --git a/docs/source/environments/repl.md b/docs/source/environments/repl.md index 0ec937bb2d..d0fb6b0c38 100644 --- a/docs/source/environments/repl.md +++ b/docs/source/environments/repl.md @@ -1,19 +1,12 @@ # REPL Environment for OpenEnv -`repl_env` is a Python REPL environment for [Recursive Language Model](https://huggingface.co/papers/2512.24601) (RLM) style execution. The model writes code that runs in a persistent namespace and can: +`repl_env` runs model-generated Python in a persistent namespace, with optional +recursive model queries, following the [RLM paradigm](https://huggingface.co/papers/2512.24601). -- inspect `context` -- execute Python across multiple turns with persistent state -- call `llm_query(...)` and `llm_query_batched(...)` to query a language model -- call `rlm_query(...)` and `rlm_query_batched(...)` for recursive child runs, when configured -- finish with `FINAL(...)`, `FINAL_VAR(...)`, or `answer = {"content": ..., "ready": True}` - -The package provides: - -- `REPLEnv`: the async client for a remote server (`.sync()` for synchronous code), with `execute(code)`, `submit_final_answer(answer)`, `get_variable(name)` and `list_variables()` on top of `reset()`/`step()` -- `LocalREPLEnv`: the same environment, in process -- `LocalRLMRunner`: a local RLM loop that prompts a model, runs its code and handles recursion +- `REPLEnv`: remote async client, or synchronous client via `.sync()`. +- `LocalREPLEnv`: the environment in process. +- `LocalRLMRunner`: the model prompting, code-execution, and recursion loop. Try it on the [`openenv/repl`](https://huggingface.co/spaces/openenv/repl) Space, or connect a client to `https://openenv-repl.hf.space`. @@ -46,20 +39,9 @@ async def main(): asyncio.run(main()) ``` -Sync: - -```python -from repl_env import REPLEnv - -with REPLEnv(base_url="http://127.0.0.1:8000").sync() as env: - result = env.reset( - context="alpha beta gamma", - task_prompt="Count the words", - ) - result = env.execute("count = len(context.split())") - result = env.execute("print(FINAL(count))") - print(result.observation.result.stdout) -``` +For synchronous remote use, enter `with REPLEnv(...).sync() as env:` and omit +`await`. Both clients also provide `submit_final_answer()`, `get_variable()`, and +`list_variables()`. ### In Process @@ -93,198 +75,143 @@ Environment variables read by [`server/app.py`](https://github.com/openenv/OpenE ## Reward -Rewards use the OpenEnv [rubric system](https://huggingface.co/docs/openenv/guides/rewards). The default `REPLRubric` combines: - -- **Outcome reward** (on terminal steps): compares `final_answer` against - `expected_answer` if provided. Returns 1.0 for match, 0.0 otherwise. -- **Process reward** (on non-terminal steps): returns -0.05 for code - execution errors, 0.0 for successful steps. -- **Failure reward**: returns -0.1 when max iterations exhausted without an answer. - -For RL training (GRPO, etc.), pass `expected_answer` at reset time: +The default `REPLRubric` gives 1.0 for a matching `expected_answer`, 0.0 for a +mismatch or successful intermediate step, -0.05 for execution errors, and -0.1 +when iterations run out without an answer. Set the expected answer at reset: ```python with LocalREPLEnv() as env: - env.reset( - context="...", - task_prompt="...", - expected_answer="42", # ground truth for rubric scoring - ) + env.reset(expected_answer="42") result = env.execute("print(FINAL(42))") - print(result.reward) # 1.0 (correct) + print(result.reward) ``` -Other rubrics: `ExactMatchRubric` (binary match), `FuzzyMatchRubric` (1.0 for an exact match, 0.5 when the expected answer is contained in the final answer), `CustomMetricRubric` (your `metric(expected, predicted) -> float`) and `CodeExecutionRubric` (per-step error penalty). Pass one at construction: - -```python -from repl_env import LocalREPLEnv, CustomMetricRubric, REPLRubric - -def my_metric(expected, predicted): - return 1.0 if expected.strip() == predicted.strip() else 0.0 - -env = LocalREPLEnv(rubric=REPLRubric(outcome=CustomMetricRubric(my_metric))) -``` +Customize `LocalREPLEnv(rubric=...)` with the match, partial-credit, custom-metric, +or execution rubrics in [`rubrics.py`](https://github.com/openenv/OpenEnv/blob/main/envs/repl_env/rubrics.py). See the +[rubric guide](https://huggingface.co/docs/openenv/guides/rewards) for composition. ## Running an RLM Locally -`LocalRLMRunner` takes any `chat_fn(messages, model=None) -> str`. It works -with HF Inference API, vLLM, SGLang, Ollama, or any OpenAI-compatible server. - -With HF Inference API: +`LocalRLMRunner` accepts any `chat_fn(messages, model=None) -> str`, including +adapters for hosted or local model servers. For HF inference, save this reusable +factory in `my_model.py` (uses `HF_TOKEN` or your cached login): ```python from huggingface_hub import InferenceClient -from repl_env import LocalRLMRunner, RLM_SYSTEM_PROMPT -client = InferenceClient(model="Qwen/Qwen3.5-9B", timeout=300) -def chat_fn(messages, model=None): - response = client.chat.completions.create( - model=model or "Qwen/Qwen3.5-9B", - messages=messages, - max_tokens=2048, - temperature=0.6, - extra_body={"chat_template_kwargs": {"enable_thinking": False}}, - ) - return response.choices[0].message.content +def make_chat(): + client = InferenceClient(model="Qwen/Qwen3.5-9B", timeout=300) -runner = LocalRLMRunner(chat_fn, max_iterations=30, max_depth=2) -result = runner.run("The answer is 42", "What number is mentioned?") -print(result.final_answer) -``` + def chat(messages, model=None): + response = client.chat.completions.create( + model=model or "Qwen/Qwen3.5-9B", + messages=messages, + max_tokens=2048, + temperature=0.6, + extra_body={"chat_template_kwargs": {"enable_thinking": False}}, + ) + return response.choices[0].message.content or "" -With a local vLLM server: + return chat +``` ```python -from openai import OpenAI +from my_model import make_chat from repl_env import LocalRLMRunner -client = OpenAI(base_url="http://localhost:8000/v1", api_key="unused") - -def chat_fn(messages, model=None): - response = client.chat.completions.create( - model=model or "Qwen/Qwen3.5-9B", - messages=messages, - max_tokens=2048, - temperature=0.6, - ) - return response.choices[0].message.content - -runner = LocalRLMRunner(chat_fn, max_iterations=30, max_depth=2) -result = runner.run(context, task) +runner = LocalRLMRunner(make_chat(), max_iterations=30, max_depth=2) +result = runner.run("The answer is 42", "What number is mentioned?") +print(result.final_answer) ``` -### Different Models for Outer and Inner Loops - -The outer loop (code generation) can use a large model while inner -`llm_query`/`rlm_query` calls use a smaller, faster model. Pass a -custom `backend_factory` to the runner: - -```python -from openai import OpenAI -from huggingface_hub import InferenceClient -from repl_env import LocalRLMRunner -from repl_env.recursive_backends import BackendLimits, LocalChildRLMBackend - -# Outer loop: large local model via vLLM -vllm = OpenAI(base_url="http://localhost:8000/v1", api_key="unused") +The runner supports child-count, batch, timeout, and result-length limits, plus +`on_subcall_start` / `on_subcall_complete` callbacks and child traces. For different +outer/inner models, supply a `backend_factory`; see +[`recursive_backends.py`](https://github.com/openenv/OpenEnv/blob/main/envs/repl_env/recursive_backends.py). -def outer_chat(messages, model=None): - r = vllm.chat.completions.create( - model="Qwen/Qwen3-32B", messages=messages, max_tokens=2048, - ) - return r.choices[0].message.content +### Managed child lifecycle (opt-in) -# Inner calls (llm_query/rlm_query): smaller HF-hosted model -hf = InferenceClient(model="Qwen/Qwen3.5-9B") +Managed children let the parent keep working while queries run in supervised +processes. Each worker builds its own client; the model server stays shared. +Use an importable factory such as `my_model.make_chat` above, not a lambda or live +client, and guard the entry point for Python's `spawn` start method: -def inner_chat(messages, model=None): - r = hf.chat.completions.create( - model=model or "Qwen/Qwen3.5-9B", messages=messages, max_tokens=2048, - extra_body={"chat_template_kwargs": {"enable_thinking": False}}, - ) - return r.choices[0].message.content - -def my_backend_factory(llm_chat_fn, **kwargs): - return LocalChildRLMBackend( - inner_chat, # inner calls use the smaller model - runner_factory=LocalRLMRunner, - system_prompt=kwargs["system_prompt"], - max_iterations=kwargs["max_iterations"], - env_max_iterations_multiplier=kwargs["env_max_iterations_multiplier"], - depth=kwargs["depth"], - limits=BackendLimits(max_depth=2), +```python +from my_model import make_chat +from repl_env import LocalRLMRunner, ManagedChildConfig + +if __name__ == "__main__": + runner = LocalRLMRunner( + make_chat(), + managed_children=ManagedChildConfig(chat_factory=make_chat), + max_depth=3, + max_iterations=10, + max_children_total=8, + max_batch_workers=4, + per_child_timeout_s=60, ) - -runner = LocalRLMRunner( - outer_chat, # outer loop: large model - backend_factory=my_backend_factory, # inner calls: small model - max_iterations=30, - max_depth=2, -) -result = runner.run(context, task) + print(runner.run("alpha beta gamma", "Count the words").final_answer) ``` -`LocalRLMRunner` also takes recursion limits (`max_children_total`, `max_children_per_batch`, `per_child_timeout_s`, `result_truncation_limit`) and lifecycle callbacks (`on_subcall_start(depth, model, prompt_preview)`, `on_subcall_complete(depth, model, duration, error_or_none)`). Its results carry lightweight child trace metadata. - -## Actions and Observations - -`REPLAction` +Inside the model's REPL: ```python -code: str = "" -is_final: bool = False -final_answer: str | None = None +child = rlm_query("Count the words: alpha beta gamma", blocking=False) +parent_count = len(context.split()) +child_answer = rlm_wait(child, timeout_s=30) +print(FINAL(parent_count)) ``` -`REPLObservation` +- `rlm_query(prompt, model=None, *, blocking=True)` waits for an answer by default. + With `blocking=False`, it returns an episode-scoped handle instead and requires + managed mode. Use `rlm_status` to inspect the handle, + `rlm_result` for a completed answer, or `rlm_wait` for a bounded wait. Failed or + unfinished results raise errors; a wait timeout does **not** cancel the child. +- Nested workers share child-count and concurrency limits. Full concurrency + rejects new spawns rather than queuing them. Child timeouts interrupt even + in-flight model calls; deeper explicit spawns are rejected at the depth limit. +- Finalization, reset, and close stop admission, request cooperative shutdown, + then escalate to SIGTERM and SIGKILL. Cleanup waits for every registered worker + to exit; failure raises `ChildCleanupError` and blocks reuse. Reset invalidates + handles and replenishes the budget; close is permanent and idempotent. +- Controller lifecycle transitions are serialized. Callbacks are asynchronous, + ordered notifications: blocked or failing observers do not delay cleanup, and + notifications may run after reset/close. REPL execution itself is not thread-safe. + +`ManagedChildConfig` exposes `shutdown_grace_s`, `terminate_grace_s`, and +`kill_timeout_s` (defaults: 1, 1, and 5 seconds). This POSIX-only backend does not +contain arbitrary subprocesses or guarantee server-side generation abort. +User-facing cancellation and follow-up messages are deferred. + +For manual/server use, pass `managed_children` to +[`create_server_recursive_controller`](https://github.com/openenv/OpenEnv/blob/main/envs/repl_env/recursive_controller.py), then supply that +controller to `LocalREPLEnv` or `REPLEnvironment` via `recursive_controller=`. +Configure its credentials and limits at construction, not through reset overrides; +infrastructure can cancel a handle with `controller.children.cancel(handle)`. +`managed_children` cannot be combined with a custom `backend_factory`. -```python -result: CodeBlockResult -context_preview: str | None -context_length: int -available_variables: list[str] -iteration: int -max_iterations: int -done: bool -reward: float | None -metadata: dict -``` +## Actions and Observations + +`REPLAction` carries `code`, or `is_final=True` with a `final_answer` for explicit +submission. `REPLObservation` includes execution results (stdout/stderr, success, +exceptions, and timing), context/iteration metadata, `done`, and `reward`. +See [`models.py`](https://github.com/openenv/OpenEnv/blob/main/envs/repl_env/models.py) for all fields. ## REPL Helpers When configured, the REPL namespace exposes: - `llm_query(prompt, model=None)` and `llm_query_batched(prompts, model=None)` -- `rlm_query(prompt, model=None)` and `rlm_query_batched(prompts, model=None)`: recursive child runs. At the maximum depth they fall back to direct model calls. -- `FINAL(value)`, `FINAL_VAR(name)` and `SHOW_VARS()` - -## Finalization Patterns - -### `FINAL(...)` - -```python -result = env.execute("answer = 42") -result = env.execute("print(FINAL(answer))") -``` +- `rlm_query(prompt, model=None, *, blocking=True)` and `rlm_query_batched(prompts, model=None)`: recursive child runs that wait for their answers. Blocking queries fall back to direct model calls at the recursion limit. In managed mode, `rlm_query(..., blocking=False)` starts a background child and returns a handle. +- `SHOW_VARS()` lists variables in the namespace. -### `FINAL_VAR(...)` - -```python -result = env.execute("my_answer = '42'") -result = env.execute('print(FINAL_VAR("my_answer"))') -``` - -### `answer` dict - -```python -result = env.execute("answer['content'] = '42'") -result = env.execute("answer['ready'] = True") -``` +Finish by executing `print(FINAL(value))`, `print(FINAL_VAR("variable_name"))`, or +`answer.update({"content": value, "ready": True})`. ## Prompts and Examples -[`prompts.py`](https://github.com/openenv/OpenEnv/blob/main/envs/repl_env/prompts.py) has the system prompts and helpers used by the runner: `RLM_SYSTEM_PROMPT`, `RLM_SYSTEM_PROMPT_QWEN`, `QueryMetadata`, `build_rlm_system_prompt(...)`, `build_user_prompt(...)`, `extract_code_blocks(...)` and `format_observations(...)`. +See [`prompts.py`](https://github.com/openenv/OpenEnv/blob/main/envs/repl_env/prompts.py) for system prompts and message/code-block helpers. Examples, which default to `Qwen/Qwen3.5-9B` through Hugging Face inference (needs `HF_TOKEN`): diff --git a/envs/repl_env/README.md b/envs/repl_env/README.md index 8b5d7d2bf2..6159fc7160 100644 --- a/envs/repl_env/README.md +++ b/envs/repl_env/README.md @@ -13,19 +13,12 @@ tags: # REPL Environment for OpenEnv -`repl_env` is a Python REPL environment for [Recursive Language Model](https://huggingface.co/papers/2512.24601) (RLM) style execution. The model writes code that runs in a persistent namespace and can: +`repl_env` runs model-generated Python in a persistent namespace, with optional +recursive model queries, following the [RLM paradigm](https://huggingface.co/papers/2512.24601). -- inspect `context` -- execute Python across multiple turns with persistent state -- call `llm_query(...)` and `llm_query_batched(...)` to query a language model -- call `rlm_query(...)` and `rlm_query_batched(...)` for recursive child runs, when configured -- finish with `FINAL(...)`, `FINAL_VAR(...)`, or `answer = {"content": ..., "ready": True}` - -The package provides: - -- `REPLEnv`: the async client for a remote server (`.sync()` for synchronous code), with `execute(code)`, `submit_final_answer(answer)`, `get_variable(name)` and `list_variables()` on top of `reset()`/`step()` -- `LocalREPLEnv`: the same environment, in process -- `LocalRLMRunner`: a local RLM loop that prompts a model, runs its code and handles recursion +- `REPLEnv`: remote async client, or synchronous client via `.sync()`. +- `LocalREPLEnv`: the environment in process. +- `LocalRLMRunner`: the model prompting, code-execution, and recursion loop. Try it on the [`openenv/repl`](https://huggingface.co/spaces/openenv/repl) Space, or connect a client to `https://openenv-repl.hf.space`. @@ -58,20 +51,9 @@ async def main(): asyncio.run(main()) ``` -Sync: - -```python -from repl_env import REPLEnv - -with REPLEnv(base_url="http://127.0.0.1:8000").sync() as env: - result = env.reset( - context="alpha beta gamma", - task_prompt="Count the words", - ) - result = env.execute("count = len(context.split())") - result = env.execute("print(FINAL(count))") - print(result.observation.result.stdout) -``` +For synchronous remote use, enter `with REPLEnv(...).sync() as env:` and omit +`await`. Both clients also provide `submit_final_answer()`, `get_variable()`, and +`list_variables()`. ### In Process @@ -105,198 +87,143 @@ Environment variables read by [`server/app.py`](server/app.py): ## Reward -Rewards use the OpenEnv [rubric system](https://huggingface.co/docs/openenv/guides/rewards). The default `REPLRubric` combines: - -- **Outcome reward** (on terminal steps): compares `final_answer` against - `expected_answer` if provided. Returns 1.0 for match, 0.0 otherwise. -- **Process reward** (on non-terminal steps): returns -0.05 for code - execution errors, 0.0 for successful steps. -- **Failure reward**: returns -0.1 when max iterations exhausted without an answer. - -For RL training (GRPO, etc.), pass `expected_answer` at reset time: +The default `REPLRubric` gives 1.0 for a matching `expected_answer`, 0.0 for a +mismatch or successful intermediate step, -0.05 for execution errors, and -0.1 +when iterations run out without an answer. Set the expected answer at reset: ```python with LocalREPLEnv() as env: - env.reset( - context="...", - task_prompt="...", - expected_answer="42", # ground truth for rubric scoring - ) + env.reset(expected_answer="42") result = env.execute("print(FINAL(42))") - print(result.reward) # 1.0 (correct) + print(result.reward) ``` -Other rubrics: `ExactMatchRubric` (binary match), `FuzzyMatchRubric` (1.0 for an exact match, 0.5 when the expected answer is contained in the final answer), `CustomMetricRubric` (your `metric(expected, predicted) -> float`) and `CodeExecutionRubric` (per-step error penalty). Pass one at construction: - -```python -from repl_env import LocalREPLEnv, CustomMetricRubric, REPLRubric - -def my_metric(expected, predicted): - return 1.0 if expected.strip() == predicted.strip() else 0.0 - -env = LocalREPLEnv(rubric=REPLRubric(outcome=CustomMetricRubric(my_metric))) -``` +Customize `LocalREPLEnv(rubric=...)` with the match, partial-credit, custom-metric, +or execution rubrics in [`rubrics.py`](rubrics.py). See the +[rubric guide](https://huggingface.co/docs/openenv/guides/rewards) for composition. ## Running an RLM Locally -`LocalRLMRunner` takes any `chat_fn(messages, model=None) -> str`. It works -with HF Inference API, vLLM, SGLang, Ollama, or any OpenAI-compatible server. - -With HF Inference API: +`LocalRLMRunner` accepts any `chat_fn(messages, model=None) -> str`, including +adapters for hosted or local model servers. For HF inference, save this reusable +factory in `my_model.py` (uses `HF_TOKEN` or your cached login): ```python from huggingface_hub import InferenceClient -from repl_env import LocalRLMRunner, RLM_SYSTEM_PROMPT -client = InferenceClient(model="Qwen/Qwen3.5-9B", timeout=300) -def chat_fn(messages, model=None): - response = client.chat.completions.create( - model=model or "Qwen/Qwen3.5-9B", - messages=messages, - max_tokens=2048, - temperature=0.6, - extra_body={"chat_template_kwargs": {"enable_thinking": False}}, - ) - return response.choices[0].message.content +def make_chat(): + client = InferenceClient(model="Qwen/Qwen3.5-9B", timeout=300) -runner = LocalRLMRunner(chat_fn, max_iterations=30, max_depth=2) -result = runner.run("The answer is 42", "What number is mentioned?") -print(result.final_answer) -``` + def chat(messages, model=None): + response = client.chat.completions.create( + model=model or "Qwen/Qwen3.5-9B", + messages=messages, + max_tokens=2048, + temperature=0.6, + extra_body={"chat_template_kwargs": {"enable_thinking": False}}, + ) + return response.choices[0].message.content or "" -With a local vLLM server: + return chat +``` ```python -from openai import OpenAI +from my_model import make_chat from repl_env import LocalRLMRunner -client = OpenAI(base_url="http://localhost:8000/v1", api_key="unused") - -def chat_fn(messages, model=None): - response = client.chat.completions.create( - model=model or "Qwen/Qwen3.5-9B", - messages=messages, - max_tokens=2048, - temperature=0.6, - ) - return response.choices[0].message.content - -runner = LocalRLMRunner(chat_fn, max_iterations=30, max_depth=2) -result = runner.run(context, task) +runner = LocalRLMRunner(make_chat(), max_iterations=30, max_depth=2) +result = runner.run("The answer is 42", "What number is mentioned?") +print(result.final_answer) ``` -### Different Models for Outer and Inner Loops - -The outer loop (code generation) can use a large model while inner -`llm_query`/`rlm_query` calls use a smaller, faster model. Pass a -custom `backend_factory` to the runner: - -```python -from openai import OpenAI -from huggingface_hub import InferenceClient -from repl_env import LocalRLMRunner -from repl_env.recursive_backends import BackendLimits, LocalChildRLMBackend - -# Outer loop: large local model via vLLM -vllm = OpenAI(base_url="http://localhost:8000/v1", api_key="unused") +The runner supports child-count, batch, timeout, and result-length limits, plus +`on_subcall_start` / `on_subcall_complete` callbacks and child traces. For different +outer/inner models, supply a `backend_factory`; see +[`recursive_backends.py`](recursive_backends.py). -def outer_chat(messages, model=None): - r = vllm.chat.completions.create( - model="Qwen/Qwen3-32B", messages=messages, max_tokens=2048, - ) - return r.choices[0].message.content +### Managed child lifecycle (opt-in) -# Inner calls (llm_query/rlm_query): smaller HF-hosted model -hf = InferenceClient(model="Qwen/Qwen3.5-9B") +Managed children let the parent keep working while queries run in supervised +processes. Each worker builds its own client; the model server stays shared. +Use an importable factory such as `my_model.make_chat` above, not a lambda or live +client, and guard the entry point for Python's `spawn` start method: -def inner_chat(messages, model=None): - r = hf.chat.completions.create( - model=model or "Qwen/Qwen3.5-9B", messages=messages, max_tokens=2048, - extra_body={"chat_template_kwargs": {"enable_thinking": False}}, - ) - return r.choices[0].message.content - -def my_backend_factory(llm_chat_fn, **kwargs): - return LocalChildRLMBackend( - inner_chat, # inner calls use the smaller model - runner_factory=LocalRLMRunner, - system_prompt=kwargs["system_prompt"], - max_iterations=kwargs["max_iterations"], - env_max_iterations_multiplier=kwargs["env_max_iterations_multiplier"], - depth=kwargs["depth"], - limits=BackendLimits(max_depth=2), +```python +from my_model import make_chat +from repl_env import LocalRLMRunner, ManagedChildConfig + +if __name__ == "__main__": + runner = LocalRLMRunner( + make_chat(), + managed_children=ManagedChildConfig(chat_factory=make_chat), + max_depth=3, + max_iterations=10, + max_children_total=8, + max_batch_workers=4, + per_child_timeout_s=60, ) - -runner = LocalRLMRunner( - outer_chat, # outer loop: large model - backend_factory=my_backend_factory, # inner calls: small model - max_iterations=30, - max_depth=2, -) -result = runner.run(context, task) + print(runner.run("alpha beta gamma", "Count the words").final_answer) ``` -`LocalRLMRunner` also takes recursion limits (`max_children_total`, `max_children_per_batch`, `per_child_timeout_s`, `result_truncation_limit`) and lifecycle callbacks (`on_subcall_start(depth, model, prompt_preview)`, `on_subcall_complete(depth, model, duration, error_or_none)`). Its results carry lightweight child trace metadata. - -## Actions and Observations - -`REPLAction` +Inside the model's REPL: ```python -code: str = "" -is_final: bool = False -final_answer: str | None = None +child = rlm_query("Count the words: alpha beta gamma", blocking=False) +parent_count = len(context.split()) +child_answer = rlm_wait(child, timeout_s=30) +print(FINAL(parent_count)) ``` -`REPLObservation` +- `rlm_query(prompt, model=None, *, blocking=True)` waits for an answer by default. + With `blocking=False`, it returns an episode-scoped handle instead and requires + managed mode. Use `rlm_status` to inspect the handle, + `rlm_result` for a completed answer, or `rlm_wait` for a bounded wait. Failed or + unfinished results raise errors; a wait timeout does **not** cancel the child. +- Nested workers share child-count and concurrency limits. Full concurrency + rejects new spawns rather than queuing them. Child timeouts interrupt even + in-flight model calls; deeper explicit spawns are rejected at the depth limit. +- Finalization, reset, and close stop admission, request cooperative shutdown, + then escalate to SIGTERM and SIGKILL. Cleanup waits for every registered worker + to exit; failure raises `ChildCleanupError` and blocks reuse. Reset invalidates + handles and replenishes the budget; close is permanent and idempotent. +- Controller lifecycle transitions are serialized. Callbacks are asynchronous, + ordered notifications: blocked or failing observers do not delay cleanup, and + notifications may run after reset/close. REPL execution itself is not thread-safe. + +`ManagedChildConfig` exposes `shutdown_grace_s`, `terminate_grace_s`, and +`kill_timeout_s` (defaults: 1, 1, and 5 seconds). This POSIX-only backend does not +contain arbitrary subprocesses or guarantee server-side generation abort. +User-facing cancellation and follow-up messages are deferred. + +For manual/server use, pass `managed_children` to +[`create_server_recursive_controller`](recursive_controller.py), then supply that +controller to `LocalREPLEnv` or `REPLEnvironment` via `recursive_controller=`. +Configure its credentials and limits at construction, not through reset overrides; +infrastructure can cancel a handle with `controller.children.cancel(handle)`. +`managed_children` cannot be combined with a custom `backend_factory`. -```python -result: CodeBlockResult -context_preview: str | None -context_length: int -available_variables: list[str] -iteration: int -max_iterations: int -done: bool -reward: float | None -metadata: dict -``` +## Actions and Observations + +`REPLAction` carries `code`, or `is_final=True` with a `final_answer` for explicit +submission. `REPLObservation` includes execution results (stdout/stderr, success, +exceptions, and timing), context/iteration metadata, `done`, and `reward`. +See [`models.py`](models.py) for all fields. ## REPL Helpers When configured, the REPL namespace exposes: - `llm_query(prompt, model=None)` and `llm_query_batched(prompts, model=None)` -- `rlm_query(prompt, model=None)` and `rlm_query_batched(prompts, model=None)`: recursive child runs. At the maximum depth they fall back to direct model calls. -- `FINAL(value)`, `FINAL_VAR(name)` and `SHOW_VARS()` - -## Finalization Patterns - -### `FINAL(...)` - -```python -result = env.execute("answer = 42") -result = env.execute("print(FINAL(answer))") -``` +- `rlm_query(prompt, model=None, *, blocking=True)` and `rlm_query_batched(prompts, model=None)`: recursive child runs that wait for their answers. Blocking queries fall back to direct model calls at the recursion limit. In managed mode, `rlm_query(..., blocking=False)` starts a background child and returns a handle. +- `SHOW_VARS()` lists variables in the namespace. -### `FINAL_VAR(...)` - -```python -result = env.execute("my_answer = '42'") -result = env.execute('print(FINAL_VAR("my_answer"))') -``` - -### `answer` dict - -```python -result = env.execute("answer['content'] = '42'") -result = env.execute("answer['ready'] = True") -``` +Finish by executing `print(FINAL(value))`, `print(FINAL_VAR("variable_name"))`, or +`answer.update({"content": value, "ready": True})`. ## Prompts and Examples -[`prompts.py`](prompts.py) has the system prompts and helpers used by the runner: `RLM_SYSTEM_PROMPT`, `RLM_SYSTEM_PROMPT_QWEN`, `QueryMetadata`, `build_rlm_system_prompt(...)`, `build_user_prompt(...)`, `extract_code_blocks(...)` and `format_observations(...)`. +See [`prompts.py`](prompts.py) for system prompts and message/code-block helpers. Examples, which default to `Qwen/Qwen3.5-9B` through Hugging Face inference (needs `HF_TOKEN`): diff --git a/envs/repl_env/__init__.py b/envs/repl_env/__init__.py index 544d2be1e5..c6542e9371 100644 --- a/envs/repl_env/__init__.py +++ b/envs/repl_env/__init__.py @@ -42,6 +42,7 @@ from .client import REPLEnv from .local import LocalREPLEnv +from .managed_children import ChildCleanupError, ManagedChildConfig from .models import CodeBlockResult, REPLAction, REPLObservation, REPLState from .prompts import ( build_rlm_system_prompt, @@ -75,6 +76,8 @@ "LocalREPLEnv", "LocalRLMRunner", "RLMRunResult", + "ManagedChildConfig", + "ChildCleanupError", # System prompts "RLM_SYSTEM_PROMPT", "RLM_SYSTEM_PROMPT_QWEN", diff --git a/envs/repl_env/local.py b/envs/repl_env/local.py index f4f8ed05b1..8565164975 100644 --- a/envs/repl_env/local.py +++ b/envs/repl_env/local.py @@ -16,6 +16,8 @@ from collections.abc import Callable from typing import Optional +from .recursive_controller import RecursiveController + try: from openenv.core.client_types import StepResult @@ -46,6 +48,7 @@ def __init__( rubric: Optional[REPLRubric] = None, rlm_max_depth: int = 1, rlm_max_iterations: int | None = None, + recursive_controller: RecursiveController | None = None, ): self._env = REPLEnvironment( max_output_length=max_output_length, @@ -57,6 +60,7 @@ def __init__( subcall_batch_fn=subcall_batch_fn, rlm_max_depth=rlm_max_depth, rlm_max_iterations=rlm_max_iterations, + recursive_controller=recursive_controller, ) def reset( diff --git a/envs/repl_env/managed_children.py b/envs/repl_env/managed_children.py new file mode 100644 index 0000000000..b9922bad01 --- /dev/null +++ b/envs/repl_env/managed_children.py @@ -0,0 +1,677 @@ +"""Opt-in, process-supervised recursive queries for the REPL environment.""" + +from __future__ import annotations + +import math +import multiprocessing +import os +import pickle +import signal +import threading +import time +from contextlib import suppress +from dataclasses import dataclass, field +from multiprocessing.connection import Connection +from queue import SimpleQueue +from typing import Callable +from uuid import uuid4 + +from .recursive_backends import BackendLimits, ChildTrace, DirectLMBackend + + +MANAGED_CHILD_TOOLS_PROMPT = """ +rlm_query(prompt, model=None, blocking=True) waits and returns an answer string. +To keep working while a child runs, use child = rlm_query(prompt, model=None, blocking=False). +The blocking option is keyword-only; blocking=False returns an episode-scoped handle. +Keep working, then inspect rlm_status(child), read rlm_result(child), or call +rlm_wait(child, timeout_s=30). Wait timeouts do not cancel work. Results are strings; +unfinished or failed children raise errors. Handles expire on environment reset. +Outstanding children are stopped when you finalize. Concurrency limits reject +new children rather than queueing them, including when nested parents occupy slots. +""" + + +class ChildCleanupError(RuntimeError): + """A managed worker's exit could not be confirmed.""" + + +@dataclass(frozen=True) +class ManagedChildConfig: + """Construct a fresh model client inside each spawned worker. + + Args: + chat_factory (`Callable`): + Picklable, importable callable returning a chat function. Live clients + and closures must not be passed across the process boundary. + shutdown_grace_s (`float`, *optional*, defaults to `1.0`): + Time to allow cooperative shutdown before sending SIGTERM. + terminate_grace_s (`float`, *optional*, defaults to `1.0`): + Time to wait after SIGTERM before sending SIGKILL. + kill_timeout_s (`float`, *optional*, defaults to `5.0`): + Final bounded process-exit wait after SIGKILL. + """ + + chat_factory: Callable[[], Callable[..., str]] = field(repr=False) + shutdown_grace_s: float = 1.0 + terminate_grace_s: float = 1.0 + kill_timeout_s: float = 5.0 + + def __post_init__(self): + if os.name != "posix": + raise ValueError("Managed children require POSIX process signals") + if not callable(self.chat_factory): + raise ValueError("chat_factory must be callable") + for duration in ( + self.shutdown_grace_s, + self.terminate_grace_s, + self.kill_timeout_s, + ): + if not math.isfinite(duration) or duration < 0: + raise ValueError("Shutdown timeouts must be finite and nonnegative") + if self.kill_timeout_s == 0: + raise ValueError("kill_timeout_s must be positive") + try: + pickle.dumps(self.chat_factory) + except (TypeError, AttributeError, pickle.PicklingError) as exc: + raise ValueError("chat_factory must be importable and picklable") from exc + + +@dataclass +class _Job: + process: multiprocessing.Process + channel: Connection + shutdown: Connection + prompt: str + depth: int + parent: _Job | None = None + model: str | None = None + accepting: bool = True + reported: bool = False + pid: int = field(init=False) + exitcode: int | None = None + started: float = field(default_factory=time.monotonic) + stop: threading.Event = field(default_factory=threading.Event) + finished: threading.Event = field(default_factory=threading.Event) + status: str = "running" + answer: str | None = None + error: str | None = None + thread: threading.Thread | None = None + + def __post_init__(self): + self.pid = self.process.pid + + +class ManagedChildren: + """Own all managed workers for one episode, including nested queries. + + Every process is spawned and reaped by this controller, not by another + worker. This registry does not cover arbitrary subprocesses created by a + custom client, nor does it guarantee cancellation on the model server. + """ + + def __init__( + self, + config: ManagedChildConfig, + limits: BackendLimits, + *, + on_subcall_start=None, + on_subcall_complete=None, + **runner_options, + ): + if limits.max_depth < 1 or limits.max_batch_workers < 1: + raise ValueError("Managed recursion depth and concurrency must be positive") + if limits.per_child_timeout_s is not None: + if ( + not math.isfinite(limits.per_child_timeout_s) + or limits.per_child_timeout_s <= 0 + ): + raise ValueError("per_child_timeout_s must be finite and positive") + for limit in ( + limits.max_children_total, + limits.max_children_per_batch, + limits.result_truncation_limit, + ): + if limit is not None and limit < 0: + raise ValueError("Child count and result limits must be nonnegative") + self.config = config + self.limits = limits + self.runner_options = runner_options + self.on_subcall_start = on_subcall_start + self.on_subcall_complete = on_subcall_complete + self.child_traces: list[ChildTrace] = [] + self._jobs: dict[str, _Job] = {} + self._lock = threading.RLock() + self._lifecycle_lock = threading.RLock() + self._callback_queue = SimpleQueue() + self._callback_thread: threading.Thread | None = None + self._accepting = True + self._closed = False + self._cleanup_failed = False + + def check_ready(self): + if self._cleanup_failed: + raise ChildCleanupError( + "Managed child cleanup failed; environment cannot be reused" + ) + if self._closed: + raise RuntimeError("Managed child controller is closed") + + @staticmethod + def _run_callbacks(notifications): + while (notification := notifications.get()) is not None: + callback, args = notification + with suppress(Exception): + callback(*args) + + def _notify(self, callback, *args): + if callback is None: + return + with self._lock: + if self._callback_thread is None: + try: + callback_thread = threading.Thread( + target=self._run_callbacks, + args=(self._callback_queue,), + daemon=True, + ) + callback_thread.start() + except RuntimeError: + return + self._callback_thread = callback_thread + self._callback_queue.put((callback, args)) + + def spawn(self, prompt: str, model: str | None = None) -> str: + """Start a child without waiting; return an episode-scoped handle.""" + return self._spawn(prompt, model, parent=None) + + def _spawn(self, prompt: str, model: str | None, *, parent: _Job | None) -> str: + if not isinstance(prompt, str) or ( + model is not None and not isinstance(model, str) + ): + raise TypeError("prompt and model must be strings") + with self._lock: + self.check_ready() + depth = parent.depth if parent is not None else 0 + if parent is not None and (parent.stop.is_set() or not parent.accepting): + raise RuntimeError("Parent child is finalizing") + if depth >= self.limits.max_depth: + raise RuntimeError("max_depth exceeded") + if not self._accepting: + raise RuntimeError("Episode is finalizing; new children are disabled") + if self.limits.max_children_total is not None: + if len(self._jobs) >= self.limits.max_children_total: + raise RuntimeError("max_children_total exceeded") + if ( + sum(not job.finished.is_set() for job in self._jobs.values()) + >= self.limits.max_batch_workers + ): + raise RuntimeError("Managed child concurrency limit exceeded") + context = multiprocessing.get_context("spawn") + host_channel, worker_channel = context.Pipe() + shutdown_reader, shutdown_writer = context.Pipe(duplex=False) + options = { + "max_depth": self.limits.max_depth, + "max_batch_workers": self.limits.max_batch_workers, + "max_children_per_batch": self.limits.max_children_per_batch, + "result_truncation_limit": self.limits.result_truncation_limit, + **self.runner_options, + } + process = context.Process( + target=_run_child, + args=( + self.config.chat_factory, + worker_channel, + shutdown_reader, + prompt, + model, + depth + 1, + options, + ), + ) + try: + process.start() + except BaseException: + host_channel.close() + shutdown_writer.close() + raise + finally: + worker_channel.close() + shutdown_reader.close() + handle = uuid4().hex + job = _Job( + process, + host_channel, + shutdown_writer, + prompt, + depth + 1, + parent, + ) + self._jobs[handle] = job + job.model = model + self._notify( + self.on_subcall_start, job.depth, model or "default", prompt[:80] + ) + try: + job.thread = threading.Thread( + target=self._watch, args=(job,), daemon=True + ) + job.thread.start() + except BaseException as exc: + job.error = str(exc) + job.stop.set() + self._watch(job) + if job.exitcode is None or job.status == "cleanup_failed": + raise ChildCleanupError( + "Failed to clean up unmonitored worker" + ) from exc + raise + return handle + + def _job(self, handle: str) -> _Job: + with self._lock: + try: + return self._jobs[handle] + except KeyError: + raise KeyError("Unknown or stale child handle") from None + + def status(self, handle: str) -> str: + """Return running, stopping, completed, failed, cancelled, or cleanup_failed.""" + job = self._job(handle) + if not job.finished.is_set(): + return "stopping" if job.stop.is_set() else "running" + return job.status + + def result(self, handle: str) -> str: + """Read a completed result without waiting.""" + job = self._job(handle) + if not job.finished.is_set(): + raise RuntimeError("Child has not finished") + if job.status != "completed": + raise RuntimeError(job.error or f"Child {job.status}") + return job.answer or "" + + def wait(self, handle: str, timeout_s: float = 30.0) -> str: + """Wait for a bounded period; a wait timeout does not cancel the child.""" + if not math.isfinite(timeout_s) or timeout_s < 0: + raise ValueError("timeout_s must be finite and nonnegative") + if not self._job(handle).finished.wait(timeout_s): + raise TimeoutError("Child wait timed out; child is still running") + return self.result(handle) + + def cancel(self, handle: str) -> None: + """Request shutdown, escalate if necessary, and confirm process exit.""" + job = self._job(handle) + job.stop.set() + self._await_cleanup([job]) + + def _signal(self, job: _Job, signum: int): + if job.exitcode is not None: + return + try: + os.kill(job.pid, signum) + except ProcessLookupError: + pass + + def _dispatch(self, job: _Job, method: str, args: tuple): + try: + if method == "spawn": + if job.stop.is_set(): + raise RuntimeError("Parent child is stopping") + value = self._spawn(*args, parent=job) + elif method == "status": + value = self.status(*args) + elif method == "result": + value = self.result(*args) + elif method == "stop": + self._stop_descendants(job) + value = None + else: + raise ValueError("Unknown child request") + response = (True, value) + except Exception as exc: + response = (False, str(exc)) + with suppress(BrokenPipeError, ConnectionResetError): + job.channel.send(response) + + def _stop_descendants(self, job: _Job): + with self._lock: + job.accepting = False + descendants = [ + child for child in self._jobs.values() if child.parent is job + ] + for child in descendants: + child.stop.set() + self._await_cleanup(descendants) + + def _watch(self, job: _Job): + shutdown_at = None + term_at = None + kill_at = None + ready_at = None + terminal_message = None + channel_ended = job.channel.closed + try: + while True: + now = time.monotonic() + if not channel_ended: + try: + if job.channel.poll(0.01): + message = job.channel.recv() + if message[0] == "ready": + ready_at = now + elif message[0] == "request": + self._dispatch(job, message[1], message[2]) + else: + terminal_message = message + except (EOFError, OSError): + channel_ended = True + if channel_ended and terminal_message is None and not job.stop.is_set(): + job.error = "Worker channel closed before completion" + job.stop.set() + timeout = self.limits.per_child_timeout_s + if ( + ready_at is not None + and timeout is not None + and now - ready_at >= timeout + ): + job.error = "Child timeout exceeded" + job.stop.set() + if ready_at is None and now - job.started >= 30: + job.error = "Child startup timeout exceeded" + job.stop.set() + if job.stop.is_set() and shutdown_at is None: + job.status = "stopping" + job.shutdown.close() + shutdown_at = now + if ( + shutdown_at is not None + and term_at is None + and now - shutdown_at >= self.config.shutdown_grace_s + ): + self._signal(job, signal.SIGTERM) + term_at = now + if ( + term_at is not None + and kill_at is None + and now - term_at >= self.config.terminate_grace_s + ): + self._signal(job, signal.SIGKILL) + kill_at = now + if job.exitcode is None: + job.process.join(timeout=0.01 if channel_ended else 0) + job.exitcode = job.process.exitcode + if job.exitcode is not None: + if ( + terminal_message is not None + or channel_ended + or not job.channel.poll() + ): + break + if kill_at is not None and now - kill_at >= self.config.kill_timeout_s: + raise ChildCleanupError("Worker exit not confirmed after SIGKILL") + if job.stop.is_set(): + job.status = "failed" if job.error else "cancelled" + elif ( + terminal_message + and terminal_message[0] == "result" + and job.exitcode == 0 + ): + job.answer = terminal_message[1] + job.status = "completed" + else: + job.status = "failed" + job.error = ( + terminal_message[1] + if terminal_message and terminal_message[0] == "error" + else f"Worker exited with code {job.exitcode}" + ) + except Exception as exc: + job.status = "cleanup_failed" + job.error = str(exc) + self._cleanup_failed = True + finally: + try: + self._stop_descendants(job) + except ChildCleanupError as exc: + job.status = "cleanup_failed" + job.error = str(exc) + self._cleanup_failed = True + if job.exitcode is not None: + job.channel.close() + job.shutdown.close() + job.process.close() + if not job.reported: + self.child_traces.append( + ChildTrace( + job.depth, + time.monotonic() - job.started, + job.prompt[:80], + job.answer[:80] if job.answer else None, + job.error, + ) + ) + self._notify( + self.on_subcall_complete, + job.depth, + job.model or "default", + time.monotonic() - job.started, + job.error, + ) + job.reported = True + job.finished.set() + + def _await_cleanup(self, jobs: list[_Job]): + deadline = ( + time.monotonic() + + self.config.shutdown_grace_s + + self.config.terminate_grace_s + + self.config.kill_timeout_s + + 1 + ) + for job in jobs: + if job.status == "cleanup_failed" and job.finished.is_set(): + job.finished.clear() + job.thread = threading.Thread( + target=self._watch, args=(job,), daemon=True + ) + job.thread.start() + if ( + not job.finished.wait(max(0, deadline - time.monotonic())) + or job.exitcode is None + or job.status == "cleanup_failed" + ): + self._cleanup_failed = True + raise ChildCleanupError("Managed worker exit could not be confirmed") + + def stop(self): + """Stop every registered worker, including nested workers, before returning.""" + with self._lifecycle_lock: + with self._lock: + self._accepting = False + jobs = list(self._jobs.values()) + for job in jobs: + job.stop.set() + self._await_cleanup(jobs) + self._cleanup_failed = False + + def reset(self): + """Drain the old episode before invalidating its handles and budget.""" + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("Managed child controller is closed") + self.stop() + with self._lock: + self._jobs.clear() + self.child_traces.clear() + self._accepting = True + + def close(self): + """Close permanently; a failed cleanup may be retried.""" + with self._lifecycle_lock: + self.stop() + if not self._closed: + self._closed = True + self._callback_queue.put(None) + + +class _ChildCancelled(BaseException): + pass + + +class _RemoteChildren: + def __init__(self, channel: Connection, shutdown: Connection): + self.channel = channel + self.shutdown = shutdown + self.lock = threading.Lock() + + def check_ready(self): + if self.shutdown.poll(): + raise _ChildCancelled() + + def _request(self, method: str, *args): + with self.lock: + self.check_ready() + self.channel.send(("request", method, args)) + while not self.channel.poll(0.01): + self.check_ready() + success, value = self.channel.recv() + if not success: + raise RuntimeError(value) + return value + + def spawn(self, prompt: str, model: str | None = None): + return self._request("spawn", prompt, model) + + def status(self, handle: str): + return self._request("status", handle) + + def result(self, handle: str): + return self._request("result", handle) + + def wait(self, handle: str, timeout_s: float = 30.0): + if not math.isfinite(timeout_s) or timeout_s < 0: + raise ValueError("timeout_s must be finite and nonnegative") + deadline = time.monotonic() + timeout_s + while self.status(handle) in ("running", "stopping"): + if time.monotonic() >= deadline: + raise TimeoutError("Child wait timed out; child is still running") + self.check_ready() + time.sleep(0.01) + return self.result(handle) + + def stop(self): + self._request("stop") + + def reset(self): + self.check_ready() + + def close(self): + self.stop() + + +class ManagedChildBackend(DirectLMBackend): + """Keep existing blocking query APIs and add managed child handles.""" + + def __init__(self, chat_fn, *, children, depth=0, limits=None): + super().__init__(chat_fn, depth=depth, limits=limits) + self.children = children + if isinstance(children, ManagedChildren): + self.child_traces = children.child_traces + + def recursive_query(self, prompt: str, model: str | None = None) -> str: + if self.depth + 1 >= self.max_depth: + return self.query(prompt, model) + try: + handle = self.children.spawn(prompt, model) + while self.children.status(handle) in ("running", "stopping"): + try: + return self.children.wait(handle, timeout_s=1) + except TimeoutError: + pass + return self.children.result(handle) + except RuntimeError as exc: + return f"Error: {exc}" + + def recursive_query_batched( + self, prompts: list[str], model: str | None = None + ) -> list[str]: + from concurrent.futures import ThreadPoolExecutor + + if self.limits.max_children_per_batch is not None: + prompts = prompts[: self.limits.max_children_per_batch] + if not prompts: + return [] + with ThreadPoolExecutor(max_workers=self.limits.max_batch_workers) as executor: + return list( + executor.map( + lambda prompt: self.recursive_query(prompt, model), prompts + ) + ) + + +def _run_child(chat_factory, channel, shutdown, prompt, model, depth, options): + children = _RemoteChildren(channel, shutdown) + try: + channel.send(("ready",)) + chat = chat_factory() + + def checked_chat(messages, model=None): + children.check_ready() + try: + result = chat(messages, model) + except TypeError: + result = chat(messages) + children.check_ready() + return result + + limits = BackendLimits( + **{ + key: options[key] + for key in ( + "max_depth", + "max_batch_workers", + "max_children_per_batch", + "result_truncation_limit", + ) + } + ) + + def backend_factory(chat_fn, **kwargs): + return ManagedChildBackend( + chat_fn, children=children, depth=kwargs["depth"], limits=limits + ) + + if depth >= limits.max_depth: + answer = DirectLMBackend(checked_chat, limits=limits).query(prompt, model) + else: + from .runner import LocalRLMRunner + + result = LocalRLMRunner( + checked_chat, + depth=depth, + backend_factory=backend_factory, + **{ + key: options[key] + for key in ( + "max_depth", + "max_batch_workers", + "system_prompt", + "max_iterations", + "env_max_iterations_multiplier", + ) + }, + ).run(prompt, prompt, model=model) + answer = result.final_answer or "" + limit = limits.result_truncation_limit + if limit is not None: + answer = answer[:limit] + with children.lock: + channel.send(("result", answer)) + except _ChildCancelled: + pass + except Exception as exc: + with children.lock: + channel.send(("error", f"{type(exc).__name__}: {exc}")) + finally: + channel.close() + shutdown.close() diff --git a/envs/repl_env/prompts.py b/envs/repl_env/prompts.py index b6f6331c09..a1260efbfb 100644 --- a/envs/repl_env/prompts.py +++ b/envs/repl_env/prompts.py @@ -39,7 +39,7 @@ class QueryMetadata: 1. A `context` variable that contains extremely important information about your query. You should check the content of the `context` variable to understand what you are working with. Make sure you look through it sufficiently as you answer your query. 2. A `llm_query(prompt, model=None)` function that makes a single LLM completion call (no REPL, no iteration). Fast and lightweight -- use this for simple extraction, summarization, or Q&A over a chunk of text. The sub-LLM can handle around 500K chars. 3. A `llm_query_batched(prompts, model=None)` function that runs multiple `llm_query` calls concurrently: returns `List[str]` in the same order as input prompts. Much faster than sequential `llm_query` calls for independent queries. -4. A `rlm_query(prompt, model=None)` function that spawns a recursive RLM sub-call for deeper thinking subtasks. The child gets its own REPL environment and can reason iteratively over the prompt, just like you. Use this when a subtask requires multi-step reasoning, code execution, or its own iterative problem-solving. Falls back to `llm_query` if recursion is not available. +4. A `rlm_query(prompt, model=None, blocking=True)` function that runs a recursive RLM sub-call for deeper thinking subtasks and waits for its final answer by default. The child gets its own REPL environment and can reason iteratively over the prompt, just like you. Use this when a subtask requires multi-step reasoning, code execution, or its own iterative problem-solving. Blocking queries fall back to `llm_query` if recursion is not available. The keyword-only `blocking=False` option requires managed children and returns a handle instead of an answer. 5. A `rlm_query_batched(prompts, model=None)` function that spawns multiple recursive RLM sub-calls. Each prompt gets its own child RLM. Falls back to `llm_query_batched` if recursion is not available. 6. A `SHOW_VARS()` function that returns all variables you have created in the REPL. Use this to check what variables exist before using FINAL_VAR. 7. The ability to use `print()` statements to view the output of your REPL code and continue your reasoning. diff --git a/envs/repl_env/recursive_controller.py b/envs/repl_env/recursive_controller.py index bc1fc8b8e4..897d1e6e69 100644 --- a/envs/repl_env/recursive_controller.py +++ b/envs/repl_env/recursive_controller.py @@ -18,6 +18,7 @@ from dataclasses import dataclass from typing import Callable +from .managed_children import ManagedChildBackend, ManagedChildConfig, ManagedChildren from .prompts import RLM_SYSTEM_PROMPT from .recursive_backends import BackendLimits, DirectLMBackend, LocalChildRLMBackend @@ -33,8 +34,15 @@ class RecursiveController: rlm_batch_fn: Callable[[list[str], str | None], list[str]] | None backend: object + @property + def children(self): + if isinstance(self.backend, ManagedChildBackend): + return self.backend.children + return None + def close(self) -> None: - pass + if self.children is not None: + self.children.close() def create_server_recursive_controller( @@ -49,6 +57,9 @@ def create_server_recursive_controller( result_truncation_limit: int | None = None, per_child_timeout_s: float | None = None, env_max_iterations_multiplier: int = 5, + managed_children: ManagedChildConfig | None = None, + on_subcall_start: Callable[[int, str, str], None] | None = None, + on_subcall_complete: Callable[[int, str, float, str | None], None] | None = None, ) -> RecursiveController: limits = BackendLimits( max_depth=max_depth, @@ -58,6 +69,24 @@ def create_server_recursive_controller( result_truncation_limit=result_truncation_limit, per_child_timeout_s=per_child_timeout_s, ) + if managed_children is not None: + children = ManagedChildren( + managed_children, + limits, + on_subcall_start=on_subcall_start, + on_subcall_complete=on_subcall_complete, + system_prompt=system_prompt, + max_iterations=max_iterations, + env_max_iterations_multiplier=env_max_iterations_multiplier, + ) + backend = ManagedChildBackend(chat_fn, children=children, limits=limits) + return RecursiveController( + llm_query_fn=backend.query, + llm_batch_fn=backend.query_batched, + rlm_query_fn=backend.recursive_query, + rlm_batch_fn=backend.recursive_query_batched, + backend=backend, + ) if max_depth > 1: from .runner import LocalRLMRunner @@ -69,6 +98,8 @@ def create_server_recursive_controller( env_max_iterations_multiplier=env_max_iterations_multiplier, depth=0, limits=limits, + on_subcall_start=on_subcall_start, + on_subcall_complete=on_subcall_complete, ) return RecursiveController( llm_query_fn=backend.query, diff --git a/envs/repl_env/runner.py b/envs/repl_env/runner.py index bf37d0e7b0..e6820797ab 100644 --- a/envs/repl_env/runner.py +++ b/envs/repl_env/runner.py @@ -15,12 +15,16 @@ from __future__ import annotations -import re import time from dataclasses import dataclass from typing import Callable from .local import LocalREPLEnv +from .managed_children import ( + MANAGED_CHILD_TOOLS_PROMPT, + ManagedChildBackend, + ManagedChildConfig, +) from .prompts import ( build_rlm_system_prompt, build_user_prompt, @@ -30,6 +34,10 @@ RLM_SYSTEM_PROMPT, ) from .recursive_backends import BackendLimits, LocalChildRLMBackend, RecursiveBackend +from .recursive_controller import ( + create_server_recursive_controller, + RecursiveController, +) ChatFn = Callable[..., str] @@ -66,6 +74,7 @@ def __init__( on_subcall_complete: Callable[[int, str, float, str | None], None] | None = None, verbose: bool = False, + managed_children: ManagedChildConfig | None = None, ) -> None: self.llm_chat_fn = llm_chat_fn self.system_prompt = system_prompt @@ -82,6 +91,9 @@ def __init__( self.on_subcall_start = on_subcall_start self.on_subcall_complete = on_subcall_complete self.verbose = verbose + self.managed_children = managed_children + if managed_children is not None and backend_factory is not None: + raise ValueError("Choose managed_children or backend_factory, not both") def _default_backend_factory( self, llm_chat_fn: ChatFn, **kwargs @@ -106,14 +118,27 @@ def _default_backend_factory( on_subcall_complete=self.on_subcall_complete, ) - def run( - self, - context: str, - task_prompt: str, - *, - model: str | None = None, - timeout_s: float | None = None, - ) -> RLMRunResult: + def _create_controller(self) -> RecursiveController: + if self.managed_children is not None: + if self.depth != 0: + raise ValueError( + "Managed child controllers must be created at the root" + ) + return create_server_recursive_controller( + self.llm_chat_fn, + max_depth=self.max_depth, + max_iterations=self.max_iterations, + system_prompt=self.system_prompt, + max_batch_workers=self.max_batch_workers, + max_children_total=self.max_children_total, + max_children_per_batch=self.max_children_per_batch, + result_truncation_limit=self.result_truncation_limit, + per_child_timeout_s=self.per_child_timeout_s, + env_max_iterations_multiplier=self.env_max_iterations_multiplier, + managed_children=self.managed_children, + on_subcall_start=self.on_subcall_start, + on_subcall_complete=self.on_subcall_complete, + ) backend = self.backend_factory( self.llm_chat_fn, system_prompt=self.system_prompt, @@ -122,12 +147,25 @@ def run( depth=self.depth, env_max_iterations_multiplier=self.env_max_iterations_multiplier, ) - with LocalREPLEnv( + return RecursiveController( llm_query_fn=backend.query, llm_batch_fn=backend.query_batched, - subcall_fn=backend.recursive_query, - subcall_batch_fn=backend.recursive_query_batched, - ) as env: + rlm_query_fn=backend.recursive_query, + rlm_batch_fn=backend.recursive_query_batched, + backend=backend, + ) + + def run( + self, + context: str, + task_prompt: str, + *, + model: str | None = None, + timeout_s: float | None = None, + ) -> RLMRunResult: + controller = self._create_controller() + backend = controller.backend + with LocalREPLEnv(recursive_controller=controller) as env: result = env.reset( context=context, task_prompt=task_prompt, @@ -141,7 +179,10 @@ def run( context_total_length=obs.context_length, context_type="str", ) - messages = build_rlm_system_prompt(self.system_prompt, query_metadata) + system_prompt = self.system_prompt + if isinstance(backend, ManagedChildBackend): + system_prompt += "\n" + MANAGED_CHILD_TOOLS_PROMPT + messages = build_rlm_system_prompt(system_prompt, query_metadata) messages.append(build_user_prompt(root_prompt=task_prompt, iteration=0)) run_start = time.perf_counter() @@ -248,7 +289,9 @@ def _default_answer( idx = response.find("FINAL(") if idx != -1: depth, start = 0, idx + len("FINAL") - for i, ch in enumerate(response[idx + len("FINAL"):], start=idx + len("FINAL")): + for i, ch in enumerate( + response[idx + len("FINAL") :], start=idx + len("FINAL") + ): if ch == "(": depth += 1 elif ch == ")": diff --git a/envs/repl_env/server/repl_environment.py b/envs/repl_env/server/repl_environment.py index 8ba645a44b..d242042bf2 100644 --- a/envs/repl_env/server/repl_environment.py +++ b/envs/repl_env/server/repl_environment.py @@ -43,16 +43,25 @@ from models import CodeBlockResult, REPLAction, REPLObservation, REPLState try: - from ..recursive_controller import create_server_recursive_controller + from ..recursive_controller import ( + create_server_recursive_controller, + RecursiveController, + ) from ..rubrics import REPLRubric from .python_executor import PythonExecutor except ImportError: try: - from repl_env.recursive_controller import create_server_recursive_controller + from repl_env.recursive_controller import ( + create_server_recursive_controller, + RecursiveController, + ) from repl_env.rubrics import REPLRubric from .python_executor import PythonExecutor except ImportError: - from recursive_controller import create_server_recursive_controller + from recursive_controller import ( + create_server_recursive_controller, + RecursiveController, + ) from rubrics import REPLRubric from .python_executor import PythonExecutor @@ -105,6 +114,7 @@ def __init__( ] = None, rlm_max_depth: int = 1, rlm_max_iterations: int | None = None, + recursive_controller: RecursiveController | None = None, ): """Initialize the REPL environment. @@ -121,6 +131,10 @@ def __init__( subcall_batch_fn: Optional function for recursive rlm_query_batched() support rlm_max_depth: Max recursion depth for server-backed rlm_query() rlm_max_iterations: Max iterations for recursive child runners + recursive_controller (`RecursiveController`, *optional*): + Preconfigured controller owned by this environment. Enables + managed helpers when its backend supports them; cannot be + combined with raw query functions. """ self.initial_context = context or os.environ.get("REPL_CONTEXT", "") self.initial_task_prompt = task_prompt or os.environ.get("REPL_TASK_PROMPT", "") @@ -145,6 +159,27 @@ def __init__( self._runtime_controller = None self._runtime_controller_chat_fn: Optional[Callable[..., str]] = None self._current_llm_model: Optional[str] = None + self._external_controller = recursive_controller + if recursive_controller is not None: + if any( + fn is not None + for fn in (llm_query_fn, llm_batch_fn, subcall_fn, subcall_batch_fn) + ): + raise ValueError( + "Use recursive_controller or raw query functions, not both" + ) + self.llm_query_fn = recursive_controller.llm_query_fn + self.llm_batch_fn = recursive_controller.llm_batch_fn + self.subcall_fn = recursive_controller.rlm_query_fn + self.subcall_batch_fn = recursive_controller.rlm_batch_fn + + @property + def _managed_children(self): + if self._external_controller is not None: + return self._external_controller.children + if isinstance(self._runtime_controller, RecursiveController): + return self._runtime_controller.children + return None @staticmethod def _resolve_model(llm_model: Optional[str]) -> str: @@ -240,6 +275,16 @@ def reset( Returns: Initial REPLObservation with environment ready message """ + if self._external_controller is not None and ( + hf_token is not None + or kwargs.get("rlm_max_depth") is not None + or kwargs.get("rlm_max_iterations") is not None + ): + raise ValueError( + "Configure credentials and recursion limits on the supplied controller" + ) + if self._managed_children is not None: + self._managed_children.reset() effective_context = context or self.initial_context effective_task_prompt = task_prompt or self.initial_task_prompt @@ -272,9 +317,7 @@ def reset( # reset() are treated as equal and don't trigger a redundant rebuild. resolved_model = self._resolve_model(llm_model) has_runtime_llm = self._runtime_controller is not None - model_changed = ( - has_runtime_llm and resolved_model != self._current_llm_model - ) + model_changed = has_runtime_llm and resolved_model != self._current_llm_model token_provided = hf_token is not None if not self.llm_query_fn or model_changed or token_provided: effective_token = ( @@ -339,7 +382,15 @@ def _call_batched_query( except TypeError: return self.llm_batch_fn(prompts) # type: ignore[misc] - def _call_recursive_query(prompt: str, model: str | None = None) -> str: + def _call_recursive_query( + prompt: str, model: str | None = None, *, blocking: bool = True + ) -> str: + if not blocking: + if self._managed_children is None: + raise RuntimeError( + "rlm_query(blocking=False) requires managed children" + ) + return self._managed_children.spawn(prompt, model) if self.subcall_fn is None: return _call_single_query(prompt, model) return self.subcall_fn(prompt, model) @@ -367,6 +418,14 @@ def _call_recursive_batched( if self.llm_batch_fn or self.subcall_batch_fn: self._executor.inject_function("rlm_query_batched", _call_recursive_batched) + if self._managed_children is not None: + for name, function in ( + ("rlm_status", self._managed_children.status), + ("rlm_result", self._managed_children.result), + ("rlm_wait", self._managed_children.wait), + ): + self._executor.inject_function(name, function) + # Inject FINAL helper function so both FINAL(x) and print(f'FINAL({x})') work # Returns the FINAL pattern as a string so it appears in stdout for detection def final_helper(value): @@ -402,6 +461,12 @@ def show_vars_helper(): # Build initial message message_parts = ["REPL environment initialized."] + if self._managed_children is not None: + message_parts.append( + "Managed children: rlm_query(prompt, blocking=False) returns a handle; " + "rlm_query(prompt) waits for an answer. Use rlm_status(handle), " + "rlm_result(handle), rlm_wait(handle, timeout_s=30)." + ) if effective_context: message_parts.append( f"Context loaded ({len(effective_context)} chars). Use 'context' variable to access it." @@ -453,6 +518,8 @@ def step( Returns: REPLObservation with execution results """ + if self._managed_children is not None: + self._managed_children.check_ready() if self._state is None or self._executor is None: raise RuntimeError("Environment not initialized. Call reset() first.") @@ -493,6 +560,8 @@ def step( if done: self._state.final_answer = final_answer + if self._managed_children is not None: + self._managed_children.stop() obs = REPLObservation( result=CodeBlockResult( @@ -541,7 +610,9 @@ def _extract_final_answer(self, stdout: str) -> Optional[str]: idx = stdout.find("FINAL(") if idx != -1: depth, start = 0, idx + len("FINAL") - for i, ch in enumerate(stdout[idx + len("FINAL"):], start=idx + len("FINAL")): + for i, ch in enumerate( + stdout[idx + len("FINAL") :], start=idx + len("FINAL") + ): if ch == "(": depth += 1 elif ch == ")": @@ -576,6 +647,8 @@ def _create_final_observation(self, success: bool, message: str) -> REPLObservat Returns: Final REPLObservation with done=True (reward set by rubric) """ + if self._managed_children is not None: + self._managed_children.stop() return REPLObservation( result=CodeBlockResult( stdout=message, @@ -616,6 +689,8 @@ def state(self) -> REPLState: def close(self) -> None: """Cleanup resources.""" + if self._external_controller is not None: + self._external_controller.close() if self._runtime_controller is not None: self._runtime_controller.close() self._runtime_controller = None diff --git a/tests/envs/test_repl_managed_children.py b/tests/envs/test_repl_managed_children.py new file mode 100644 index 0000000000..27792e1251 --- /dev/null +++ b/tests/envs/test_repl_managed_children.py @@ -0,0 +1,978 @@ +"""Unit contracts and real-process cleanup coverage for managed child queries.""" + +import functools +import os +import signal +import sys +import threading +import time +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +pytest.importorskip("smolagents") + +import repl_env.runner as runner_module +from repl_env import LocalREPLEnv, LocalRLMRunner +from repl_env.managed_children import ChildCleanupError, ManagedChildConfig +from repl_env.recursive_controller import create_server_recursive_controller + + +pytestmark = pytest.mark.skipif( + sys.platform != "linux", reason="POSIX worker supervision" +) + + +def fake_chat_factory(directory): + def chat(messages, model=None): + Path(directory, f"{os.getpid()}.started").touch() + prompt = messages[-1]["content"] + if "blocked" in prompt: + if "ignore-term" in prompt: + signal.signal( + signal.SIGTERM, + lambda *_: Path(directory, f"{os.getpid()}.term").touch(), + ) + Path(directory, f"{os.getpid()}.ready").touch() + while not Path(directory, "release").exists(): + time.sleep(0.01) + if "raise-error" in prompt: + raise ValueError("scripted model failure") + if "launch-nested" in prompt: + return "```repl\nchild = rlm_query('blocked ignore-term', blocking=False); llm_query('blocked ignore-term')\n```" + if "nested" in prompt: + return "```repl\nprint(FINAL(rlm_query('answer')))\n```" + return "```repl\nprint(FINAL(42))\n```" + + return chat + + +def make_controller(tmp_path, **kwargs): + config = ManagedChildConfig( + chat_factory=functools.partial(fake_chat_factory, str(tmp_path)), + shutdown_grace_s=0.1, + terminate_grace_s=0.15, + kill_timeout_s=2, + ) + return create_server_recursive_controller( + fake_chat_factory(str(tmp_path)), + managed_children=config, + max_depth=kwargs.pop("max_depth", 3), + max_iterations=2, + **kwargs, + ) + + +def wait_ready(tmp_path, count=1, *, manager=None, phase="ready"): + deadline = time.monotonic() + 30 * count + details = [] + while time.monotonic() < deadline: + files = list(tmp_path.glob(f"*.{phase}")) + if len(files) >= count: + return [int(path.stem) for path in files] + if manager is not None: + with manager._lock: + entries = list(manager._jobs.items()) + jobs = [job for _, job in entries] + details = [ + (job.pid, manager.status(handle), job.exitcode, job.error) + for handle, job in entries + ] + if any( + job.finished.is_set() and job.status in ("failed", "cleanup_failed") + for job in jobs + ): + pytest.fail(f"Worker failed before {phase}: {details}") + time.sleep(0.01) + pytest.fail(f"Expected {count} workers at {phase}; found {len(files)}: {details}") + + +@pytest.fixture +def fault_controller(tmp_path): + controller = make_controller(tmp_path) + yield controller + try: + controller.close() + finally: + for job in controller.children._jobs.values(): + if job.exitcode is None: + job.process.kill() + job.process.join(5) + job.exitcode = job.process.exitcode + assert job.exitcode is not None + if job.thread is not None and job.thread.ident is not None: + job.thread.join(5) + job.channel.close() + job.shutdown.close() + job.process.close() + + +def assert_reaped(pids): + for pid in pids: + with pytest.raises(ProcessLookupError): + os.kill(pid, 0) + with pytest.raises(ChildProcessError): + os.waitpid(pid, os.WNOHANG) + + +def start_call(function, name): + done = threading.Event() + errors = [] + + def invoke(): + try: + function() + except BaseException as exc: + errors.append(exc) + finally: + done.set() + + thread = threading.Thread(target=invoke, name=name) + thread.start() + return thread, done, errors + + +@pytest.mark.parametrize( + "model_argument, expected_model", + [ + ("", None), + (", 'fake-model'", "fake-model"), + (", model='fake-model'", "fake-model"), + ], +) +def test_nonblocking_query_routes_model_arguments( + tmp_path, monkeypatch, model_argument, expected_model +): + controller = make_controller(tmp_path) + spawn = Mock(return_value="child-handle") + wait = Mock(return_value="42") + monkeypatch.setattr(controller.children, "spawn", spawn) + monkeypatch.setattr(controller.children, "wait", wait) + monkeypatch.setattr(controller.children, "status", Mock(return_value="completed")) + monkeypatch.setattr(controller.children, "result", Mock(return_value="42")) + with LocalREPLEnv(recursive_controller=controller) as env: + reset = env.reset() + assert "blocking=False" in reset.observation.result.stdout + assert "rlm_spawn" not in env.list_variables() + result = env.execute( + f"child = rlm_query('answer'{model_argument}, blocking=False); parent_work = 6 * 7" + ) + assert result.observation.result.success + spawn.assert_called_once_with("answer", expected_model) + assert ( + env.execute("print(parent_work)").observation.result.stdout.strip() == "42" + ) + result = env.execute("print(rlm_wait(child, timeout_s=15))") + assert result.observation.result.stdout.strip() == "42" + wait.assert_called_once_with("child-handle", timeout_s=15) + assert ( + env.execute("print(rlm_status(child))").observation.result.stdout.strip() + == "completed" + ) + assert ( + env.execute("print(rlm_result(child))").observation.result.stdout.strip() + == "42" + ) + assert not controller.children._jobs + + +def test_nonblocking_query_runs_supervised_worker(tmp_path): + controller = make_controller(tmp_path) + with LocalREPLEnv(recursive_controller=controller) as env: + env.reset() + result = env.execute( + "child = rlm_query('answer', model='fake-model', blocking=False); parent_work = 6 * 7" + ) + assert result.observation.result.success + result = env.execute("print(parent_work, rlm_wait(child, timeout_s=15))") + assert result.observation.result.success + assert result.observation.result.stdout.strip() == "42 42" + result = env.execute("print(rlm_status(child), rlm_result(child))") + assert result.observation.result.stdout.strip() == "completed 42" + job = next(iter(controller.children._jobs.values())) + assert job.model == "fake-model" + assert job.pid != os.getpid() + assert job.finished.is_set() + assert_reaped([job.pid]) + + +@pytest.mark.parametrize("blocking_argument", ["", ", blocking=True"]) +def test_blocking_query_returns_answer_by_default( + tmp_path, monkeypatch, blocking_argument +): + controller = make_controller(tmp_path) + query = Mock(return_value="42") + monkeypatch.setattr(controller, "rlm_query_fn", query) + with LocalREPLEnv(recursive_controller=controller) as env: + env.reset() + result = env.execute( + f"print(rlm_query('answer', 'fake-model'{blocking_argument}))" + ) + assert result.observation.result.success + assert result.observation.result.stdout.strip() == "42" + query.assert_called_once_with("answer", "fake-model") + assert not controller.children._jobs + + +@pytest.mark.parametrize("blocking_argument", ["", ", blocking=True"]) +@pytest.mark.parametrize("recursive", [False, True]) +def test_unmanaged_blocking_query_keeps_existing_call_signature( + blocking_argument, recursive +): + calls = [] + + def query(prompt, model=None): + calls.append((prompt, model)) + return "answer" + + with LocalREPLEnv( + llm_query_fn=query, subcall_fn=query if recursive else None + ) as env: + env.reset() + result = env.execute( + f"print(rlm_query('question', 'existing-model'{blocking_argument}))" + ) + assert result.observation.result.success + assert result.observation.result.stdout.strip() == "answer" + assert calls == [("question", "existing-model")] + + +@pytest.mark.parametrize("recursive", [False, True]) +def test_nonblocking_query_requires_managed_mode(recursive): + calls = [] + + def query(prompt, model=None): + calls.append((prompt, model)) + return "answer" + + with LocalREPLEnv( + llm_query_fn=query, subcall_fn=query if recursive else None + ) as env: + env.reset() + result = env.execute("rlm_query('question', blocking=False)") + assert not result.observation.result.success + assert "requires managed children" in result.observation.result.stderr + assert calls == [] + + +def test_blocking_option_is_keyword_only(): + calls = [] + with LocalREPLEnv(llm_query_fn=lambda *args: calls.append(args)) as env: + env.reset() + result = env.execute("rlm_query('question', None, False)") + assert not result.observation.result.success + assert "positional" in result.observation.result.stderr + assert calls == [] + + +@pytest.mark.parametrize("limit", ["max_children_total", "max_batch_workers"]) +def test_nonblocking_query_respects_supervisor_limits(tmp_path, monkeypatch, limit): + controller = make_controller( + tmp_path, **{limit: 0 if limit == "max_children_total" else 1} + ) + launcher = Mock(side_effect=AssertionError("Unexpected worker launch")) + monkeypatch.setattr( + "repl_env.managed_children.multiprocessing.get_context", launcher + ) + with LocalREPLEnv(recursive_controller=controller) as env: + env.reset() + if limit == "max_batch_workers": + controller.children._jobs["running-child"] = SimpleNamespace( + finished=threading.Event() + ) + admitted = len(controller.children._jobs) + try: + result = env.execute("rlm_query('answer', blocking=False)") + assert not result.observation.result.success + expected_error = ( + "max_children_total" if limit == "max_children_total" else "concurrency" + ) + assert expected_error in result.observation.result.stderr + assert len(controller.children._jobs) == admitted + launcher.assert_not_called() + finally: + controller.children._jobs.clear() + + +def test_nonblocking_query_at_depth_limit_still_uses_worker(tmp_path): + controller = make_controller(tmp_path, max_depth=1) + with LocalREPLEnv(recursive_controller=controller) as env: + env.reset() + result = env.execute("child = rlm_query('blocked ignore-term', blocking=False)") + assert result.observation.result.success + pids = wait_ready(tmp_path, manager=controller.children) + handle = next(iter(controller.children._jobs)) + with pytest.raises(TimeoutError): + controller.children.wait(handle, timeout_s=0.01) + assert controller.children.status(handle) == "running" + with pytest.raises(RuntimeError, match="not finished"): + controller.children.result(handle) + controller.children.cancel(handle) + with pytest.raises(RuntimeError, match="cancelled"): + controller.children.result(handle) + env.close() + assert_reaped(pids) + + +@pytest.mark.parametrize( + "ignore_term, exitcode", [(False, -signal.SIGTERM), (True, -signal.SIGKILL)] +) +def test_reset_stops_blocked_call_and_invalidates_handles( + tmp_path, ignore_term, exitcode +): + controller = make_controller(tmp_path) + with LocalREPLEnv(recursive_controller=controller) as env: + env.reset(context="old") + prompt = "blocked ignore-term" if ignore_term else "blocked" + handle = controller.children.spawn(prompt) + pids = wait_ready(tmp_path) + job = controller.children._jobs[handle] + env.reset(context="new") + assert job.exitcode == exitcode + with pytest.raises(ValueError, match="closed"): + _ = job.process.sentinel + assert_reaped(pids) + with pytest.raises(KeyError, match="handle"): + controller.children.status(handle) + assert env.state().context == "new" + + +def test_close_stops_all_workers_and_is_idempotent(tmp_path): + controller = make_controller(tmp_path) + env = LocalREPLEnv(recursive_controller=controller) + env.reset() + for _ in range(2): + controller.children.spawn("blocked ignore-term") + pids = wait_ready(tmp_path, 2) + env.close() + assert_reaped(pids) + env.close() + with pytest.raises(RuntimeError, match="closed"): + env.reset() + + +def test_cooperative_shutdown_exits_before_term(tmp_path): + controller = make_controller(tmp_path) + controller.children.config = replace(controller.children.config, shutdown_grace_s=2) + handle = controller.children.spawn("blocked") + pids = wait_ready(tmp_path) + job = controller.children._jobs[handle] + timer = threading.Timer(0.03, lambda: Path(tmp_path, "release").touch()) + timer.start() + try: + controller.children.cancel(handle) + assert job.exitcode == 0 + assert controller.children.status(handle) == "cancelled" + assert_reaped(pids) + finally: + timer.join() + controller.close() + + +@pytest.mark.parametrize( + "final_code", + ["print(FINAL(42))", "answer.update({'content': '42', 'ready': True})"], +) +def test_parent_finalization_stops_children(tmp_path, final_code): + controller = make_controller(tmp_path) + with LocalREPLEnv(recursive_controller=controller) as env: + env.reset() + controller.children.spawn("blocked ignore-term") + pids = wait_ready(tmp_path) + assert env.execute(final_code).done + assert_reaped(pids) + assert not env.execute( + "rlm_query('answer', blocking=False)" + ).observation.result.success + + +@pytest.mark.parametrize( + "final_code", + ["print(FINAL(42))", "answer.update({'content': '42', 'ready': True})"], +) +def test_final_answer_survives_cleanup_failure(tmp_path, monkeypatch, final_code): + controller = make_controller(tmp_path) + manager = controller.children + env = LocalREPLEnv(recursive_controller=controller) + env.reset(context="original episode") + finished = Mock(spec=threading.Event) + finished.wait.return_value = False + job = SimpleNamespace( + stop=threading.Event(), finished=finished, status="running", exitcode=None + ) + manager._jobs["unconfirmed-child"] = job + launcher = Mock(side_effect=AssertionError("Unexpected worker launch")) + monkeypatch.setattr( + "repl_env.managed_children.multiprocessing.get_context", launcher + ) + try: + with pytest.raises(ChildCleanupError, match="exit could not be confirmed"): + env.execute(final_code) + assert env.state().final_answer == "42" + assert manager._cleanup_failed + assert job.stop.is_set() + finished.wait.assert_called_once() + with pytest.raises(ChildCleanupError, match="cannot be reused"): + env.execute("print('must not execute')") + with pytest.raises(ChildCleanupError, match="cannot be reused"): + manager.spawn("answer") + with pytest.raises(ChildCleanupError, match="exit could not be confirmed"): + env.reset(context="replacement episode") + assert env.state().context == "original episode" + assert env.state().final_answer == "42" + launcher.assert_not_called() + finally: + manager._jobs.clear() + env.close() + + +def test_iteration_limit_stops_children(tmp_path): + controller = make_controller(tmp_path) + with LocalREPLEnv(recursive_controller=controller) as env: + env.reset(max_iterations=1) + controller.children.spawn("blocked ignore-term") + pids = wait_ready(tmp_path) + assert env.execute("pass").done + assert_reaped(pids) + + +def test_worker_failure_is_observable(tmp_path): + controller = make_controller(tmp_path) + try: + failed = controller.children.spawn("raise-error") + with pytest.raises(RuntimeError, match="scripted model failure"): + controller.children.wait(failed, timeout_s=15) + assert controller.children.status(failed) == "failed" + assert_reaped([controller.children._jobs[failed].pid]) + finally: + controller.close() + + +def test_cleanup_failure_prevents_reuse_and_can_be_retried(tmp_path, monkeypatch): + controller = make_controller(tmp_path) + completions = [] + callback_done = threading.Event() + + def on_complete(*args): + completions.append(args) + callback_done.set() + + controller.children.on_subcall_complete = on_complete + env = LocalREPLEnv(recursive_controller=controller) + env.reset(context="old") + controller.children.spawn("blocked ignore-term") + pids = wait_ready(tmp_path) + send_signal = controller.children._signal + monkeypatch.setattr(controller.children, "_signal", lambda *args: None) + try: + with pytest.raises(ChildCleanupError): + env.reset(context="new") + assert env.state().context == "old" + with pytest.raises(ChildCleanupError): + env.execute("print(42)") + with pytest.raises(ChildCleanupError): + controller.children.spawn("answer") + finally: + monkeypatch.setattr(controller.children, "_signal", send_signal) + env.close() + assert_reaped(pids) + assert callback_done.wait(5) + assert len(completions) == 1 + + +def test_nested_recursion_shares_child_budget(tmp_path): + controller = make_controller(tmp_path, max_children_total=1) + try: + handle = controller.children.spawn("nested") + assert "max_children_total" in controller.children.wait(handle, timeout_s=15) + with pytest.raises(RuntimeError, match="max_children_total"): + controller.children.spawn("answer") + finally: + controller.close() + + +@pytest.mark.parametrize("outcome", ["completed", "failed", "cancelled"]) +def test_child_budget_counts_all_admitted_jobs_until_reset( + tmp_path, monkeypatch, outcome +): + controller = make_controller(tmp_path, max_children_total=1) + manager = controller.children + finished = threading.Event() + finished.set() + manager._jobs["prior-child"] = SimpleNamespace( + finished=finished, + stop=threading.Event(), + status=outcome, + exitcode=0, + ) + launcher = Mock(side_effect=RuntimeError("Worker launch reached")) + monkeypatch.setattr( + "repl_env.managed_children.multiprocessing.get_context", launcher + ) + try: + assert manager.status("prior-child") == outcome + with pytest.raises(RuntimeError, match="max_children_total"): + manager.spawn("answer") + launcher.assert_not_called() + manager.reset() + with pytest.raises(KeyError, match="handle"): + manager.status("prior-child") + with pytest.raises(RuntimeError, match="Worker launch reached"): + manager.spawn("answer") + launcher.assert_called_once_with("spawn") + finally: + controller.close() + + +def test_zero_child_budget_never_admits_a_job(tmp_path): + controller = make_controller(tmp_path, max_children_total=0) + try: + with pytest.raises(RuntimeError, match="max_children_total"): + controller.children.spawn("answer") + controller.children.reset() + with pytest.raises(RuntimeError, match="max_children_total"): + controller.children.spawn("answer") + assert not controller.children._jobs + finally: + controller.close() + + +def test_runner_reuses_its_managed_controller(tmp_path, monkeypatch): + created = [] + received = [] + create_controller = runner_module.create_server_recursive_controller + local_env = runner_module.LocalREPLEnv + + def track_controller(*args, **kwargs): + controller = create_controller(*args, **kwargs) + created.append(controller) + return controller + + def track_environment(*args, **kwargs): + received.append(kwargs["recursive_controller"]) + return local_env(*args, **kwargs) + + monkeypatch.setattr( + runner_module, "create_server_recursive_controller", track_controller + ) + monkeypatch.setattr(runner_module, "LocalREPLEnv", track_environment) + runner = LocalRLMRunner( + fake_chat_factory(str(tmp_path)), + max_iterations=2, + max_depth=3, + max_children_total=1, + managed_children=ManagedChildConfig( + chat_factory=functools.partial(fake_chat_factory, str(tmp_path)), + ), + ) + result = runner.run("", "nested") + assert result.final_answer == "42" + assert len(result.child_traces) == 1 + assert len(created) == len(received) == 1 + assert received[0] is created[0] + assert created[0].backend.limits.max_children_total == 1 + assert created[0].children._closed + assert all(job.exitcode is not None for job in created[0].children._jobs.values()) + + +def test_nested_recursion_returns_result_from_supervised_worker(tmp_path): + controller = make_controller(tmp_path) + try: + handle = controller.children.spawn("nested") + wait_ready(tmp_path, 2, manager=controller.children, phase="started") + assert controller.children.wait(handle, timeout_s=15) == "42" + assert len(controller.children._jobs) == 2 + assert all(job.exitcode == 0 for job in controller.children._jobs.values()) + finally: + controller.close() + + +def test_reset_tracks_and_reaps_nested_managed_workers(tmp_path): + controller = make_controller(tmp_path, max_depth=4) + with LocalREPLEnv(recursive_controller=controller) as env: + env.reset() + controller.children.spawn("launch-nested") + pids = wait_ready(tmp_path, 2, manager=controller.children) + assert len(controller.children._jobs) == 2 + env.reset() + assert_reaped(pids) + + +def test_parent_exit_is_not_proof_nested_worker_stopped(tmp_path, monkeypatch): + controller = make_controller(tmp_path, max_depth=4) + send_signal = controller.children._signal + try: + handle = controller.children.spawn("launch-nested") + pids = wait_ready(tmp_path, 2, manager=controller.children) + monkeypatch.setattr( + controller.children, + "_signal", + lambda job, signum: send_signal(job, signum) if job.depth == 1 else None, + ) + with pytest.raises(ChildCleanupError): + controller.children.cancel(handle) + assert controller.children._jobs[handle].exitcode is not None + with pytest.raises(ChildCleanupError): + controller.children.spawn("answer") + finally: + monkeypatch.setattr(controller.children, "_signal", send_signal) + controller.close() + assert_reaped(pids) + + +def test_watcher_start_failure_reaps_started_worker( + tmp_path, monkeypatch, fault_controller +): + manager = fault_controller.children + original_start = threading.Thread.start + + def start(thread): + if thread._target == manager._watch: + wait_ready(tmp_path, manager=manager) + raise RuntimeError("watcher unavailable") + original_start(thread) + + with monkeypatch.context() as patch: + patch.setattr(threading.Thread, "start", start) + with pytest.raises(RuntimeError, match="watcher unavailable"): + manager.spawn("blocked ignore-term") + job = next(iter(manager._jobs.values())) + assert job.exitcode == -signal.SIGKILL + assert job.finished.is_set() + assert_reaped([job.pid]) + manager.close() + + +@pytest.mark.parametrize("operation", ["reset", "close", "wait"]) +def test_broken_channel_still_stops_and_reaps_worker( + tmp_path, fault_controller, operation +): + manager = fault_controller.children + env = LocalREPLEnv(recursive_controller=fault_controller) + env.reset() + handle = manager.spawn("blocked ignore-term") + pids = wait_ready(tmp_path, manager=manager) + job = manager._jobs[handle] + job.channel.close() + operations = {"reset": env.reset, "close": env.close} + if operation == "wait": + assert job.finished.wait(5) + assert manager.status(handle) == "failed" + else: + operations[operation]() + assert job.finished.is_set() + assert job.exitcode == -signal.SIGKILL + assert_reaped(pids) + if operation == "reset": + result = env.execute( + "print(rlm_wait(rlm_query('answer', blocking=False), timeout_s=15))" + ) + assert result.observation.result.success + assert result.observation.result.stdout.strip() == "42" + + +def test_unmanaged_controller_forwards_subcall_callbacks(tmp_path): + events = [] + controller = create_server_recursive_controller( + fake_chat_factory(str(tmp_path)), + max_depth=3, + max_iterations=2, + on_subcall_start=lambda *args: events.append("start"), + on_subcall_complete=lambda *args: events.append("complete"), + ) + try: + assert controller.rlm_query_fn("answer") == "42" + assert events == ["start", "complete"] + finally: + controller.close() + + +def test_callback_dispatcher_start_failure_does_not_leak_worker( + tmp_path, monkeypatch, fault_controller +): + manager = fault_controller.children + manager.on_subcall_start = lambda *args: None + manager.on_subcall_complete = lambda *args: None + original_start = threading.Thread.start + + def start(thread): + if thread._target == manager._run_callbacks: + raise RuntimeError("observer unavailable") + original_start(thread) + + with monkeypatch.context() as patch: + patch.setattr(threading.Thread, "start", start) + handle = manager.spawn("blocked ignore-term") + pids = wait_ready(tmp_path, manager=manager) + manager.close() + assert manager._jobs[handle].exitcode == -signal.SIGKILL + assert manager._callback_thread is None + assert_reaped(pids) + + +def test_concurrency_limit_rejects_without_starting_extra_worker(tmp_path): + controller = make_controller(tmp_path, max_batch_workers=1) + try: + controller.children.spawn("blocked ignore-term") + pids = wait_ready(tmp_path) + with pytest.raises(RuntimeError, match="concurrency"): + controller.children.spawn("blocked ignore-term") + assert len(list(tmp_path.glob("*.ready"))) == 1 + finally: + controller.close() + assert_reaped(pids) + + +def test_child_timeout_interrupts_inflight_call(tmp_path): + controller = make_controller(tmp_path, per_child_timeout_s=0.3) + try: + handle = controller.children.spawn("blocked ignore-term") + pids = wait_ready(tmp_path) + with pytest.raises(RuntimeError, match="timeout"): + controller.children.wait(handle, timeout_s=15) + assert_reaped(pids) + finally: + controller.close() + + +def test_runner_exception_cleans_up_managed_children(tmp_path): + def root_chat(messages, model=None): + if len(messages) > 2: + wait_ready(tmp_path) + raise ValueError("root failed") + return "```repl\nchild = rlm_query('blocked ignore-term', blocking=False)\n```" + + runner = LocalRLMRunner( + root_chat, + max_depth=3, + managed_children=make_controller(tmp_path).children.config, + ) + with pytest.raises(ValueError, match="root failed"): + runner.run("", "question") + assert_reaped([int(path.stem) for path in tmp_path.glob("*.ready")]) + + +def test_non_importable_factory_rejected(): + with pytest.raises(ValueError, match="pickl"): + ManagedChildConfig(chat_factory=lambda: lambda *_: "answer") + + +@pytest.mark.parametrize( + "kwargs", + [ + {"max_depth": 0}, + {"max_batch_workers": 0}, + {"max_children_total": -1}, + {"per_child_timeout_s": float("nan")}, + ], +) +def test_invalid_limits_rejected_before_spawning(tmp_path, kwargs): + with pytest.raises(ValueError): + make_controller(tmp_path, **kwargs) + + +def test_callbacks_and_trace_run_in_supervisor(tmp_path, monkeypatch): + controller = make_controller(tmp_path) + events = [] + callback_done = threading.Event() + + def on_complete(*args): + events.append((os.getpid(), "complete", args, threading.get_ident())) + callback_done.set() + + controller.children.on_subcall_start = lambda *args: events.append( + (os.getpid(), "start", args, threading.get_ident()) + ) + controller.children.on_subcall_complete = on_complete + try: + handle = controller.children.spawn("answer", model="fake-model") + assert controller.children.wait(handle, timeout_s=15) == "42" + assert callback_done.wait(5) + assert [event[:2] for event in events] == [ + (os.getpid(), "start"), + (os.getpid(), "complete"), + ] + assert all(event[2][1] == "fake-model" for event in events) + assert events[0][3] == events[1][3] + assert events[0][3] != threading.get_ident() + assert events[0][3] != controller.children._jobs[handle].thread.ident + assert controller.backend.child_traces[0].result_preview == "42" + assert_reaped([controller.children._jobs[handle].pid]) + + def unexpected_signal(*args): + pytest.fail("A reaped worker's PID may already belong to another process") + + monkeypatch.setattr(os, "kill", unexpected_signal) + controller.children._signal(controller.children._jobs[handle], signal.SIGKILL) + finally: + controller.close() + + +def test_invalid_prompt_does_not_start_worker(tmp_path): + controller = make_controller(tmp_path) + try: + with pytest.raises(TypeError): + controller.children.spawn(None) + assert not controller.children._jobs + finally: + controller.close() + + +@pytest.mark.parametrize("first_operation", ["reset", "close"]) +def test_reset_and_close_are_serialized(tmp_path, monkeypatch, first_operation): + controller = make_controller(tmp_path) + manager = controller.children + drained = threading.Event() + release_first = threading.Event() + second_entered = threading.Event() + original_stop = manager.stop + + def pause_after_drain(): + original_stop() + if threading.current_thread().name == "first-lifecycle": + drained.set() + if not release_first.wait(10): + raise TimeoutError("lifecycle test barrier was not released") + + operations = {"reset": manager.reset, "close": manager.close} + second_operation = "close" if first_operation == "reset" else "reset" + + def second_call(): + second_entered.set() + operations[second_operation]() + + monkeypatch.setattr(manager, "stop", pause_after_drain) + second_thread = None + first_thread, first_done, first_errors = start_call( + operations[first_operation], "first-lifecycle" + ) + try: + assert drained.wait(10) + second_thread, second_done, second_errors = start_call( + second_call, "second-lifecycle" + ) + assert second_entered.wait(5) + assert not second_done.wait(0.25) + finally: + release_first.set() + first_thread.join(10) + if second_thread is not None: + second_thread.join(10) + monkeypatch.setattr(manager, "stop", original_stop) + manager.close() + + assert first_done.is_set() and not first_thread.is_alive() + assert second_done.is_set() and not second_thread.is_alive() + assert not first_errors + if second_operation == "reset": + assert len(second_errors) == 1 + assert isinstance(second_errors[0], RuntimeError) + assert "closed" in str(second_errors[0]) + else: + assert not second_errors + with pytest.raises(RuntimeError, match="closed"): + manager.spawn("answer") + + +@pytest.mark.parametrize("callback_phase", ["start", "complete"]) +@pytest.mark.parametrize("operation", ["reset", "close"]) +def test_blocked_callbacks_do_not_block_lifecycle(tmp_path, callback_phase, operation): + controller = make_controller(tmp_path) + manager = controller.children + env = LocalREPLEnv(recursive_controller=controller) + env.reset() + entered = threading.Event() + release_callback = threading.Event() + callback_done = threading.Event() + events = [] + handles = [] + + def observe(phase): + events.append(phase) + if phase == callback_phase: + entered.set() + if not release_callback.wait(20): + raise TimeoutError("callback test barrier was not released") + + def on_complete(*args): + observe("complete") + callback_done.set() + + manager.on_subcall_start = lambda *args: observe("start") + manager.on_subcall_complete = on_complete + prompt = "blocked ignore-term" if callback_phase == "start" else "answer" + spawn_thread, spawn_done, spawn_errors = start_call( + lambda: handles.append(manager.spawn(prompt)), "callback-spawn" + ) + operation_thread = None + try: + assert entered.wait(15) + assert spawn_done.wait(0.5) + assert not spawn_errors + job = manager._jobs[handles[0]] + pids = wait_ready(tmp_path) if callback_phase == "start" else [job.pid] + operations = {"reset": env.reset, "close": env.close} + operation_thread, operation_done, operation_errors = start_call( + operations[operation], "callback-lifecycle" + ) + assert operation_done.wait(4) + assert not operation_errors + assert job.finished.is_set() + assert_reaped(pids) + assert not release_callback.is_set() + assert not callback_done.is_set() + finally: + release_callback.set() + spawn_thread.join(10) + if operation_thread is not None: + operation_thread.join(10) + env.close() + + assert not spawn_thread.is_alive() + assert not operation_thread.is_alive() + assert callback_done.wait(5) + assert events == ["start", "complete"] + manager._callback_thread.join(5) + assert not manager._callback_thread.is_alive() + + +@pytest.mark.parametrize("callback_phase", ["start", "complete"]) +def test_callback_exceptions_do_not_stop_dispatcher(tmp_path, callback_phase): + controller = make_controller(tmp_path) + events = [] + callback_done = threading.Event() + + def observe(phase): + events.append(phase) + if phase == callback_phase: + raise ValueError("observer failed") + + def on_complete(*args): + try: + observe("complete") + finally: + callback_done.set() + + controller.children.on_subcall_start = lambda *args: observe("start") + controller.children.on_subcall_complete = on_complete + after_failure = threading.Event() + try: + controller.children._notify( + controller.children.on_subcall_start, 1, "default", "answer" + ) + controller.children._notify( + controller.children.on_subcall_complete, 1, "default", 0.0, None + ) + controller.children._notify(after_failure.set) + assert callback_done.wait(5) + assert after_failure.wait(5) + assert events == ["start", "complete"] + controller.children.check_ready() + assert not controller.children._jobs + finally: + controller.close() + controller.children._callback_thread.join(5) + assert not controller.children._callback_thread.is_alive()