comfy-env 0.0.8__py3-none-any.whl

comfy_env/runner.py

@@ -0,0 +1,273 @@
+"""
+Generic runner for isolated subprocess execution.
+
+This module is the entry point for subprocess execution. The runner handles
+requests for ANY @isolated class in the environment, importing classes on demand.
+
+Usage (Unix Domain Socket - recommended):
+    python -m comfy_env.runner \
+        --node-dir /path/to/ComfyUI-SAM3DObjects/nodes \
+        --comfyui-base /path/to/ComfyUI \
+        --import-paths ".,../vendor" \
+        --socket /tmp/comfyui-isolation-myenv-12345.sock
+
+Usage (Legacy stdin/stdout):
+    python -m comfy_env.runner \
+        --node-dir /path/to/ComfyUI-SAM3DObjects/nodes \
+        --comfyui-base /path/to/ComfyUI \
+        --import-paths ".,../vendor"
+
+The runner:
+1. Sets COMFYUI_ISOLATION_WORKER=1 (so @isolated decorator becomes no-op)
+2. Adds paths to sys.path
+3. Connects to Unix Domain Socket (or uses stdin/stdout)
+4. Dynamically imports classes as needed (cached)
+5. Calls methods and returns responses
+"""
+
+import os
+import sys
+import json
+import argparse
+import traceback
+import warnings
+import logging
+import importlib
+from typing import Any, Dict, Optional
+
+# Suppress warnings that could interfere with JSON IPC
+warnings.filterwarnings("ignore")
+os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
+logging.disable(logging.WARNING)
+
+# Mark that we're in worker mode - this makes @isolated decorator a no-op
+os.environ["COMFYUI_ISOLATION_WORKER"] = "1"
+
+
+def setup_paths(node_dir: str, comfyui_base: Optional[str], import_paths: Optional[str]):
+    """Setup sys.path for imports."""
+    from pathlib import Path
+
+    node_path = Path(node_dir)
+
+    # Set COMFYUI_BASE env var for stubs to use
+    if comfyui_base:
+        os.environ["COMFYUI_BASE"] = comfyui_base
+
+    # Add comfyui-isolation stubs directory (provides folder_paths, etc.)
+    stubs_dir = Path(__file__).parent / "stubs"
+    sys.path.insert(0, str(stubs_dir))
+
+    # Add import paths
+    if import_paths:
+        for p in import_paths.split(","):
+            p = p.strip()
+            if p:
+                full_path = node_path / p
+                sys.path.insert(0, str(full_path))
+
+    # Add node_dir itself
+    sys.path.insert(0, str(node_path))
+
+
+def serialize_result(obj: Any) -> Any:
+    """Serialize result for JSON transport."""
+    from comfy_env.ipc.protocol import encode_object
+    return encode_object(obj)
+
+
+def deserialize_arg(obj: Any) -> Any:
+    """Deserialize argument from JSON transport."""
+    from comfy_env.ipc.protocol import decode_object
+    return decode_object(obj)
+
+
+# Cache for imported classes and instances
+_class_cache: Dict[str, type] = {}
+_instance_cache: Dict[str, object] = {}
+
+
+def get_instance(module_name: str, class_name: str) -> object:
+    """Get or create an instance of a class."""
+    cache_key = f"{module_name}.{class_name}"
+
+    if cache_key not in _instance_cache:
+        # Import the class if not cached
+        if cache_key not in _class_cache:
+            print(f"[Runner] Importing {class_name} from {module_name}...", file=sys.stderr)
+            module = importlib.import_module(module_name)
+            cls = getattr(module, class_name)
+            _class_cache[cache_key] = cls
+
+        # Create instance
+        cls = _class_cache[cache_key]
+        _instance_cache[cache_key] = cls()
+        print(f"[Runner] Created instance of {class_name}", file=sys.stderr)
+
+    return _instance_cache[cache_key]
+
+
+def run_worker(
+    node_dir: str,
+    comfyui_base: Optional[str],
+    import_paths: Optional[str],
+    socket_path: Optional[str] = None,
+):
+    """
+    Main worker loop - handles JSON-RPC requests via transport.
+
+    Args:
+        node_dir: Path to node package directory
+        comfyui_base: Path to ComfyUI base directory
+        import_paths: Comma-separated import paths
+        socket_path: Unix domain socket path (if None, uses stdin/stdout)
+    """
+    from comfy_env.ipc.transport import UnixSocketTransport, StdioTransport
+
+    # Setup paths first
+    setup_paths(node_dir, comfyui_base, import_paths)
+
+    # Create transport
+    if socket_path:
+        # Unix Domain Socket transport (recommended)
+        print(f"[Runner] Connecting to socket: {socket_path}", file=sys.stderr)
+        transport = UnixSocketTransport.connect(socket_path)
+        use_uds = True
+    else:
+        # Legacy stdin/stdout transport
+        print("[Runner] Using stdin/stdout transport", file=sys.stderr)
+        transport = StdioTransport()
+        use_uds = False
+
+    try:
+        # Send ready signal
+        transport.send({"status": "ready"})
+
+        # Main loop - read requests, execute, respond
+        while True:
+            response = {"jsonrpc": "2.0", "id": None}
+
+            try:
+                request = transport.recv()
+                response["id"] = request.get("id")
+
+                method_name = request.get("method")
+                params = request.get("params", {})
+
+                if method_name == "shutdown":
+                    # Clean shutdown
+                    response["result"] = {"status": "shutdown"}
+                    transport.send(response)
+                    break
+
+                # Get module/class from request
+                module_name = request.get("module")
+                class_name = request.get("class")
+
+                if not module_name or not class_name:
+                    response["error"] = {
+                        "code": -32602,
+                        "message": "Missing 'module' or 'class' in request",
+                    }
+                    transport.send(response)
+                    continue
+
+                # Get or create instance
+                try:
+                    instance = get_instance(module_name, class_name)
+                except Exception as e:
+                    response["error"] = {
+                        "code": -32000,
+                        "message": f"Failed to import {module_name}.{class_name}: {e}",
+                        "data": {"traceback": traceback.format_exc()}
+                    }
+                    transport.send(response)
+                    continue
+
+                # Get the method
+                method = getattr(instance, method_name, None)
+                if method is None:
+                    response["error"] = {
+                        "code": -32601,
+                        "message": f"Method not found: {method_name}",
+                    }
+                    transport.send(response)
+                    continue
+
+                # Deserialize arguments
+                deserialized_params = {}
+                for key, value in params.items():
+                    deserialized_params[key] = deserialize_arg(value)
+
+                # For legacy stdio transport, redirect stdout to stderr during execution
+                # This prevents print() in node code from corrupting JSON protocol
+                # (UDS transport doesn't need this since it uses a separate socket)
+                if not use_uds:
+                    original_stdout = sys.stdout
+                    sys.stdout = sys.stderr
+
+                    # Also redirect at file descriptor level for C libraries
+                    stdout_fd = original_stdout.fileno()
+                    stderr_fd = sys.stderr.fileno()
+                    stdout_fd_copy = os.dup(stdout_fd)
+                    os.dup2(stderr_fd, stdout_fd)
+
+                # Call the method
+                print(f"[Runner] Calling {class_name}.{method_name}...", file=sys.stderr)
+                try:
+                    result = method(**deserialized_params)
+                finally:
+                    if not use_uds:
+                        # Restore file descriptor first, then Python stdout
+                        os.dup2(stdout_fd_copy, stdout_fd)
+                        os.close(stdout_fd_copy)
+                        sys.stdout = original_stdout
+
+                # Serialize result
+                serialized_result = serialize_result(result)
+                response["result"] = serialized_result
+
+                print(f"[Runner] {class_name}.{method_name} completed", file=sys.stderr)
+
+            except ConnectionError as e:
+                # Socket closed - normal shutdown
+                print(f"[Runner] Connection closed: {e}", file=sys.stderr)
+                break
+            except Exception as e:
+                tb = traceback.format_exc()
+                print(f"[Runner] Error: {e}", file=sys.stderr)
+                print(tb, file=sys.stderr)
+                response["error"] = {
+                    "code": -32000,
+                    "message": str(e),
+                    "data": {"traceback": tb}
+                }
+
+            try:
+                transport.send(response)
+            except ConnectionError:
+                break
+
+    finally:
+        transport.close()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Isolated node runner")
+    parser.add_argument("--node-dir", required=True, help="Node package directory")
+    parser.add_argument("--comfyui-base", help="ComfyUI base directory")
+    parser.add_argument("--import-paths", help="Comma-separated import paths")
+    parser.add_argument("--socket", help="Unix domain socket path (if not provided, uses stdin/stdout)")
+
+    args = parser.parse_args()
+
+    run_worker(
+        node_dir=args.node_dir,
+        comfyui_base=args.comfyui_base,
+        import_paths=args.import_paths,
+        socket_path=args.socket,
+    )
+
+
+if __name__ == "__main__":
+    main()

