Merge pull request #171 from alexzhang13/fix_docker

Add docker changes to support persistence + parallel sub-calling

Author: alexzhang13

README.md

@@ -100,7 +100,9 @@ The default `local` environment `LocalREPL` runs in the same process as the RLM
 `IPythonREPL` runs cells inside a real IPython session — either in-process (default) or in a separate `ipykernel` subprocess. Subprocess mode adds hard `cell_timeout` enforcement and full namespace isolation from the RLM host. See the [IPythonREPL docs](https://alexzhang13.github.io/rlm/environments/ipython) for details.
 
 #### Docker <img src="https://github.com/docker.png" alt="Docker" height="20" style="vertical-align: middle;"/> (*requires [Docker installed](https://docs.docker.com/desktop/setup/install/)*)
-We also support a Docker-based environment called `DockerREPL` that launches the REPL environment as a Docker image. By default, we use the `python:3.11-slim` image, but the user can specify custom images as well.
+We also support a Docker-based environment called `DockerREPL` that launches the REPL environment as a Docker image. By default, we use the `python:3.11-slim` image, but the user can specify custom images as well. The container runs fully isolated from the host; a lightweight host-side proxy bridges LM access back into the container.
+
+`DockerREPL` supports the full feature set of the local environment: single LM calls (`llm_query` / `llm_query_batched`), recursive sub-RLM calls (`rlm_query` / `rlm_query_batched`, including parallel batched sub-calls bounded by `max_concurrent_subcalls`), `custom_tools` / `custom_sub_tools`, `persistent=True` multi-turn sessions (versioned `context_N` / `history_N` reused across `completion()` calls), and `compaction=True` auto-summarization of the running `history`. For isolated environments, custom tools should be passed as Python code strings or JSON-serializable values (host callables cannot cross the process boundary).
 
 ### Isolated Environments
 We support several different REPL environments that run on separate, cloud-based machines. Whenever a recursive sub-call is made in these instances, it is requested from the host process.

docs/getting-started.md

@@ -214,7 +214,9 @@ rlm = RLM(
 
 ### Docker
 
-Code runs in a Docker container with full isolation.
+Code runs in a Docker container with full isolation. A lightweight host-side
+proxy bridges LM access back into the container, so the container never needs
+API keys or direct network access to the model provider.
 
 ```python
 rlm = RLM(
@@ -227,7 +229,19 @@ rlm = RLM(
 )
 ```
 
-**Pros:** Containerized isolation, reproducible  
+`DockerREPL` supports the same capabilities as the local environment:
+
+- `llm_query` / `llm_query_batched` — single LM completions
+- `rlm_query` / `rlm_query_batched` — recursive sub-RLM calls (with `max_depth > 1`),
+  including parallel batched sub-calls bounded by `max_concurrent_subcalls`
+- `custom_tools` / `custom_sub_tools` — injected functions/data (pass as Python
+  code strings or JSON-serializable values, since host callables cannot cross
+  the container boundary)
+- `persistent=True` — multi-turn sessions; the container and its namespace are
+  kept alive across `completion()` calls, with versioned `context_N` / `history_N`
+- `compaction=True` — auto-summarizes the running `history` when context fills up
+
+**Pros:** Containerized isolation, reproducible, full feature parity with local  
 **Cons:** Requires Docker, slower startup
 
 ### Modal

examples/docker_repl_example.py

@@ -1,16 +1,24 @@
 """
-Docker REPL example with code execution and LLM queries.
+Docker REPL example with code execution, LLM queries, recursive sub-RLM calls,
+and custom tools.
 
 Setup:
     1. Ensure Docker is running
     2. Run: python -m examples.docker_repl_example
 
 The default image (python:3.11-slim) will be pulled automatically.
+
+Persistence and compaction:
+    DockerREPL also supports multi-turn persistence and compaction. Drive them
+    through the RLM rather than the env directly, e.g.:
+
+        rlm = RLM(backend=..., environment="docker", persistent=True)   # multi-turn
+        rlm = RLM(backend=..., environment="docker", compaction=True)   # auto-summarize
 """
 
 from rlm.clients.base_lm import BaseLM
 from rlm.core.lm_handler import LMHandler
-from rlm.core.types import ModelUsageSummary, UsageSummary
+from rlm.core.types import ModelUsageSummary, RLMChatCompletion, UsageSummary
 from rlm.environments.docker_repl import DockerREPL
 
 
@@ -61,6 +69,32 @@ def main():
             result = repl.execute_code("print(len(rs))")
             print(f"  Batched count: {result.stdout.strip()}")
 
+    # Recursive sub-RLM calls + custom tools
+    print("\n[3] Recursive sub-RLM (rlm_query) + custom tools")
+
+    def mock_subcall(prompt, model=None):
+        # In real use this is RLM._subcall, which spawns a child RLM.
+        return RLMChatCompletion(
+            root_model=model or "child",
+            prompt=prompt,
+            response=f"child-answer({prompt})",
+            usage_summary=UsageSummary(model_usage_summaries={}),
+            execution_time=0.0,
+        )
+
+    with DockerREPL(
+        subcall_fn=mock_subcall,
+        custom_tools={"triple": "def triple(x):\n    return x * 3", "CONST": 7},
+    ) as repl:
+        result = repl.execute_code("print(triple(CONST))")
+        print(f"  custom tools triple(CONST) → {result.stdout.strip()}")
+
+        result = repl.execute_code('print(rlm_query("solve subtask"))')
+        print(f"  rlm_query → {result.stdout.strip()}")
+
+        result = repl.execute_code('print(rlm_query_batched(["a", "b", "c"]))')
+        print(f"  rlm_query_batched → {result.stdout.strip()}")
+
     print("\n" + "=" * 50)
     print("Done!")
 

rlm/core/rlm.py

@@ -272,15 +272,17 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]):
             env_kwargs["lm_handler_address"] = (lm_handler.host, lm_handler.port)
             env_kwargs["context_payload"] = prompt
             env_kwargs["depth"] = self.depth + 1  # Environment depth is RLM depth + 1
