@aliwey/bmo 2.0.0

package/core/worker_manager.py

@@ -0,0 +1,71 @@
+import logging
+import os
+from typing import Callable, Optional
+
+from core.worker_subprocess import AsyncSubprocessWorker
+from core.worker_multiproc import MultiprocessWorker
+
+logger = logging.getLogger(__name__)
+
+
+class WorkerManager:
+    """Unified facade over both worker backends.
+    Selects backend via env WORKER_BACKEND (subprocess|multiproc).
+    """
+
+    def __init__(self, backend: Optional[str] = None):
+        if backend is None:
+            backend = os.getenv("WORKER_BACKEND", "subprocess")
+        self.backend = backend.lower()
+
+        if self.backend == "subprocess":
+            self._worker = AsyncSubprocessWorker()
+        elif self.backend == "multiproc":
+            self._worker = MultiprocessWorker()
+        else:
+            raise ValueError(f"Unknown worker backend: {backend}")
+
+        logger.info(f"WorkerManager using backend: {self.backend}")
+
+    async def start(self):
+        await self._worker.start()
+
+    async def send_query(
+        self,
+        session_id: str,
+        payload: dict,
+        base_url: str,
+        poll_interval: float = 2.0,
+        poll_timeout: float = 600.0,
+        callback: Optional[Callable] = None,
+    ) -> str:
+        return await self._worker.send_query(
+            session_id=session_id,
+            payload=payload,
+            base_url=base_url,
+            poll_interval=poll_interval,
+            poll_timeout=poll_timeout,
+            callback=callback,
+        )
+
+    async def connect(self, base_url: str):
+        await self._worker.connect(base_url)
+
+    async def shutdown(self):
+        await self._worker.shutdown()
+
+    async def cancel_current(self):
+        """Cancel the in-flight request by terminating the underlying worker.
+
+        Delegates to the active backend (subprocess or multiproc). Called by
+        the CLI SIGINT handler so Ctrl+C interrupts the actual model request,
+        not just the asyncio task wrapping it.
+        """
+        try:
+            await self._worker.cancel_current()
+        except Exception as e:
+            logger.warning("WorkerManager.cancel_current error: %s", e)
+
+    @property
+    def is_alive(self) -> bool:
+        return self._worker.is_alive

package/core/worker_multiproc.py