comfy_env/stubs/__init__.py

@@ -0,0 +1 @@
+# ComfyUI stubs for isolated workers

comfy_env/stubs/folder_paths.py

@@ -0,0 +1,57 @@
+"""
+Minimal folder_paths stub for isolated worker processes.
+
+Provides the same interface as ComfyUI's folder_paths module
+without importing any ComfyUI dependencies.
+"""
+
+import os
+from pathlib import Path
+
+_comfyui_base = None
+
+def _find_comfyui_base():
+    """Find ComfyUI base from COMFYUI_BASE env var or by walking up."""
+    global _comfyui_base
+    if _comfyui_base:
+        return _comfyui_base
+
+    # Check env var first
+    if os.environ.get("COMFYUI_BASE"):
+        _comfyui_base = Path(os.environ["COMFYUI_BASE"])
+        return _comfyui_base
+
+    # Walk up from cwd looking for ComfyUI
+    current = Path.cwd().resolve()
+    for _ in range(10):
+        if (current / "main.py").exists() and (current / "comfy").exists():
+            _comfyui_base = current
+            return _comfyui_base
+        current = current.parent
+
+    return None
+
+# Models directory
+@property
+def models_dir():
+    base = _find_comfyui_base()
+    return str(base / "models") if base else None
+
+# Make models_dir work as both attribute and property
+class _ModuleProxy:
+    @property
+    def models_dir(self):
+        base = _find_comfyui_base()
+        return str(base / "models") if base else None
+
+    def get_output_directory(self):
+        base = _find_comfyui_base()
+        return str(base / "output") if base else None
+
+    def get_input_directory(self):
+        base = _find_comfyui_base()
+        return str(base / "input") if base else None
+
+# Replace module with proxy instance
+import sys
+sys.modules[__name__] = _ModuleProxy()