-            # For local/ipython environments with max_depth > 1, pass subcall callback for recursive RLM calls
-            if self.environment_type in ("local", "ipython") and self.max_depth > 1:
+            # For environments that support recursive RLM calls, pass the subcall
+            # callback when max_depth > 1. local/ipython invoke it in-process;
+            # docker invokes it via its host-side proxy (/rlm_query endpoints).
+            if self.environment_type in ("local", "ipython", "docker") and self.max_depth > 1:
                 env_kwargs["subcall_fn"] = self._subcall
             # Pass custom tools to the environment
             if self.custom_tools is not None:
                 env_kwargs["custom_tools"] = self.custom_tools
             if self.custom_sub_tools is not None:
                 env_kwargs["custom_sub_tools"] = self.custom_sub_tools
-            if self.compaction and self.environment_type == "local":
+            if self.compaction and self.environment_type in ("local", "docker"):
                 env_kwargs["compaction"] = True
             env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls
             environment: BaseEnv = get_environment(self.environment_type, env_kwargs)
@@ -876,13 +878,13 @@ def _validate_persistent_environment_support(self) -> None:
         - add_context(payload, index): Add new context for multi-turn conversations
         - get_context_count(): Return the number of loaded contexts
 
-        Currently only 'local' (LocalREPL) supports these methods.
+        Currently 'local', 'ipython', and 'docker' support these methods.
 
         Raises:
             ValueError: If the environment type does not support persistent mode.
         """
         # Known environments that support persistence
-        persistent_supported_environments = {"local", "ipython"}
+        persistent_supported_environments = {"local", "ipython", "docker"}
 
         if self.environment_type not in persistent_supported_environments:
             raise ValueError(