@@ -0,0 +1,155 @@
+import asyncio
+import logging
+import multiprocessing
+import os
+import signal
+import sys
+import time
+from typing import Callable, Optional
+
+from core.request_worker import worker_main
+from core.worker_protocol import (
+    MSG_RESULT, MSG_UPDATE, MSG_ERROR, MSG_CANCELLED,
+)
+
+logger = logging.getLogger(__name__)
+
+HEARTBEAT_INTERVAL = 5.0
+HEARTBEAT_TIMEOUT = 12.0
+MAX_RESPAWN_DELAY = 30.0
+
+
+class MultiprocessWorker:
+    """Approach 2: Multiprocessing.Process + Queue worker.
+    Uses loop.run_in_executor for non-blocking result reads.
+    """
+
+    def __init__(self):
+        self.task_queue: multiprocessing.Queue = multiprocessing.Queue()
+        self.result_queue: multiprocessing.Queue = multiprocessing.Queue()
+        self.process: Optional[multiprocessing.Process] = None
+        self._running = False
+        self._spawn_count = 0
+        self._closed = False
+
+    async def start(self):
+        if self._running:
+            return
+        self._running = True
+        self._spawn_worker()
+
+    def _spawn_worker(self):
+        if self.process and self.process.is_alive():
+            self._kill_worker()
+
+        self._spawn_count += 1
+        self.process = multiprocessing.Process(
+            target=worker_main,
+            args=(self.task_queue, self.result_queue),
+            daemon=True,
+            name=f"request-worker-{self._spawn_count}",
+        )
+        self.process.start()
+        logger.info(f"Spawned multiproc worker (PID {self.process.pid}, spawn #{self._spawn_count})")
+
+    def _kill_worker(self):
+        if self.process and self.process.is_alive():
+            try:
+                if sys.platform == "win32":
+                    self.process.terminate()
+                else:
+                    os.kill(self.process.pid, signal.SIGKILL)
+                self.process.join(timeout=5)
+            except Exception as e:
+                logger.warning(f"Error killing worker: {e}")
+            finally:
+                self.process = None
+
+    def _on_worker_dead(self):
+        if self._closed:
+            return
+        logger.warning("Multiproc worker dead, respawning...")
+        self._kill_worker()
+        self._spawn_worker()
+
+    async def send_query(
+        self,
+        session_id: str,
+        payload: dict,
+        base_url: str,
+        poll_interval: float = 2.0,
+        poll_timeout: float = 600.0,
+        callback: Optional[Callable] = None,
+    ) -> str:
+        task = {
+            "type": "query",
+            "session_id": session_id,
+            "payload": payload,
+            "base_url": base_url,
+            "poll_interval": poll_interval,
+            "poll_timeout": poll_timeout,
+        }
+        self.task_queue.put(task)
+
+        loop = asyncio.get_event_loop()
+        full_text = ""
+        deadline = time.monotonic() + poll_timeout
+
+        while time.monotonic() < deadline:
+            try:
+                result = await loop.run_in_executor(
+                    None, self.result_queue.get, True, 5.0
+                )
+            except Exception:
+                if not self.process or not self.process.is_alive():
+                    return "Error: Worker process died"
+                continue
+
+            if not isinstance(result, dict):
+                continue
+
+            msg_type = result.get("type")
+            if msg_type == MSG_UPDATE:
+                diff = result.get("content", "")
+                full_text = result.get("full_text", full_text + diff)
+                if callback and diff:
+                    callback(full_text)
+            elif msg_type == MSG_RESULT:
+                content = result.get("content", "")
+                if callback and content:
+                    callback(content)
+                return content
+            elif msg_type == MSG_ERROR:
+                return f"Error: {result.get('message', 'Unknown worker error')}"
+            elif msg_type == MSG_CANCELLED:
+                return "Error: Request was cancelled."
+
+        return "Error: Request timed out"
+
+    async def connect(self, base_url: str):
+        self.task_queue.put({"type": "connect", "base_url": base_url})
+
+    async def shutdown(self):
+        self._closed = True
+        self._running = False
+        try:
+            self.task_queue.put({"type": "shutdown"})
+        except Exception:
+            pass
+        self._kill_worker()
+
+    async def cancel_current(self):
+        """Cancel the in-flight request by killing and respawning the multiproc worker.
+
+        Mirrors AsyncSubprocessWorker.cancel_current() for the multiprocessing
+        backend so the SIGINT handler in CLI works uniformly.
+        """
+        logger.info("Cancelling in-flight request by killing multiproc worker (PID %s)", self.process.pid if self.process else None)
+        self._kill_worker()
+        # Don't respawn immediately — let the next send_query call it via
+        # _on_worker_dead pattern, or let BotClient ensure a fresh one.
+        self._running = False
+
+    @property
+    def is_alive(self) -> bool:
+        return self.process is not None and self.process.is_alive()

package/core/worker_protocol.py

@@ -0,0 +1,30 @@
+import json
+import logging
+
+logger = logging.getLogger(__name__)
+
+MSG_QUERY = "query"
+MSG_RESULT = "result"
+MSG_UPDATE = "update"
+MSG_ERROR = "error"
+MSG_CANCELLED = "cancelled"
+MSG_CONNECT = "connect"
+MSG_PING = "ping"
+MSG_PONG = "pong"
+MSG_SHUTDOWN = "shutdown"
+
+
+def encode_msg(msg_id: int, msg_type: str, **kwargs) -> str:
+    data = {"id": msg_id, "type": msg_type, **kwargs}
+    return json.dumps(data, default=str) + "\n"
+
+
+def decode_msg(line: str) -> dict | None:
+    line = line.strip().lstrip("\ufeff")
+    if not line:
+        return None
+    try:
+        return json.loads(line)
+    except json.JSONDecodeError as e:
+        logger.warning("Failed to decode worker message: %s", e)
+        return None