comfy_env/workers/__init__.py

@@ -0,0 +1,49 @@
+"""
+Workers - Simple, explicit process isolation for ComfyUI nodes.
+
+This module provides three isolation tiers:
+
+Tier 1: TorchMPWorker (same Python, fresh CUDA context)
+    - Uses torch.multiprocessing.Queue
+    - Zero-copy tensor transfer via CUDA IPC
+    - ~30ms overhead per call
+    - Use for: Memory isolation, fresh CUDA context
+
+Tier 2: VenvWorker (different Python/venv)
+    - Uses subprocess + torch.save/load via /dev/shm
+    - One memcpy per tensor direction
+    - ~100-500ms overhead per call
+    - Use for: Different PyTorch versions, incompatible deps
+
+Tier 3: ContainerWorker (full isolation) [future]
+    - Docker with GPU passthrough
+    - Use for: Different CUDA versions, hermetic environments
+
+Usage:
+    from comfy_env.workers import get_worker, TorchMPWorker
+
+    # Get a named worker from the pool
+    worker = get_worker("sam3d")
+    result = worker.call(my_function, image=tensor)
+
+    # Or create directly
+    worker = TorchMPWorker()
+    result = worker.call(my_function, arg1, arg2)
+"""
+
+from .base import Worker
+from .torch_mp import TorchMPWorker
+from .venv import VenvWorker, PersistentVenvWorker
+from .pool import WorkerPool, get_worker, register_worker, shutdown_workers, list_workers
+
+__all__ = [
+    "Worker",
+    "TorchMPWorker",
+    "VenvWorker",
+    "PersistentVenvWorker",
+    "WorkerPool",
+    "get_worker",
+    "register_worker",
+    "shutdown_workers",
+    "list_workers",
+]

