groknroll 2.0.0__py3-none-any.whl

groknroll/environments/modal_repl.py

@@ -0,0 +1,502 @@
+import base64
+import json
+import textwrap
+import threading
+import time
+
+import modal
+import requests
+
+from groknroll.core.comms_utils import LMRequest, send_lm_request, send_lm_request_batched
+from groknroll.core.types import REPLResult, RLMChatCompletion
+from groknroll.environments.base_env import IsolatedEnv
+from groknroll.environments.constants import APT_PACKAGES, PIP_PACKAGES
+
+# =============================================================================
+# Default Modal Image
+# =============================================================================
+
+
+def get_default_image() -> modal.Image:
+    """
+    Build a default Modal image with common libraries for data science,
+    math, and general Python work.
+    """
+    return (
+        modal.Image.debian_slim(python_version="3.11")
+        .apt_install(*APT_PACKAGES)
+        .pip_install(*PIP_PACKAGES)
+    )
+
+
+# =============================================================================
+# Broker Server Script (runs inside sandbox, handles LLM request queue)
+# =============================================================================
+
+_BROKER_SCRIPT = textwrap.dedent(
+    '''
+import json
+import threading
+import uuid
+from flask import Flask, request, jsonify
+
+app = Flask(__name__)
+
+# Request queue: {request_id: {"request": {...}, "response": None, "event": Event}}
+pending_requests = {}
+lock = threading.Lock()
+
+@app.route("/health")
+def health():
+    return jsonify({"status": "ok"})
+
+@app.route("/enqueue", methods=["POST"])
+def enqueue():
+    """Called by sandbox code to submit an LLM request and wait for response."""
+    data = request.json
+    request_id = str(uuid.uuid4())
+    event = threading.Event()
+
+    with lock:
+        pending_requests[request_id] = {
+            "request": data,
+            "response": None,
+            "event": event,
+        }
+
+    # Wait for response (with timeout)
+    event.wait(timeout=300)
+
+    with lock:
+        entry = pending_requests.pop(request_id, None)
+
+    if entry and entry["response"] is not None:
+        return jsonify(entry["response"])
+    else:
+        return jsonify({"error": "Request timed out"}), 504
+
+@app.route("/pending")
+def get_pending():
+    """Called by ModalREPL to get pending requests."""
+    with lock:
+        pending = [
+            {"id": rid, "request": entry["request"]}
+            for rid, entry in pending_requests.items()
+            if entry["response"] is None
+        ]
+    return jsonify({"pending": pending})
+
+@app.route("/respond", methods=["POST"])
+def respond():
+    """Called by ModalREPL to submit a response."""
+    data = request.json
+    request_id = data.get("id")
+    response = data.get("response")
+
+    with lock:
+        if request_id in pending_requests:
+            pending_requests[request_id]["response"] = response
+            pending_requests[request_id]["event"].set()
+            return jsonify({"status": "ok"})
+
+    return jsonify({"error": "Request not found"}), 404
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=8080, threaded=True)
+'''
+)
+
+
+# =============================================================================
+# Execution Script (runs inside the sandbox for each code block)
+# =============================================================================
+
+
+def _build_exec_script(code: str, broker_port: int = 8080, depth: int = 1) -> str:
+    """
+    Build a script that executes code with state persistence.
+    LLM queries go through the local broker server.
+    """
+    code_b64 = base64.b64encode(code.encode()).decode()
+
+    return textwrap.dedent(
+        f'''
+import sys
+import io
+import json
+import base64
+import traceback
+import os
+import requests
+
+try:
+    import dill
+except ImportError:
+    import pickle as dill
+
+# =============================================================================
+# LLM Query Functions (via local broker)
+# =============================================================================
+
+BROKER_URL = "http://127.0.0.1:{broker_port}"
+
+def llm_query(prompt, model=None):
+    """Query the LM via the broker."""
+    try:
+        response = requests.post(
+            f"{{BROKER_URL}}/enqueue",
+            json={{"type": "single", "prompt": prompt, "model": model, "depth": {depth}}},
+            timeout=300,
+        )
+        data = response.json()
+        if data.get("error"):
+            return f"Error: {{data['error']}}"
+        return data.get("response", "Error: No response")
+    except Exception as e:
+        return f"Error: LM query failed - {{e}}"
+
+
+def llm_query_batched(prompts, model=None):
+    """Query the LM with multiple prompts."""
+    try:
+        response = requests.post(
+            f"{{BROKER_URL}}/enqueue",
+            json={{"type": "batched", "prompts": prompts, "model": model, "depth": {depth}}},
+            timeout=300,
+        )
+        data = response.json()
+        if data.get("error"):
+            return [f"Error: {{data['error']}}"] * len(prompts)
+        return data.get("responses", ["Error: No response"] * len(prompts))
+    except Exception as e:
+        return [f"Error: LM query failed - {{e}}"] * len(prompts)
+
+
+# =============================================================================
+# State Management
+# =============================================================================
+
+STATE_FILE = "/tmp/rlm_state.dill"
+
+def load_state():
+    if os.path.exists(STATE_FILE):
+        try:
+            with open(STATE_FILE, "rb") as f:
+                return dill.load(f)
+        except:
+            pass
+    return {{}}
+
+def save_state(state):
+    clean_state = {{}}
+    for k, v in state.items():
+        if k.startswith("_"):
+            continue
+        try:
+            dill.dumps(v)
+            clean_state[k] = v
+        except:
+            pass
+    with open(STATE_FILE, "wb") as f:
+        dill.dump(clean_state, f)
+
+def serialize_locals(state):
+    result = {{}}
+    for k, v in state.items():
+        if k.startswith("_"):
+            continue
+        try:
+            result[k] = repr(v)
+        except:
+            result[k] = f"<{{type(v).__name__}}>"
+    return result
+
+# =============================================================================
+# Execution
+# =============================================================================
+
+_locals = load_state()
+
+def FINAL_VAR(variable_name):
+    variable_name = variable_name.strip().strip("\\"\\'")
+    if variable_name in _locals:
+        return str(_locals[variable_name])
+    return f"Error: Variable '{{variable_name}}' not found"
+
+_globals = {{
+    "__builtins__": __builtins__,
+    "__name__": "__main__",
+    "llm_query": llm_query,
+    "llm_query_batched": llm_query_batched,
+    "FINAL_VAR": FINAL_VAR,
+}}
+
+code = base64.b64decode("{code_b64}").decode()
+
+stdout_buf = io.StringIO()
+stderr_buf = io.StringIO()
+old_stdout, old_stderr = sys.stdout, sys.stderr
+
+try:
+    sys.stdout = stdout_buf
+    sys.stderr = stderr_buf
+    combined = {{**_globals, **_locals}}
+    exec(code, combined, combined)
+    for key, value in combined.items():
+        if key not in _globals and not key.startswith("_"):
+            _locals[key] = value
+except Exception as e:
+    traceback.print_exc(file=stderr_buf)
+finally:
+    sys.stdout = old_stdout
+    sys.stderr = old_stderr
+
+save_state(_locals)
+
+result = {{
+    "stdout": stdout_buf.getvalue(),
+    "stderr": stderr_buf.getvalue(),
+    "locals": serialize_locals(_locals),
+}}
+print(json.dumps(result))
+'''
+    )
+
+
+class ModalREPL(IsolatedEnv):
+    """
+    Modal REPL environment that runs Python code in a Modal Sandbox.
+
+    Uses Modal tunnels for LLM communication:
+    - Sandbox runs a broker server exposed via encrypted_ports
+    - ModalREPL polls the broker for pending LLM requests
+    - ModalREPL forwards requests to the LM handler and posts responses back
+    """
+
+    BROKER_PORT = 8080
+
+    def __init__(
+        self,
+        app_name: str = "rlm-sandbox",
+        image: modal.Image | None = None,
+        timeout: int = 600,
+        lm_handler_address: tuple[str, int] | None = None,
+        context_payload: dict | list | str | None = None,
+        setup_code: str | None = None,
+        persistent: bool = False,
+        depth: int = 1,
+        **kwargs,
+    ):
+        if persistent:
+            raise NotImplementedError(
+                "Persistent REPLs are currently not supported for environment: ModalREPL"
+            )
+        super().__init__(persistent=persistent, depth=depth, **kwargs)
+
+        self.app_name = app_name
+        self.timeout = timeout
+        self.lm_handler_address = lm_handler_address
+
+        self.image = image or get_default_image()
+
+        self.app = None
+        self.sandbox = None
+        self.broker_process = None
+        self.broker_url: str | None = None
+        self.poller_thread: threading.Thread | None = None
+        self.poller_stop = threading.Event()
+        self.pending_llm_calls: list[RLMChatCompletion] = []
+        self._calls_lock = threading.Lock()
+
+        self.setup()
+
+        if context_payload is not None:
+            self.load_context(context_payload)
+
+        if setup_code:
+            self.execute_code(setup_code)
+
+    def setup(self):
+        """Create the Modal app, sandbox, broker, and start polling."""
+        self.app = modal.App.lookup(self.app_name, create_if_missing=True)
+
+        # Create sandbox with encrypted port for broker
+        self.sandbox = modal.Sandbox.create(
+            app=self.app,
+            image=self.image,
+            timeout=self.timeout,
+            encrypted_ports=[self.BROKER_PORT],
+        )
+
+        # Start the broker server in the sandbox
+        self.broker_process = self.sandbox.exec(
+            "python",
+            "-c",
+            _BROKER_SCRIPT,
+        )
+
+        # Wait for broker to be ready
+        time.sleep(2)
+
+        # Get the tunnel URL
+        tunnels = self.sandbox.tunnels()
+        if self.BROKER_PORT in tunnels:
+            self.broker_url = tunnels[self.BROKER_PORT].url
+
+        # Start polling thread if we have an LM handler
+        if self.lm_handler_address and self.broker_url:
+            self.poller_stop.clear()
+            self.poller_thread = threading.Thread(target=self._poll_broker, daemon=True)
+            self.poller_thread.start()
+
+    def _poll_broker(self):
+        """Poll the broker for pending LLM requests and handle them."""
+        while not self.poller_stop.is_set():
+            try:
+                # Get pending requests
+                resp = requests.get(
+                    f"{self.broker_url}/pending",
+                    timeout=5,
+                )
+                pending = resp.json().get("pending", [])
+
+                for item in pending:
+                    request_id = item["id"]
+                    req_data = item["request"]
+
+                    # Handle the request
+                    response = self._handle_llm_request(req_data)
+
+                    # Send response back
+                    requests.post(
+                        f"{self.broker_url}/respond",
+                        json={"id": request_id, "response": response},
+                        timeout=10,
+                    )
+
+            except requests.exceptions.RequestException:
+                pass
+            except Exception:
+                pass
+
+            time.sleep(0.1)
+
+    def _handle_llm_request(self, req_data: dict) -> dict:
+        """Handle an LLM request from the sandbox."""
+        req_type = req_data.get("type")
+        model = req_data.get("model")
+
+        if req_type == "single":
+            prompt = req_data.get("prompt")
+            request = LMRequest(prompt=prompt, model=model, depth=self.depth)
+            response = send_lm_request(self.lm_handler_address, request)
+
+            if not response.success:
+                return {"error": response.error}
+
+            # Track the call
+            with self._calls_lock:
+                self.pending_llm_calls.append(response.chat_completion)
+
+            return {"response": response.chat_completion.response}
+
+        elif req_type == "batched":
+            prompts = req_data.get("prompts", [])
+            responses = send_lm_request_batched(
+                self.lm_handler_address, prompts, model=model, depth=self.depth
+            )
+
+            results = []
+            for resp in responses:
+                if not resp.success:
+                    results.append(f"Error: {resp.error}")
+                else:
+                    with self._calls_lock:
+                        self.pending_llm_calls.append(resp.chat_completion)
+                    results.append(resp.chat_completion.response)
+
+            return {"responses": results}
+
+        return {"error": "Unknown request type"}
+
+    def load_context(self, context_payload: dict | list | str):
+        """Load context into the sandbox environment."""
+        if isinstance(context_payload, str):
+            escaped = context_payload.replace("\\", "\\\\").replace('"""', '\\"\\"\\"')
+            context_code = f'context = """{escaped}"""'
+        else:
+            context_json = json.dumps(context_payload)
+            escaped_json = context_json.replace("\\", "\\\\").replace("'", "\\'")
+            context_code = f"import json; context = json.loads('{escaped_json}')"
+
+        self.execute_code(context_code)
+
+    def execute_code(self, code: str) -> REPLResult:
+        """Execute code in the Modal sandbox and return result."""
+        start_time = time.perf_counter()
+
+        # Clear pending LLM calls
+        with self._calls_lock:
+            self.pending_llm_calls.clear()
+
+        # Build and execute the script
+        script = _build_exec_script(code, self.BROKER_PORT, self.depth)
+        process = self.sandbox.exec("python", "-c", script)
+
+        # Read output
+        stdout = process.stdout.read()
+        stderr = process.stderr.read()
+
+        # Collect LLM calls made during this execution
+        with self._calls_lock:
+            pending_calls = self.pending_llm_calls.copy()
+            self.pending_llm_calls.clear()
+
+        execution_time = time.perf_counter() - start_time
+
+        # Parse the JSON result
+        try:
+            lines = stdout.strip().split("\n")
+            result_json = lines[-1] if lines else "{}"
+            result = json.loads(result_json)
+
+            return REPLResult(
+                stdout=result.get("stdout", ""),
+                stderr=result.get("stderr", "") + stderr,
+                locals=result.get("locals", {}),
+                execution_time=execution_time,
+                rlm_calls=pending_calls,
+            )
+        except json.JSONDecodeError:
+            return REPLResult(
+                stdout=stdout,
+                stderr=stderr or "Failed to parse execution result",
+                locals={},
+                execution_time=execution_time,
+                rlm_calls=pending_calls,
+            )
+
+    def cleanup(self):
+        """Terminate the sandbox and stop polling."""
+        # Stop the poller thread
+        if self.poller_thread is not None:
+            self.poller_stop.set()
+            self.poller_thread.join(timeout=2)
+            self.poller_thread = None
+
+        if self.sandbox is not None:
+            try:
+                self.sandbox.terminate()
+            except Exception:
+                pass
+            self.sandbox = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.cleanup()
+        return False
+
+    def __del__(self):
+        self.cleanup()