Spaces:
Runtime error
Runtime error
| """ | |
| Docker container management for ToolStore. | |
| Two separate container models: | |
| 1. **Single shared worker** — one persistent Python worker for ALL toolset | |
| executions. The container starts once and stays alive forever. Toolset | |
| code is piped in/out per invocation via newline-delimited JSON. | |
| 2. **MCP-server containers** — one container per MCP server (started by | |
| ``DockerTransport`` in ``transport.py``). JSON-RPC over stdin/stdout. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import subprocess | |
| import threading | |
| import time | |
| from typing import Dict, Any, Optional, List | |
| # --------------------------------------------------------------------------- | |
| # Worker bootstrap (injected into every warm Python container) | |
| # --------------------------------------------------------------------------- | |
| WORKER_BOOTSTRAP = r''' | |
| import sys as _sys, json as _json, traceback as _tb | |
| from io import StringIO as _StringIO | |
| _modules = {} # module_name -> namespace holding its functions | |
| _sys.__stdout__.write("READY\n") | |
| _sys.__stdout__.flush() | |
| while True: | |
| _line = _sys.stdin.readline() | |
| if not _line: | |
| break | |
| try: | |
| _req = _json.loads(_line) | |
| except _json.JSONDecodeError: | |
| continue | |
| _rid = _req["id"] | |
| _action = _req.get("action", "call") | |
| if _action == "load": | |
| # Load a module — exec the code, store its namespace. | |
| _name = _req["module"] | |
| _code = _req["code"] | |
| _ns = {} | |
| try: | |
| exec(_code, _ns) | |
| except Exception: | |
| _res = {"id": _rid, "ok": False, "output": _tb.format_exc()} | |
| else: | |
| _modules[_name] = _ns | |
| _res = {"id": _rid, "ok": True, "output": f"module '{_name}' loaded"} | |
| elif _action == "call": | |
| _mod = _req["module"] | |
| _func = _req["function"] | |
| _args = _req.get("args", {}) | |
| _ns = _modules.get(_mod) | |
| if _ns is None: | |
| _res = {"id": _rid, "ok": False, | |
| "output": f"module '{_mod}' not loaded"} | |
| else: | |
| _fn = _ns.get(_func) | |
| if _fn is None: | |
| _res = {"id": _rid, "ok": False, | |
| "output": f"function '{_func}' not found in module '{_mod}'"} | |
| elif not callable(_fn): | |
| _res = {"id": _rid, "ok": False, | |
| "output": f"'{_func}' in module '{_mod}' is not callable"} | |
| else: | |
| _cap = _StringIO() | |
| _old = _sys.stdout | |
| _sys.stdout = _cap | |
| try: | |
| _result = _fn(**_args) | |
| _output = _cap.getvalue() | |
| if _result is not None: | |
| _output += _json.dumps(_result, default=str) | |
| _res = {"id": _rid, "ok": True, "output": _output} | |
| except Exception: | |
| _res = {"id": _rid, "ok": False, | |
| "output": _cap.getvalue() + _tb.format_exc()} | |
| finally: | |
| _sys.stdout = _old | |
| else: | |
| _res = {"id": _rid, "ok": False, | |
| "output": f"unknown action '{_action}'"} | |
| _sys.__stdout__.write(_json.dumps(_res) + "\n") | |
| _sys.__stdout__.flush() | |
| ''' | |
| # --------------------------------------------------------------------------- | |
| # Helpers | |
| # --------------------------------------------------------------------------- | |
| def check_docker_available() -> Optional[str]: | |
| """Return ``None`` if Docker is usable, otherwise an error message.""" | |
| import shutil | |
| if shutil.which("docker") is None: | |
| return "Docker is not installed or not on PATH." | |
| try: | |
| subprocess.run( | |
| ["docker", "info"], | |
| capture_output=True, text=True, timeout=10, | |
| ) | |
| except subprocess.TimeoutExpired: | |
| return "Docker daemon is unresponsive (timeout)." | |
| except FileNotFoundError: | |
| return "Docker CLI not found." | |
| except Exception as exc: | |
| return f"Cannot communicate with Docker daemon: {exc}" | |
| return None | |
| def detect_dind() -> bool: | |
| """Return ``True`` if this process is running inside a Docker container.""" | |
| import os | |
| if os.path.exists("/.dockerenv"): | |
| return True | |
| try: | |
| with open("/proc/self/cgroup", "r") as fh: | |
| if "docker" in fh.read() or "kubepods" in fh.read(): | |
| return True | |
| except Exception: | |
| pass | |
| if os.environ.get("THINKTOOL_DOCKER") == "1": | |
| return True | |
| return False | |
| def dind_socket_check() -> Optional[str]: | |
| """If we are inside Docker but ``/var/run/docker.sock`` is missing, | |
| return a helpful error message; otherwise ``None``.""" | |
| import os | |
| if detect_dind() and not os.path.exists("/var/run/docker.sock"): | |
| return ( | |
| "Error: Running inside Docker but /var/run/docker.sock is not mounted.\n" | |
| "Add this to docker-compose.yml:\n" | |
| " volumes:\n" | |
| " - /var/run/docker.sock:/var/run/docker.sock" | |
| ) | |
| return None | |
| # --------------------------------------------------------------------------- | |
| # WarmContainer — the single shared Python worker for all toolset executions | |
| # --------------------------------------------------------------------------- | |
| class WarmContainer: | |
| """The single persistent Docker container that runs ALL toolset executions. | |
| Started once (lazily, on first use) and kept alive forever. Toolset | |
| modules are loaded once and cached in the worker; subsequent calls | |
| invoke named functions from those modules. | |
| """ | |
| def __init__(self, image: str, container_name: str) -> None: | |
| self.image = image | |
| self.name = container_name | |
| self.proc: Optional[subprocess.Popen] = None | |
| self._request_id: int = 0 | |
| self._lock = threading.Lock() | |
| self._alive = False | |
| self._loaded_modules: set = set() # track what we've sent load for | |
| # -- lifecycle ----------------------------------------------------------- | |
| def start(self) -> None: | |
| """Run *docker run -i* and bootstrap the worker.""" | |
| self.proc = subprocess.Popen( | |
| [ | |
| "docker", "run", "-i", "--rm", | |
| "--name", self.name, | |
| "--network", "none", | |
| "--cpus", "1", | |
| "--memory", "256m", | |
| self.image, | |
| "python", "-u", | |
| ], | |
| stdin=subprocess.PIPE, | |
| stdout=subprocess.PIPE, | |
| stderr=subprocess.PIPE, | |
| text=True, | |
| ) | |
| # Bootstrap: load the worker loop as a module via the load action. | |
| self._write_line(json.dumps({ | |
| "id": "_bootstrap", | |
| "action": "load", | |
| "module": "_worker", | |
| "code": WORKER_BOOTSTRAP, | |
| })) | |
| # Discard the bootstrap response, then send a dummy call to | |
| # confirm the worker is running and reading. | |
| self._read_line(timeout=15) | |
| self._write_line(json.dumps({ | |
| "id": "_ping", | |
| "action": "call", | |
| "module": "_worker", | |
| "function": "_whoami_does_not_exist", | |
| "args": {}, | |
| })) | |
| raw = self._read_line(timeout=15) | |
| if raw is None: | |
| self.stop() | |
| raise RuntimeError("Worker did not respond after bootstrap") | |
| # Any response (even error) means the worker is alive. | |
| self._alive = True | |
| def stop(self) -> None: | |
| self._alive = False | |
| if self.proc: | |
| try: | |
| self.proc.stdin.close() | |
| except Exception: | |
| pass | |
| try: | |
| self.proc.terminate() | |
| self.proc.wait(timeout=5) | |
| except Exception: | |
| try: | |
| self.proc.kill() | |
| except Exception: | |
| pass | |
| self.proc = None | |
| def is_alive(self) -> bool: | |
| return ( | |
| self._alive | |
| and self.proc is not None | |
| and self.proc.poll() is None | |
| ) | |
| # -- load / call --------------------------------------------------------- | |
| def load_module(self, name: str, code: str, timeout_s: int = 15) -> str: | |
| """Load *code* into the worker as a named module (idempotent). | |
| Returns the worker response string (usually a confirmation or error). | |
| """ | |
| if name in self._loaded_modules: | |
| return "" # already loaded | |
| with self._lock: | |
| rid = self._request_id | |
| self._request_id += 1 | |
| self._write_line(json.dumps({ | |
| "id": rid, "action": "load", "module": name, "code": code, | |
| })) | |
| raw = self._read_line(timeout=timeout_s) | |
| if raw is None: | |
| return f"Error: load timed out for module '{name}'" | |
| try: | |
| result = json.loads(raw) | |
| except json.JSONDecodeError: | |
| return f"Error: invalid load response: {raw!r}" | |
| if result.get("ok"): | |
| self._loaded_modules.add(name) | |
| return "" | |
| return result.get("output", "unknown load error") | |
| def call_function( | |
| self, module: str, function: str, | |
| args: Dict[str, Any], timeout_s: int, | |
| ) -> str: | |
| """Call *function* inside *module* with keyword *args*.""" | |
| with self._lock: | |
| rid = self._request_id | |
| self._request_id += 1 | |
| self._write_line(json.dumps({ | |
| "id": rid, "action": "call", | |
| "module": module, "function": function, | |
| "args": args, | |
| })) | |
| raw = self._read_line(timeout=timeout_s + 5) | |
| if raw is None: | |
| return f"Error: call timed out after {timeout_s}s" | |
| try: | |
| result = json.loads(raw) | |
| except json.JSONDecodeError: | |
| return f"Error: invalid call response: {raw!r}" | |
| if result.get("ok"): | |
| out = result.get("output", "") | |
| return out.strip() if out.strip() else "(no output)" | |
| return f"Error:\n{result.get('output', 'unknown error')}" | |
| # -- I/O helpers --------------------------------------------------------- | |
| def _write_line(self, text: str) -> None: | |
| assert self.proc and self.proc.stdin | |
| self.proc.stdin.write(text + "\n") | |
| self.proc.stdin.flush() | |
| def _read_line(self, timeout: float) -> Optional[str]: | |
| """Read a single line from the worker stdout with a timeout. | |
| Uses a separate thread so we can enforce a wall-clock deadline | |
| without relying on file-descriptor-level timeouts. | |
| """ | |
| assert self.proc and self.proc.stdout | |
| result: List[Optional[str]] = [None] | |
| def _target() -> None: | |
| try: | |
| result[0] = self.proc.stdout.readline() # type: ignore[union-attr] | |
| except Exception: | |
| result[0] = None | |
| t = threading.Thread(target=_target, daemon=True) | |
| t.start() | |
| t.join(timeout) | |
| if t.is_alive(): | |
| # Timed out — the thread will eventually read and discard. | |
| return None | |
| return result[0] | |
| # --------------------------------------------------------------------------- | |
| # Module-level singleton — the one shared worker | |
| # --------------------------------------------------------------------------- | |
| _worker: Optional[WarmContainer] = None | |
| _worker_lock = threading.Lock() | |
| def get_worker() -> WarmContainer: | |
| """Return (or lazily start) the single shared worker container. | |
| The image comes from the user's config (``default_docker_image``). | |
| """ | |
| global _worker | |
| if _worker is not None and _worker.is_alive(): | |
| return _worker | |
| with _worker_lock: | |
| # Double-check inside the lock | |
| if _worker is not None and _worker.is_alive(): | |
| return _worker | |
| # Import here to avoid circularity at module level | |
| from toolstore.config_manager import ConfigManager | |
| cfg = ConfigManager() | |
| cfg.load() | |
| image = cfg.get_default_docker_image() | |
| name = f"ts_worker" | |
| _worker = WarmContainer(image, name) | |
| _worker.start() | |
| return _worker | |
| def shutdown_worker() -> None: | |
| """Stop the shared worker (e.g. at process exit).""" | |
| global _worker | |
| with _worker_lock: | |
| if _worker: | |
| _worker.stop() | |
| _worker = None | |