package/core/worker_subprocess.py

@@ -0,0 +1,222 @@
+import asyncio
+import json
+import logging
+import os
+import sys
+import time
+from typing import Callable, Optional
+
+from core.worker_protocol import (
+    encode_msg, decode_msg, MSG_QUERY, MSG_RESULT, MSG_UPDATE,
+    MSG_ERROR, MSG_CANCELLED, MSG_CONNECT, MSG_SHUTDOWN,
+)
+
+logger = logging.getLogger(__name__)
+
+WORKER_SCRIPT = os.path.join(os.path.dirname(__file__), "request_worker.py")
+
+
+class AsyncSubprocessWorker:
+    """Approach 1: Spawn worker via asyncio.create_subprocess_exec.
+    Communicates via JSON-RPC over stdin/stdout.
+    """
+
+    def __init__(self):
+        self._process: Optional[asyncio.subprocess.Process] = None
+        self._msg_id = 0
+        self._pending: dict[int, asyncio.Future] = {}
+        self._running = False
+        self._reader_task: Optional[asyncio.Task] = None
+
+    async def start(self):
+        if self._running:
+            return
+        self._running = True
+
+        python = sys.executable
+        self._process = await asyncio.create_subprocess_exec(
+            python, WORKER_SCRIPT, "--mode", "stdio",
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+        )
+        logger.info(f"Subprocess worker started (PID {self._process.pid})")
+
+        self._reader_task = asyncio.create_task(self._read_stdout())
+
+    async def _read_stdout(self):
+        while self._running and self._process and not self._process.stdout.at_eof():
+            line = await self._process.stdout.readline()
+            if not line:
+                break
+            msg = decode_msg(line.decode("utf-8", errors="replace"))
+            if not msg:
+                continue
+
+            msg_id = msg.get("id")
+            if msg_id in self._pending:
+                future = self._pending.pop(msg_id)
+                if not future.done():
+                    future.set_result(msg)
+
+        logger.info("Subprocess worker stdout closed")
+
+    async def send_query(
+        self,
+        session_id: str,
+        payload: dict,
+        base_url: str,
+        poll_interval: float = 2.0,
+        poll_timeout: float = 600.0,
+        callback: Optional[Callable] = None,
+    ) -> str:
+        self._msg_id += 1
+        msg_id = self._msg_id
+        future: asyncio.Future = asyncio.get_event_loop().create_future()
+        self._pending[msg_id] = future
+
+        cmd = {
+            "id": msg_id,
+            "type": MSG_QUERY,
+            "session_id": session_id,
+            "payload": payload,
+            "base_url": base_url,
+            "poll_interval": poll_interval,
+            "poll_timeout": poll_timeout,
+        }
+        await self._write_cmd(cmd)
+
+        full_text = ""
+        deadline = time.monotonic() + poll_timeout
+        while time.monotonic() < deadline:
+            try:
+                result = await asyncio.wait_for(
+                    asyncio.shield(future), timeout=5.0
+                )
+            except asyncio.TimeoutError:
+                if msg_id in self._pending:
+                    future = self._pending[msg_id]
+                    if future.done():
+                        result = future.result()
+                    else:
+                        continue
+                else:
+                    break
+
+            msg_type = result.get("type")
+            if msg_type == MSG_UPDATE:
+                diff = result.get("content", "")
+                full_text = result.get("full_text", full_text + diff)
+                if callback:
+                    callback(full_text)
+                new_future = asyncio.get_event_loop().create_future()
+                self._pending[msg_id] = new_future
+                future = new_future
+            elif msg_type == MSG_RESULT:
+                content = result.get("content", "")
+                if callback:
+                    callback(content)
+                return content
+            elif msg_type == MSG_ERROR:
+                return f"Error: {result.get('message', 'Unknown worker error')}"
+            elif msg_type == MSG_CANCELLED:
+                return "Error: Request was cancelled."
+            else:
+                if msg_id in self._pending:
+                    del self._pending[msg_id]
+                return full_text or "Error: Unexpected response"
+
+        if msg_id in self._pending:
+            del self._pending[msg_id]
+        return full_text or "Error: Request timed out"
+
+    async def connect(self, base_url: str):
+        self._msg_id += 1
+        cmd = {"id": self._msg_id, "type": MSG_CONNECT, "base_url": base_url}
+        await self._write_cmd(cmd)
+
+    async def _write_cmd(self, cmd: dict):
+        line = json.dumps(cmd, default=str) + "\n"
+        if self._process and self._process.stdin:
+            self._process.stdin.write(line.encode("utf-8"))
+            await self._process.stdin.drain()
+
+    async def shutdown(self):
+        self._running = False
+        try:
+            if self._process:
+                # Close stdin first to signal EOF to worker (graceful path)
+                try:
+                    if self._process.stdin and not self._process.stdin.is_closing():
+                        self._process.stdin.close()
+                except Exception:
+                    pass
+                # Send shutdown command (may fail if stdin already closed)
+                try:
+                    await self._write_cmd({"type": MSG_SHUTDOWN})
+                except Exception:
+                    pass
+                try:
+                    await asyncio.wait_for(self._process.wait(), timeout=5.0)
+                except asyncio.TimeoutError:
+                    self._process.kill()
+                    await self._process.wait()
+        except Exception as e:
+            logger.warning(f"Subprocess worker shutdown error: {e}")
+        finally:
+            if self._reader_task:
+                self._reader_task.cancel()
+                try:
+                    await self._reader_task
+                except (asyncio.CancelledError, Exception):
+                    pass
+                self._reader_task = None
+            # Explicitly close all pipe transports to prevent ResourceWarning
+            # when the proactor loop tears them down at GC time. On Windows
+            # these are `_ProactorBasePipeTransport` whose __del__ tries
+            # `fileno()` on a closed pipe and raises ValueError.
+            if self._process is not None:
+                for pipe_name in ('stdin', 'stdout', 'stderr'):
+                    pipe = getattr(self._process, pipe_name, None)
+                    if pipe is not None:
+                        try:
+                            transport = getattr(pipe, '_transport', None) or pipe
+                            if hasattr(transport, 'close'):
+                                transport.close()
+                        except Exception:
+                            pass
+            self._process = None
+
+    async def cancel_current(self):
+        """Cancel the in-flight request by killing and respawning the subprocess.
+
+        Used by SIGINT handler in CLI when user presses Ctrl+C. The pending
+        future in self._pending will never resolve (subprocess is dead), so
+        the awaiter will hit its timeout and propagate the cancellation.
+        """
+        if not self._process or self._process.returncode is not None:
+            return
+        logger.info("Cancelling in-flight request by terminating subprocess (PID %s)", self._process.pid)
+        try:
+            self._process.kill()
+        except Exception as e:
+            logger.warning("Error killing subprocess: %s", e)
+        # Mark all pending futures as cancelled so awaiters unblock immediately
+        for msg_id, future in list(self._pending.items()):
+            if not future.done():
+                future.cancel()
+            self._pending.pop(msg_id, None)
+        # Wait for the process to actually die, then respawn so the next
+        # request gets a fresh worker (the old one is in a bad state).
+        try:
+            await asyncio.wait_for(self._process.wait(), timeout=2.0)
+        except asyncio.TimeoutError:
+            pass
+        self._process = None
+        self._running = False
+        # The next send_query() will trigger a fresh start() via BotClient.ensure_worker()
+
+    @property
+    def is_alive(self) -> bool:
+        return self._process is not None and self._process.returncode is None