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

comfy_env/ipc/transport.py

@@ -1,318 +0,0 @@
-"""
-Transport Layer - Pluggable IPC transports for host-worker communication.
-
-Supports:
-- QueueTransport: multiprocessing.Queue with zero-copy tensor support (recommended)
-- UnixSocketTransport: Unix Domain Sockets with length-prefixed JSON (Linux/macOS only)
-- StdioTransport: Legacy stdin/stdout JSON lines (fallback)
-
-The transport abstraction allows swapping IPC mechanisms without changing
-the serialization protocol.
-"""
-
-import json
-import os
-import socket
-import struct
-import sys
-import threading
-from abc import ABC, abstractmethod
-from pathlib import Path
-from typing import Any, Optional, Protocol, Tuple, runtime_checkable
-
-
-@runtime_checkable
-class Transport(Protocol):
-    """Protocol for IPC transport mechanisms."""
-
-    def send(self, obj: Any) -> None:
-        """Send a JSON-serializable object to the remote endpoint."""
-        ...
-
-    def recv(self) -> Any:
-        """Receive a JSON object from the remote endpoint. Blocks until available."""
-        ...
-
-    def close(self) -> None:
-        """Close the transport. Further send/recv calls may fail."""
-        ...
-
-
-def create_queue_pair(share_torch: bool = True):
-    """
-    Create a pair of queues for bidirectional IPC.
-
-    When share_torch=True, uses torch.multiprocessing for zero-copy tensor transfer.
-    This works cross-platform (Windows, Linux, macOS).
-
-    Args:
-        share_torch: If True, use torch.multiprocessing for zero-copy tensors.
-
-    Returns:
-        Tuple of (to_worker_queue, from_worker_queue, multiprocessing_module)
-    """
-    if share_torch:
-        try:
-            import torch.multiprocessing as mp
-        except ImportError:
-            import multiprocessing as mp
-    else:
-        import multiprocessing as mp
-
-    # Ensure spawn method for proper isolation
-    start_method = mp.get_start_method(allow_none=True)
-    if start_method is None:
-        mp.set_start_method("spawn")
-
-    to_worker = mp.Queue()
-    from_worker = mp.Queue()
-    return to_worker, from_worker, mp
-
-
-class QueueTransport:
-    """
-    Cross-platform transport using multiprocessing.Queue.
-
-    This is the recommended transport as it:
-    - Works on Windows, Linux, and macOS (no AF_UNIX dependency)
-    - Supports zero-copy tensor transfer when using torch.multiprocessing
-    - Has simpler code than socket-based transports
-    - Is battle-tested (used by pyisolate)
-
-    When share_torch=True, PyTorch tensors are transferred via:
-    - CPU tensors: shared memory (share_memory_())
-    - GPU tensors: CUDA IPC handles (zero-copy)
-    """
-
-    def __init__(self, send_queue, recv_queue):
-        """
-        Initialize with send and receive queues.
-
-        Args:
-            send_queue: Queue for sending messages (to remote)
-            recv_queue: Queue for receiving messages (from remote)
-        """
-        self._send = send_queue
-        self._recv = recv_queue
-
-    def send(self, obj: Any) -> None:
-        """Send an object to the remote endpoint."""
-        self._send.put(obj)
-
-    def recv(self, timeout: Optional[float] = None) -> Any:
-        """
-        Receive an object from the remote endpoint.
-
-        Args:
-            timeout: Timeout in seconds (None = block forever)
-
-        Returns:
-            Received object
-
-        Raises:
-            Empty: If timeout expires with no message
-        """
-        import queue
-        try:
-            return self._recv.get(timeout=timeout)
-        except queue.Empty:
-            raise
-
-    def close(self) -> None:
-        """Close the transport (no-op for queues, cleanup handled by creator)."""
-        pass
-
-
-class UnixSocketTransport:
-    """
-    Transport using Unix Domain Sockets with length-prefixed JSON messages.
-
-    This is the recommended transport as it:
-    - Doesn't interfere with stdout/stderr (no C library output issues)
-    - Supports binary-safe length-prefixed framing
-    - Is more efficient than line-based JSON
-
-    Message format: [4-byte big-endian length][JSON payload]
-    """
-
-    def __init__(self, sock: socket.socket):
-        """
-        Initialize with an already-connected socket.
-
-        Args:
-            sock: Connected Unix domain socket
-        """
-        self._sock = sock
-        self._send_lock = threading.Lock()
-        self._recv_lock = threading.Lock()
-
-    @classmethod
-    def create_server(cls, sock_path: str) -> tuple[socket.socket, "UnixSocketTransport"]:
-        """
-        Create a server socket, wait for one connection, return transport.
-
-        Args:
-            sock_path: Path for the Unix domain socket
-
-        Returns:
-            Tuple of (server_socket, transport) - caller should close server_socket
-        """
-        # Remove existing socket file if present
-        if os.path.exists(sock_path):
-            os.unlink(sock_path)
-
-        server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-        server.bind(sock_path)
-        server.listen(1)
-
-        conn, _ = server.accept()
-        return server, cls(conn)
-
-    @classmethod
-    def connect(cls, sock_path: str, timeout: float = 30.0) -> "UnixSocketTransport":
-        """
-        Connect to an existing Unix domain socket.
-
-        Args:
-            sock_path: Path to the Unix domain socket
-            timeout: Connection timeout in seconds
-
-        Returns:
-            Connected transport
-        """
-        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-        sock.settimeout(timeout)
-        sock.connect(sock_path)
-        sock.settimeout(None)  # Back to blocking mode
-        return cls(sock)
-
-    def send(self, obj: Any) -> None:
-        """Send a JSON-serializable object with length prefix."""
-        data = json.dumps(obj).encode('utf-8')
-        header = struct.pack('>I', len(data))
-
-        with self._send_lock:
-            self._sock.sendall(header + data)
-
-    def recv(self) -> Any:
-        """Receive a length-prefixed JSON message."""
-        with self._recv_lock:
-            # Read 4-byte length header
-            header = self._recvall(4)
-            if not header or len(header) < 4:
-                raise ConnectionError("Socket closed or incomplete header")
-
-            msg_len = struct.unpack('>I', header)[0]
-
-            # Sanity check - 100MB limit
-            if msg_len > 100 * 1024 * 1024:
-                raise ValueError(f"Message too large: {msg_len} bytes")
-
-            # Read payload
-            data = self._recvall(msg_len)
-            if len(data) < msg_len:
-                raise ConnectionError(f"Incomplete message: {len(data)}/{msg_len} bytes")
-
-            return json.loads(data.decode('utf-8'))
-
-    def _recvall(self, n: int) -> bytes:
-        """Receive exactly n bytes from the socket."""
-        chunks = []
-        remaining = n
-        while remaining > 0:
-            chunk = self._sock.recv(min(remaining, 65536))
-            if not chunk:
-                break
-            chunks.append(chunk)
-            remaining -= len(chunk)
-        return b''.join(chunks)
-
-    def close(self) -> None:
-        """Close the socket."""
-        try:
-            self._sock.close()
-        except Exception:
-            pass
-
-    def fileno(self) -> int:
-        """Return socket file descriptor for select()."""
-        return self._sock.fileno()
-
-
-class StdioTransport:
-    """
-    Legacy transport using stdin/stdout with JSON lines.
-
-    This transport has issues with C libraries that print to stdout,
-    which is why UnixSocketTransport is preferred. Kept for compatibility.
-
-    Note: Requires fd-level redirection during method execution to prevent
-    C library output from corrupting the JSON stream.
-    """
-
-    def __init__(
-        self,
-        stdin: Any = None,
-        stdout: Any = None,
-    ):
-        """
-        Initialize with file handles.
-
-        Args:
-            stdin: Input stream (default: sys.stdin)
-            stdout: Output stream (default: sys.stdout)
-        """
-        self._stdin = stdin or sys.stdin
-        self._stdout = stdout or sys.stdout
-        self._send_lock = threading.Lock()
-
-    def send(self, obj: Any) -> None:
-        """Send a JSON object as a single line."""
-        with self._send_lock:
-            line = json.dumps(obj) + '\n'
-            self._stdout.write(line)
-            self._stdout.flush()
-
-    def recv(self) -> Any:
-        """Receive a JSON object from a single line."""
-        line = self._stdin.readline()
-        if not line:
-            raise ConnectionError("stdin closed")
-        return json.loads(line.strip())
-
-    def close(self) -> None:
-        """No-op for stdio transport."""
-        pass
-
-
-def get_socket_path(env_name: str, pid: Optional[int] = None) -> str:
-    """
-    Generate a unique socket path for an isolated environment.
-
-    Args:
-        env_name: Name of the isolated environment
-        pid: Process ID (default: current process)
-
-    Returns:
-        Path string for Unix domain socket
-    """
-    if pid is None:
-        pid = os.getpid()
-
-    # Use /tmp on Linux, or a temp directory on other platforms
-    if sys.platform == 'linux':
-        base = '/tmp'
-    else:
-        import tempfile
-        base = tempfile.gettempdir()
-
-    return str(Path(base) / f"comfyui-isolation-{env_name}-{pid}.sock")
-
-
-def cleanup_socket(sock_path: str) -> None:
-    """Remove a socket file if it exists."""
-    try:
-        if os.path.exists(sock_path):
-            os.unlink(sock_path)
-    except OSError:
-        pass

