refactor packages into constants

Author: 13point5

AGENTS.md

@@ -18,6 +18,9 @@ uv pip install -e .
 
 # For Modal sandbox support
 uv pip install -e ".[modal]"
+
+# For Prime sandbox support
+uv pip install -e ".[prime]"
 ```
 
 ## General Guidelines

rlm/constants/__init__.py

@@ -0,0 +1,2 @@
+# Constants module
+

rlm/constants/repl.py

@@ -0,0 +1,32 @@
+# Default packages for isolated REPL environments (Modal, Prime, etc.)
+
+APT_PACKAGES = [
+    "build-essential",
+    "git",
+    "curl",
+    "wget",
+    "libopenblas-dev",
+    "liblapack-dev",
+]
+
+PIP_PACKAGES = [
+    # Data science essentials
+    "numpy>=1.26.0",
+    "pandas>=2.1.0",
+    "scipy>=1.11.0",
+    # Math & symbolic computation
+    "sympy>=1.12",
+    # HTTP & APIs
+    "requests>=2.31.0",
+    "httpx>=0.25.0",
+    "flask>=3.0.0",
+    # Data formats
+    "pyyaml>=6.0",
+    "toml>=0.10.2",
+    # Utilities
+    "tqdm>=4.66.0",
+    "python-dateutil>=2.8.2",
+    "regex>=2023.0.0",
+    # For state serialization
+    "dill>=0.3.7",
+]

rlm/environments/modal_repl.py

@@ -7,6 +7,7 @@
 import modal
 import requests
 
+from rlm.constants.repl import APT_PACKAGES, PIP_PACKAGES
 from rlm.core.comms_utils import LMRequest, send_lm_request, send_lm_request_batched
 from rlm.core.types import REPLResult, RLMChatCompletion
 from rlm.environments.base_env import IsolatedEnv
@@ -23,35 +24,8 @@ def get_default_image() -> modal.Image:
     """
     return (
         modal.Image.debian_slim(python_version="3.11")
-        .apt_install(
-            "build-essential",
-            "git",
-            "curl",
-            "wget",
-            "libopenblas-dev",
-            "liblapack-dev",
-        )
-        .pip_install(
-            # Data science essentials
-            "numpy>=1.26.0",
-            "pandas>=2.1.0",
-            "scipy>=1.11.0",
-            # Math & symbolic computation
-            "sympy>=1.12",
-            # HTTP & APIs
-            "requests>=2.31.0",
-            "httpx>=0.25.0",
-            "flask>=3.0.0",
-            # Data formats
-            "pyyaml>=6.0",
-            "toml>=0.10.2",
-            # Utilities
-            "tqdm>=4.66.0",
-            "python-dateutil>=2.8.2",
-            "regex>=2023.0.0",
-            # For state serialization
-            "dill>=0.3.7",
-        )
+        .apt_install(*APT_PACKAGES)
+        .pip_install(*PIP_PACKAGES)
     )
 
 

rlm/environments/prime_repl.py

@@ -20,6 +20,7 @@
     SandboxClient,
 )
 
+from rlm.constants.repl import APT_PACKAGES, PIP_PACKAGES
 from rlm.core.comms_utils import LMRequest, send_lm_request, send_lm_request_batched
 from rlm.core.types import REPLResult, RLMChatCompletion
 from rlm.environments.base_env import IsolatedEnv
@@ -328,11 +329,13 @@ def setup(self):
         # Wait for sandbox to be ready
         self.client.wait_for_creation(self.sandbox_id, max_attempts=self.timeout_minutes * 60)
 
-        # Install dependencies for the broker
-        self.client.execute_command(
-            self.sandbox_id,
-            "pip install flask requests dill",
-        )
+        # Install apt dependencies
+        apt_cmd = "apt-get update && apt-get install -y " + " ".join(APT_PACKAGES)
+        self.client.execute_command(self.sandbox_id, apt_cmd)
+
+        # Install pip dependencies
+        pip_cmd = "pip install " + " ".join(f'"{pkg}"' for pkg in PIP_PACKAGES)
+        self.client.execute_command(self.sandbox_id, pip_cmd)
 
         # Write the broker script to the sandbox.
         # Unlike Modal's sandbox.exec() which accepts separate args, Prime's
@@ -410,6 +413,7 @@ def _poll_broker(self):
                 # Get pending requests
                 resp = requests.get(
                     f"{self.broker_url}/pending",
+                    timeout=10,
                 )
                 pending = resp.json().get("pending", [])
 
@@ -424,6 +428,7 @@ def _poll_broker(self):
                     requests.post(
                         f"{self.broker_url}/respond",
                         json={"id": request_id, "response": response},
+                        timeout=10,
                     )
 
             except requests.exceptions.RequestException: