comfy-env 0.0.64__py3-none-any.whl → 0.0.66__py3-none-any.whl

comfy_env/ipc/bridge.py

@@ -1,476 +0,0 @@
-"""
-WorkerBridge - Main IPC class for communicating with isolated workers.
-
-This is the primary interface that ComfyUI node developers use.
-"""
-
-import json
-import subprocess
-import threading
-import uuid
-from pathlib import Path
-from typing import Any, Callable, Dict, Optional
-
-from ..env.config import IsolatedEnv
-from ..env.manager import IsolatedEnvManager
-from .protocol import encode_object, decode_object
-
-
-class WorkerBridge:
-    """
-    Bridge for communicating with a worker process in an isolated environment.
-
-    This class manages the worker process lifecycle and handles IPC.
-
-    Features:
-    - Lazy worker startup (starts on first call)
-    - Singleton pattern (one worker per environment)
-    - Auto-restart on crash
-    - Graceful shutdown
-    - Timeout support
-
-    Example:
-        from comfy_env import IsolatedEnv, WorkerBridge
-
-        env = IsolatedEnv(
-            name="my-node",
-            python="3.10",
-            cuda="12.8",
-            requirements=["torch==2.8.0"],
-        )
-
-        bridge = WorkerBridge(env, worker_script=Path("worker.py"))
-
-        # Direct method calls (preferred)
-        result = bridge.process(image=my_image)
-
-        # Or use explicit call() method
-        result = bridge.call("process", image=my_image)
-    """
-
-    def __init__(
-        self,
-        env: IsolatedEnv,
-        worker_script: Path,
-        base_dir: Optional[Path] = None,
-        log_callback: Optional[Callable[[str], None]] = None,
-        auto_start: bool = True,
-    ):
-        """
-        Initialize the bridge.
-
-        Args:
-            env: Isolated environment configuration
-            worker_script: Path to the worker Python script
-            base_dir: Base directory for environments (default: worker_script's parent)
-            log_callback: Optional callback for logging (default: print)
-            auto_start: Whether to auto-start worker on first call (default: True)
-        """
-        self.env = env
-        self.worker_script = Path(worker_script)
-        self.base_dir = base_dir or self.worker_script.parent
-        self.log = log_callback or print
-        self.auto_start = auto_start
-
-        self._manager = IsolatedEnvManager(self.base_dir, log_callback=log_callback)
-        self._process: Optional[subprocess.Popen] = None
-        self._process_lock = threading.Lock()
-        self._stderr_thread: Optional[threading.Thread] = None
-
-    @classmethod
-    def from_config_file(
-        cls,
-        node_dir: Path,
-        worker_script: Optional[Path] = None,
-        config_file: Optional[str] = None,
-        log_callback: Optional[Callable[[str], None]] = None,
-        auto_start: bool = True,
-    ) -> "WorkerBridge":
-        """
-        Create WorkerBridge from a config file.
-
-        This is a convenience method for loading environment configuration
-        from a TOML file instead of defining it programmatically.
-
-        Args:
-            node_dir: Directory containing the config file (and base_dir for env)
-            worker_script: Path to worker script (optional - auto-discovered from config or convention)
-            config_file: Specific config file name (default: auto-discover)
-            log_callback: Optional callback for logging
-            auto_start: Whether to auto-start worker on first call (default: True)
-
-        Returns:
-            Configured WorkerBridge instance
-
-        Raises:
-            FileNotFoundError: If no config file found or worker script not found
-            ImportError: If tomli not installed (Python < 3.11)
-
-        Example:
-            # Auto-discover everything (recommended)
-            bridge = WorkerBridge.from_config_file(node_dir=Path(__file__).parent)
-
-            # Explicit worker script (backward compatible)
-            bridge = WorkerBridge.from_config_file(
-                node_dir=Path(__file__).parent,
-                worker_script=Path(__file__).parent / "worker.py",
-            )
-        """
-        from ..env.config_file import load_env_from_file, discover_env_config, CONFIG_FILE_NAMES
-
-        node_dir = Path(node_dir)
-
-        if config_file:
-            env = load_env_from_file(node_dir / config_file, node_dir)
-        else:
-            env = discover_env_config(node_dir)
-            if env is None:
-                raise FileNotFoundError(
-                    f"No isolation config found in {node_dir}. "
-                    f"Create one of: {', '.join(CONFIG_FILE_NAMES)}"
-                )
-
-        # Resolve worker script from explicit param, config, or convention
-        if worker_script is None:
-            worker_script = cls._resolve_worker_script(node_dir, env)
-
-        return cls(
-            env=env,
-            worker_script=worker_script,
-            base_dir=node_dir,
-            log_callback=log_callback,
-            auto_start=auto_start,
-        )
-
-    @staticmethod
-    def _resolve_worker_script(node_dir: Path, env: "IsolatedEnv") -> Path:
-        """
-        Resolve worker script from config or convention.
-
-        Discovery order:
-        1. Explicit package in config ([worker] package = "worker")
-        2. Explicit script in config ([worker] script = "worker.py")
-        3. Convention: worker/ directory with __main__.py
-        4. Convention: worker.py file
-
-        Args:
-            node_dir: Node directory
-            env: IsolatedEnv with optional worker_package/worker_script
-
-        Returns:
-            Path to worker script
-
-        Raises:
-            FileNotFoundError: If no worker found
-        """
-        # 1. Explicit package in config
-        if env.worker_package:
-            pkg_dir = node_dir / env.worker_package
-            main_py = pkg_dir / "__main__.py"
-            if main_py.exists():
-                return main_py
-            raise FileNotFoundError(
-                f"Worker package '{env.worker_package}' not found. "
-                f"Expected: {main_py}"
-            )
-
-        # 2. Explicit script in config
-        if env.worker_script:
-            script = node_dir / env.worker_script
-            if script.exists():
-                return script
-            raise FileNotFoundError(
-                f"Worker script '{env.worker_script}' not found. "
-                f"Expected: {script}"
-            )
-
-        # 3. Convention: worker/ directory
-        worker_pkg = node_dir / "worker" / "__main__.py"
-        if worker_pkg.exists():
-            return worker_pkg
-
-        # 4. Convention: worker.py file
-        worker_file = node_dir / "worker.py"
-        if worker_file.exists():
-            return worker_file
-
-        raise FileNotFoundError(
-            f"No worker found in {node_dir}. "
-            f"Create worker/__main__.py, worker.py, or specify in config:\n"
-            f"  [worker]\n"
-            f"  package = \"worker\"  # or\n"
-            f"  script = \"worker.py\""
-        )
-
-    @property
-    def python_exe(self) -> Path:
-        """Get the Python executable path for the isolated environment."""
-        return self._manager.get_python(self.env)
-
-    @property
-    def is_running(self) -> bool:
-        """Check if worker process is currently running."""
-        with self._process_lock:
-            return self._process is not None and self._process.poll() is None
-
-    def ensure_environment(self, verify_packages: Optional[list] = None) -> None:
-        """
-        Ensure the isolated environment exists and is ready.
-
-        Args:
-            verify_packages: Optional list of packages to verify
-        """
-        self._manager.setup(self.env, verify_packages=verify_packages)
-
-    def start(self) -> None:
-        """
-        Start the worker process.
-
-        Does nothing if worker is already running.
-        """
-        with self._process_lock:
-            if self._process is not None and self._process.poll() is None:
-                return  # Already running
-
-            python_exe = self.python_exe
-            if not python_exe.exists():
-                raise RuntimeError(
-                    f"Python executable not found: {python_exe}\n"
-                    f"Run ensure_environment() first or check your env configuration."
-                )
-
-            if not self.worker_script.exists():
-                raise RuntimeError(f"Worker script not found: {self.worker_script}")
-
-            self.log(f"Starting worker process...")
-            self.log(f"  Python: {python_exe}")
-            self.log(f"  Script: {self.worker_script}")
-
-            self._process = subprocess.Popen(
-                [str(python_exe), str(self.worker_script)],
-                stdin=subprocess.PIPE,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                text=True,
-                bufsize=1,  # Line buffered
-            )
-
-            # Collect stderr for error reporting
-            self._stderr_lines = []
-
-            # Start stderr reader thread
-            self._stderr_thread = threading.Thread(
-                target=self._read_stderr,
-                daemon=True,
-                name=f"worker-stderr-{self.env.name}",
-            )
-            self._stderr_thread.start()
-
-            # Give the process a moment to crash if it's going to (e.g., import error)
-            import time
-            time.sleep(0.5)
-
-            # Check if process died immediately
-            if self._process.poll() is not None:
-                exit_code = self._process.returncode
-                time.sleep(0.2)  # Let stderr thread collect output
-                stderr_output = "\n".join(self._stderr_lines[-20:])  # Last 20 lines
-                raise RuntimeError(
-                    f"Worker crashed on startup (exit code {exit_code}).\n"
-                    f"Stderr output:\n{stderr_output}"
-                )
-
-            # Test connection
-            try:
-                response = self._send_raw({"method": "ping"}, timeout=30.0)
-                if response.get("result") != "pong":
-                    raise RuntimeError(f"Worker ping failed: {response}")
-                self.log("Worker started successfully")
-            except Exception as e:
-                # Collect any stderr output for debugging
-                time.sleep(0.2)
-                stderr_output = "\n".join(self._stderr_lines[-20:])
-                self.stop()
-                raise RuntimeError(f"Worker failed to start: {e}\nStderr:\n{stderr_output}")
-
-    def _read_stderr(self) -> None:
-        """Read stderr from worker and forward to log callback."""
-        if not self._process or not self._process.stderr:
-            return
-
-        for line in self._process.stderr:
-            line = line.rstrip()
-            if line:
-                self.log(line)
-                # Also collect for error reporting
-                if hasattr(self, '_stderr_lines'):
-                    self._stderr_lines.append(line)
-
-    def stop(self) -> None:
-        """
-        Stop the worker process gracefully.
-        """
-        with self._process_lock:
-            if self._process is None or self._process.poll() is not None:
-                return
-
-            self.log("Stopping worker...")
-
-            # Send shutdown command
-            try:
-                self._send_raw({"method": "shutdown"}, timeout=5.0)
-            except Exception:
-                pass
-
-            # Wait for graceful shutdown
-            try:
-                self._process.wait(timeout=5.0)
-            except subprocess.TimeoutExpired:
-                self.log("Worker didn't stop gracefully, terminating...")
-                self._process.terminate()
-                try:
-                    self._process.wait(timeout=2.0)
-                except subprocess.TimeoutExpired:
-                    self._process.kill()
-
-            self._process = None
-            self.log("Worker stopped")
-
-    def _send_raw(self, request: dict, timeout: float = 300.0) -> dict:
-        """
-        Send a raw request and wait for response.
-
-        Args:
-            request: Request dict
-            timeout: Timeout in seconds
-
-        Returns:
-            Response dict
-        """
-        if self._process is None or self._process.poll() is not None:
-            raise RuntimeError("Worker process is not running")
-
-        # Add request ID
-        if "id" not in request:
-            request["id"] = str(uuid.uuid4())[:8]
-
-        # Send request
-        request_json = json.dumps(request) + "\n"
-        self._process.stdin.write(request_json)
-        self._process.stdin.flush()
-
-        # Read response
-        import time
-        start_time = time.time()
-
-        while True:
-            if time.time() - start_time > timeout:
-                raise TimeoutError(f"Timeout waiting for response after {timeout}s")
-
-            response_line = self._process.stdout.readline()
-            if not response_line:
-                raise RuntimeError("Worker process closed unexpectedly")
-
-            response_line = response_line.strip()
-            if not response_line:
-                continue
-
-            # Skip non-JSON lines (library output that escaped)
-            if response_line.startswith('[') and not response_line.startswith('[{'):
-                continue
-
-            try:
-                return json.loads(response_line)
-            except json.JSONDecodeError:
-                continue
-
-    def call(
-        self,
-        method: str,
-        timeout: float = 300.0,
-        **kwargs,
-    ) -> Any:
-        """
-        Call a method on the worker.
-
-        Args:
-            method: Method name to call
-            timeout: Timeout in seconds (default: 5 minutes)
-            **kwargs: Arguments to pass to the method
-
-        Returns:
-            The method's return value
-
-        Raises:
-            RuntimeError: If worker returns an error
-            TimeoutError: If call times out
-        """
-        # Auto-start if needed
-        if self.auto_start and not self.is_running:
-            self.start()
-
-        # Encode arguments
-        encoded_args = encode_object(kwargs)
-
-        # Send request
-        request = {
-            "method": method,
-            "args": encoded_args,
-        }
-        response = self._send_raw(request, timeout=timeout)
-
-        # Check for error
-        if "error" in response and response["error"]:
-            error_msg = response["error"]
-            tb = response.get("traceback", "")
-            raise RuntimeError(f"Worker error: {error_msg}\n{tb}")
-
-        # Decode and return result
-        return decode_object(response.get("result"))
-
-    def __enter__(self) -> "WorkerBridge":
-        """Context manager entry."""
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
-        """Context manager exit - stops worker."""
-        self.stop()
-
-    def __del__(self) -> None:
-        """Destructor - stops worker."""
-        try:
-            self.stop()
-        except Exception:
-            pass
-
-    def __getattr__(self, name: str) -> Callable[..., Any]:
-        """
-        Enable direct method calls on the bridge.
-
-        Instead of:
-            result = bridge.call("inference", image=img)
-
-        You can do:
-            result = bridge.inference(image=img)
-
-        Args:
-            name: Method name to call on the worker
-
-        Returns:
-            A callable that forwards to self.call()
-        """
-        # Don't intercept private/magic attributes
-        if name.startswith("_"):
-            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
-
-        # Return a wrapper that calls the method on the worker
-        def method_wrapper(*args, timeout: float = 300.0, **kwargs) -> Any:
-            # If positional args provided, raise helpful error
-            if args:
-                raise TypeError(
-                    f"bridge.{name}() doesn't accept positional arguments. "
-                    f"Use keyword arguments: bridge.{name}(arg1=val1, arg2=val2)"
-                )
-            return self.call(name, timeout=timeout, **kwargs)
-
-        return method_wrapper