comfy-env 0.0.8__py3-none-any.whl → 0.0.14__py3-none-any.whl

comfy_env/runner.py

@@ -1,273 +0,0 @@
-"""
-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()