comfy_env/ipc/worker.py

@@ -1,221 +0,0 @@
-"""
-BaseWorker - Base class for worker scripts that run in isolated environments.
-
-Node developers extend this class to define their worker's functionality.
-"""
-
-import sys
-import json
-import traceback
-from typing import Any, Callable, Dict, Optional
-from functools import wraps
-
-from .protocol import encode_object, decode_object
-
-
-# Global registry of methods
-_method_registry: Dict[str, Callable] = {}
-
-
-def register(name: Optional[str] = None):
-    """
-    Decorator to register a method as callable from the bridge.
-
-    Args:
-        name: Optional method name (defaults to function name)
-
-    Example:
-        class MyWorker(BaseWorker):
-            @register("process")
-            def process_image(self, image):
-                return processed_image
-
-            @register()  # Uses function name "do_something"
-            def do_something(self, x):
-                return x * 2
-    """
-    def decorator(func: Callable) -> Callable:
-        method_name = name if name else func.__name__
-        _method_registry[method_name] = func
-
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            return func(*args, **kwargs)
-        return wrapper
-
-    return decorator
-
-
-class BaseWorker:
-    """
-    Base class for isolated worker processes.
-
-    Subclass this to create a worker that handles requests from WorkerBridge.
-
-    The worker runs a main loop that:
-    1. Reads JSON requests from stdin
-    2. Dispatches to registered methods
-    3. Writes JSON responses to stdout
-
-    Example:
-        class MyWorker(BaseWorker):
-            def setup(self):
-                # Called once when worker starts
-                import torch
-                self.model = load_my_model()
-
-            @register("inference")
-            def run_inference(self, image, params):
-                return self.model(image, **params)
-
-        if __name__ == "__main__":
-            MyWorker().run()
-    """
-
-    def __init__(self):
-        """Initialize worker."""
-        self._methods: Dict[str, Callable] = {}
-        self._running = True
-
-        # Register decorated methods
-        for method_name, func in _method_registry.items():
-            # Bind method to this instance
-            self._methods[method_name] = lambda *args, f=func, **kwargs: f(self, *args, **kwargs)
-
-        # Also check instance methods decorated with @register
-        for name in dir(self):
-            method = getattr(self, name)
-            if hasattr(method, '__wrapped__') and name in _method_registry:
-                self._methods[name] = method
-
-    def setup(self) -> None:
-        """
-        Called once when the worker starts, before processing any requests.
-
-        Override this to load models, initialize state, etc.
-        """
-        pass
-
-    def teardown(self) -> None:
-        """
-        Called when the worker is shutting down.
-
-        Override this to cleanup resources.
-        """
-        pass
-
-    def log(self, message: str) -> None:
-        """
-        Log a message to stderr (visible in main process).
-
-        Args:
-            message: Message to log
-        """
-        print(f"[Worker] {message}", file=sys.stderr, flush=True)
-
-    def handle_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Handle a single request.
-
-        Args:
-            request: Request dict with 'method' and 'args' keys
-
-        Returns:
-            Response dict with 'result' or 'error' key
-        """
-        method_name = request.get("method")
-        args = request.get("args", {})
-        request_id = request.get("id", "unknown")
-
-        # Handle built-in commands
-        if method_name == "ping":
-            return {"id": request_id, "result": "pong"}
-
-        if method_name == "shutdown":
-            self._running = False
-            return {"id": request_id, "result": "shutting_down"}
-
-        if method_name == "list_methods":
-            return {"id": request_id, "result": list(self._methods.keys())}
-
-        # Find and call registered method
-        if method_name not in self._methods:
-            return {
-                "id": request_id,
-                "error": f"Unknown method: {method_name}",
-                "traceback": f"Available methods: {list(self._methods.keys())}",
-            }
-
-        try:
-            # Decode any encoded objects in args
-            decoded_args = decode_object(args)
-
-            # Call method
-            result = self._methods[method_name](**decoded_args)
-
-            # Encode result for JSON
-            encoded_result = encode_object(result)
-
-            return {"id": request_id, "result": encoded_result}
-
-        except Exception as e:
-            return {
-                "id": request_id,
-                "error": str(e),
-                "traceback": traceback.format_exc(),
-            }
-
-    def run(self) -> None:
-        """
-        Main worker loop - reads from stdin, writes to stdout.
-
-        This method blocks until shutdown is requested.
-        """
-        # Suppress library output that could interfere with JSON protocol
-        import warnings
-        import logging
-        import os
-
-        warnings.filterwarnings("ignore")
-        os.environ.setdefault('TF_CPP_MIN_LOG_LEVEL', '3')
-        logging.disable(logging.WARNING)
-
-        self.log("Worker starting...")
-
-        # Run setup
-        try:
-            self.setup()
-            self.log("Setup complete, ready for requests")
-        except Exception as e:
-            self.log(f"Setup failed: {e}")
-            traceback.print_exc(file=sys.stderr)
-            sys.exit(1)
-
-        # Main loop
-        try:
-            for line in sys.stdin:
-                line = line.strip()
-                if not line:
-                    continue
-
-                try:
-                    request = json.loads(line)
-                    response = self.handle_request(request)
-                    print(json.dumps(response), flush=True)
-
-                    if not self._running:
-                        break
-
-                except json.JSONDecodeError as e:
-                    error_response = {
-                        "id": "unknown",
-                        "error": f"Invalid JSON: {e}",
-                    }
-                    print(json.dumps(error_response), flush=True)
-
-        except KeyboardInterrupt:
-            self.log("Interrupted")
-        finally:
-            self.log("Shutting down...")
-            self.teardown()
-            self.log("Goodbye")