comfy_env/workers/base.py

@@ -0,0 +1,82 @@
+"""
+Base Worker Interface - Protocol for all worker implementations.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Optional
+
+
+class Worker(ABC):
+    """
+    Abstract base class for process isolation workers.
+
+    All workers must implement:
+    - call(): Execute a function in the isolated process
+    - shutdown(): Clean up resources
+
+    Workers should be used as context managers when possible:
+
+        with TorchMPWorker() as worker:
+            result = worker.call(my_func, arg1, arg2)
+    """
+
+    @abstractmethod
+    def call(
+        self,
+        func: Callable,
+        *args,
+        timeout: Optional[float] = None,
+        **kwargs
+    ) -> Any:
+        """
+        Execute a function in the isolated worker process.
+
+        Args:
+            func: The function to execute. Must be picklable (top-level or staticmethod).
+            *args: Positional arguments passed to func.
+            timeout: Optional timeout in seconds (None = no timeout).
+            **kwargs: Keyword arguments passed to func.
+
+        Returns:
+            The return value of func(*args, **kwargs).
+
+        Raises:
+            TimeoutError: If execution exceeds timeout.
+            RuntimeError: If worker process dies or raises exception.
+        """
+        pass
+
+    @abstractmethod
+    def shutdown(self) -> None:
+        """
+        Shut down the worker and release resources.
+
+        Safe to call multiple times. After shutdown, further calls to
+        call() will raise RuntimeError.
+        """
+        pass
+
+    @abstractmethod
+    def is_alive(self) -> bool:
+        """Check if the worker process is still running."""
+        pass
+
+    def __enter__(self) -> "Worker":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.shutdown()
+
+
+class WorkerError(Exception):
+    """Exception raised when a worker encounters an error."""
+
+    def __init__(self, message: str, traceback: Optional[str] = None):
+        super().__init__(message)
+        self.worker_traceback = traceback
+
+    def __str__(self):
+        msg = super().__str__()
+        if self.worker_traceback:
+            msg += f"\n\nWorker traceback:\n{self.worker_traceback}"
+        return msg

comfy_env/workers/pool.py

@@ -0,0 +1,241 @@
+"""
+WorkerPool - Global registry and management of named workers.
+
+Provides a simple API for getting workers by name:
+
+    from comfy_env.workers import get_worker
+
+    worker = get_worker("sam3d")
+    result = worker.call_module("my_module", "my_func", image=tensor)
+
+Workers are registered at startup and reused across calls:
+
+    from comfy_env.workers import register_worker, TorchMPWorker
+
+    register_worker("default", TorchMPWorker())
+    register_worker("sam3d", PersistentVenvWorker(
+        python="/path/to/venv/bin/python",
+        working_dir="/path/to/nodes",
+    ))
+"""
+
+import atexit
+import threading
+from typing import Dict, Optional, Union
+from pathlib import Path
+
+from .base import Worker
+
+
+class WorkerPool:
+    """
+    Singleton pool of named workers.
+
+    Manages worker lifecycle, provides access by name, handles cleanup.
+    """
+
+    _instance: Optional["WorkerPool"] = None
+    _lock = threading.Lock()
+
+    def __new__(cls):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+                    cls._instance._initialized = False
+        return cls._instance
+
+    def __init__(self):
+        if self._initialized:
+            return
+        self._initialized = True
+        self._workers: Dict[str, Worker] = {}
+        self._factories: Dict[str, callable] = {}
+        self._worker_lock = threading.Lock()
+
+    def register(
+        self,
+        name: str,
+        worker: Optional[Worker] = None,
+        factory: Optional[callable] = None,
+    ) -> None:
+        """
+        Register a worker or worker factory.
+
+        Args:
+            name: Name to register under.
+            worker: Pre-created worker instance.
+            factory: Callable that creates worker on first use (lazy).
+
+        Only one of worker or factory should be provided.
+        """
+        if worker is not None and factory is not None:
+            raise ValueError("Provide either worker or factory, not both")
+        if worker is None and factory is None:
+            raise ValueError("Must provide worker or factory")
+
+        with self._worker_lock:
+            # Shutdown existing worker if replacing
+            if name in self._workers:
+                try:
+                    self._workers[name].shutdown()
+                except:
+                    pass
+
+            if worker is not None:
+                self._workers[name] = worker
+                self._factories.pop(name, None)
+            else:
+                self._factories[name] = factory
+                self._workers.pop(name, None)
+
+    def get(self, name: str) -> Worker:
+        """
+        Get a worker by name.
+
+        Args:
+            name: Registered worker name.
+
+        Returns:
+            The worker instance.
+
+        Raises:
+            KeyError: If no worker registered with that name.
+        """
+        with self._worker_lock:
+            # Check for existing worker
+            if name in self._workers:
+                worker = self._workers[name]
+                if worker.is_alive():
+                    return worker
+                # Worker died, try to recreate from factory
+                if name not in self._factories:
+                    raise RuntimeError(f"Worker '{name}' died and no factory to recreate")
+
+            # Create from factory
+            if name in self._factories:
+                worker = self._factories[name]()
+                self._workers[name] = worker
+                return worker
+
+            raise KeyError(f"No worker registered with name: {name}")
+
+    def shutdown(self, name: Optional[str] = None) -> None:
+        """
+        Shutdown workers.
+
+        Args:
+            name: If provided, shutdown only this worker.
+                  If None, shutdown all workers.
+        """
+        with self._worker_lock:
+            if name is not None:
+                if name in self._workers:
+                    try:
+                        self._workers[name].shutdown()
+                    except:
+                        pass
+                    del self._workers[name]
+            else:
+                for worker in self._workers.values():
+                    try:
+                        worker.shutdown()
+                    except:
+                        pass
+                self._workers.clear()
+
+    def list_workers(self) -> Dict[str, str]:
+        """
+        List all registered workers.
+
+        Returns:
+            Dict of name -> status string.
+        """
+        with self._worker_lock:
+            result = {}
+            for name, worker in self._workers.items():
+                status = "alive" if worker.is_alive() else "dead"
+                result[name] = f"{type(worker).__name__} ({status})"
+            for name in self._factories:
+                if name not in result:
+                    result[name] = f"factory (not started)"
+            return result
+
+
+# Global pool instance
+_pool = WorkerPool()
+
+
+def get_worker(name: str) -> Worker:
+    """
+    Get a worker by name from the global pool.
+
+    Args:
+        name: Registered worker name.
+
+    Returns:
+        Worker instance.
+
+    Example:
+        worker = get_worker("sam3d")
+        result = worker.call_module("my_module", "my_func", image=tensor)
+    """
+    return _pool.get(name)
+
+
+def register_worker(
+    name: str,
+    worker: Optional[Worker] = None,
+    factory: Optional[callable] = None,
+) -> None:
+    """
+    Register a worker in the global pool.
+
+    Args:
+        name: Name to register under.
+        worker: Pre-created worker instance.
+        factory: Callable that creates worker on demand.
+
+    Example:
+        # Register pre-created worker
+        register_worker("default", TorchMPWorker())
+
+        # Register factory for lazy creation
+        register_worker("sam3d", factory=lambda: PersistentVenvWorker(
+            python="/path/to/venv/bin/python",
+        ))
+    """
+    _pool.register(name, worker=worker, factory=factory)
+
+
+def shutdown_workers(name: Optional[str] = None) -> None:
+    """
+    Shutdown workers in the global pool.
+
+    Args:
+        name: If provided, shutdown only this worker.
+              If None, shutdown all workers.
+    """
+    _pool.shutdown(name)
+
+
+def list_workers() -> Dict[str, str]:
+    """
+    List all registered workers.
+
+    Returns:
+        Dict of name -> status description.
+    """
+    return _pool.list_workers()
+
+
+# Register default worker (TorchMPWorker) on import
+def _register_default():
+    from .torch_mp import TorchMPWorker
+    register_worker("default", factory=lambda: TorchMPWorker(name="default"))
+
+
+_register_default()
+
+# Cleanup on exit
+atexit.register(lambda: shutdown_workers())