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

comfy_env/ipc/protocol.py

@@ -1,265 +0,0 @@
-"""
-IPC Protocol - Message format for bridge-worker communication.
-
-Uses JSON for simplicity and debuggability. Large binary data (images, tensors)
-is serialized efficiently:
-- Tensors: Zero-copy via CUDA IPC or shared memory (see tensor.py)
-- Images: PNG encoded + base64
-- Other: pickle + base64 fallback
-"""
-
-import json
-import base64
-import pickle
-import logging
-from dataclasses import dataclass, field, asdict
-from typing import Any, Dict, Optional
-
-logger = logging.getLogger(__name__)
-
-# Flag to enable/disable IPC tensor sharing (set based on process context)
-_use_tensor_ipc = True
-
-
-def set_tensor_ipc_enabled(enabled: bool) -> None:
-    """Enable or disable IPC tensor sharing."""
-    global _use_tensor_ipc
-    _use_tensor_ipc = enabled
-
-
-def get_tensor_ipc_enabled() -> bool:
-    """Check if IPC tensor sharing is enabled."""
-    return _use_tensor_ipc
-
-
-@dataclass
-class Request:
-    """
-    Request message from bridge to worker.
-
-    Attributes:
-        id: Unique request ID for matching responses
-        method: Method name to call on worker
-        args: Keyword arguments for the method
-    """
-    id: str
-    method: str
-    args: Dict[str, Any] = field(default_factory=dict)
-
-    def to_json(self) -> str:
-        """Serialize to JSON string."""
-        return json.dumps(asdict(self))
-
-    @classmethod
-    def from_json(cls, data: str) -> "Request":
-        """Deserialize from JSON string."""
-        d = json.loads(data)
-        return cls(**d)
-
-
-@dataclass
-class Response:
-    """
-    Response message from worker to bridge.
-
-    Attributes:
-        id: Request ID this is responding to
-        result: Result value (None if error)
-        error: Error message (None if success)
-        traceback: Full traceback string (only if error)
-    """
-    id: str
-    result: Any = None
-    error: Optional[str] = None
-    traceback: Optional[str] = None
-
-    @property
-    def success(self) -> bool:
-        """Check if response indicates success."""
-        return self.error is None
-
-    def to_json(self) -> str:
-        """Serialize to JSON string."""
-        return json.dumps(asdict(self))
-
-    @classmethod
-    def from_json(cls, data: str) -> "Response":
-        """Deserialize from JSON string."""
-        d = json.loads(data)
-        return cls(**d)
-
-
-def encode_binary(data: bytes) -> str:
-    """Encode binary data as base64 string."""
-    return base64.b64encode(data).decode('utf-8')
-
-
-def decode_binary(encoded: str) -> bytes:
-    """Decode base64 string to binary data."""
-    return base64.b64decode(encoded)
-
-
-def encode_object(obj: Any) -> Dict[str, Any]:
-    """
-    Encode a Python object for JSON serialization.
-
-    Returns a dict with _type and _data keys for special types,
-    or the original object if it's JSON-serializable.
-
-    Special handling:
-    - PyTorch tensors: Zero-copy via CUDA IPC or shared memory
-    - PIL Images: PNG encoded
-    - Complex objects: pickle fallback
-    """
-    if obj is None:
-        return None
-
-    # Handle torch tensors - try zero-copy IPC first
-    if hasattr(obj, 'cpu') and hasattr(obj, 'numpy'):
-        try:
-            import torch
-            if isinstance(obj, torch.Tensor) and _use_tensor_ipc:
-                try:
-                    from .tensor import serialize_tensor
-                    return serialize_tensor(obj)
-                except Exception as e:
-                    # Fall back to pickle method if IPC fails
-                    logger.debug(f"Tensor IPC failed, using pickle: {e}")
-        except ImportError:
-            pass
-
-        # Fallback: pickle the numpy array
-        arr = obj.cpu().numpy()
-        return {
-            "_type": "tensor_pickle",
-            "_dtype": str(arr.dtype),
-            "_shape": list(arr.shape),
-            "_data": encode_binary(pickle.dumps(arr)),
-        }
-
-    # Handle numpy arrays
-    if hasattr(obj, '__array__'):
-        import numpy as np
-        arr = np.asarray(obj)
-        return {
-            "_type": "numpy",
-            "_dtype": str(arr.dtype),
-            "_shape": list(arr.shape),
-            "_data": encode_binary(pickle.dumps(arr)),
-        }
-
-    # Handle PIL Images
-    if hasattr(obj, 'save') and hasattr(obj, 'mode'):
-        import io
-        buffer = io.BytesIO()
-        obj.save(buffer, format="PNG")
-        return {
-            "_type": "image",
-            "_format": "PNG",
-            "_data": encode_binary(buffer.getvalue()),
-        }
-
-    # Handle bytes
-    if isinstance(obj, bytes):
-        return {
-            "_type": "bytes",
-            "_data": encode_binary(obj),
-        }
-
-    # Handle lists/tuples recursively
-    if isinstance(obj, (list, tuple)):
-        encoded = [encode_object(item) for item in obj]
-        return {
-            "_type": "list" if isinstance(obj, list) else "tuple",
-            "_data": encoded,
-        }
-
-    # Handle dicts recursively
-    if isinstance(obj, dict):
-        return {k: encode_object(v) for k, v in obj.items()}
-
-    # For simple objects with __dict__, serialize as dict
-    # This avoids pickle module path issues across process boundaries
-    if hasattr(obj, '__dict__') and not hasattr(obj, '__slots__'):
-        return {
-            "_type": "object",
-            "_class": obj.__class__.__name__,
-            "_data": {k: encode_object(v) for k, v in obj.__dict__.items()},
-        }
-
-    # For complex objects that can't be JSON serialized, use pickle
-    try:
-        json.dumps(obj)
-        return obj  # JSON-serializable, return as-is
-    except (TypeError, ValueError):
-        return {
-            "_type": "pickle",
-            "_data": encode_binary(pickle.dumps(obj)),
-        }
-
-
-def decode_object(obj: Any) -> Any:
-    """
-    Decode a JSON-deserialized object back to Python types.
-
-    Reverses the encoding done by encode_object.
-    """
-    if obj is None:
-        return None
-
-    if not isinstance(obj, dict):
-        return obj
-
-    # Check for special encoded types
-    obj_type = obj.get("_type")
-
-    # Handle zero-copy tensor IPC
-    if obj_type == "tensor_ipc":
-        try:
-            from .tensor import deserialize_tensor
-            return deserialize_tensor(obj)
-        except Exception as e:
-            logger.error(f"Failed to deserialize tensor via IPC: {e}")
-            raise
-
-    # Handle pickle fallback for tensors
-    if obj_type == "tensor_pickle":
-        import torch
-        arr = pickle.loads(decode_binary(obj["_data"]))
-        return torch.from_numpy(arr)
-
-    # Legacy types for backwards compatibility
-    if obj_type == "numpy":
-        return pickle.loads(decode_binary(obj["_data"]))
-
-    if obj_type in ("tensor", "comfyui_image", "comfyui_mask"):
-        import torch
-        arr = pickle.loads(decode_binary(obj["_data"]))
-        return torch.from_numpy(arr)
-
-    if obj_type == "image":
-        import io
-        from PIL import Image
-        buffer = io.BytesIO(decode_binary(obj["_data"]))
-        return Image.open(buffer)
-
-    if obj_type == "bytes":
-        return decode_binary(obj["_data"])
-
-    if obj_type == "pickle":
-        return pickle.loads(decode_binary(obj["_data"]))
-
-    # Simple object serialized as dict - restore as SimpleNamespace
-    if obj_type == "object":
-        from types import SimpleNamespace
-        data = {k: decode_object(v) for k, v in obj["_data"].items()}
-        ns = SimpleNamespace(**data)
-        ns._class_name = obj.get("_class", "unknown")
-        return ns
-
-    if obj_type in ("list", "tuple"):
-        decoded = [decode_object(item) for item in obj["_data"]]
-        return decoded if obj_type == "list" else tuple(decoded)
-
-    # Regular dict - decode values recursively
-    return {k: decode_object(v) for k, v in obj.items()}

comfy_env/ipc/tensor.py

@@ -1,371 +0,0 @@
-"""
-Tensor Serialization - Zero-copy tensor sharing via CUDA IPC and shared memory.
-
-This module provides efficient tensor transfer between processes:
-- CUDA tensors: Use CUDA IPC handles (zero-copy, ~0ms for any size)
-- CPU tensors: Use shared memory via file_system strategy (zero-copy)
-
-Based on patterns from pyisolate's tensor_serializer.py.
-"""
-
-import base64
-import collections
-import logging
-import sys
-import threading
-import time
-from typing import Any, Dict, Optional
-
-logger = logging.getLogger(__name__)
-
-
-# ---------------------------------------------------------------------------
-# TensorKeeper - Prevents GC race conditions
-# ---------------------------------------------------------------------------
-
-class TensorKeeper:
-    """
-    Keeps strong references to serialized tensors to prevent premature GC.
-
-    When we serialize a tensor for IPC, we return a handle/path to shared memory.
-    If the original tensor is garbage collected before the receiver opens the
-    shared memory, the data is lost. TensorKeeper holds references for a short
-    window to prevent this race condition.
-
-    Based on pyisolate's TensorKeeper pattern.
-    """
-
-    def __init__(self, retention_seconds: float = 30.0):
-        """
-        Args:
-            retention_seconds: How long to keep references (default 30s)
-        """
-        self.retention_seconds = retention_seconds
-        self._keeper: collections.deque = collections.deque()
-        self._lock = threading.Lock()
-
-    def keep(self, tensor: Any) -> None:
-        """
-        Keep a reference to a tensor.
-
-        Args:
-            tensor: The tensor to keep alive
-        """
-        now = time.time()
-        with self._lock:
-            self._keeper.append((now, tensor))
-
-            # Cleanup old references
-            while self._keeper:
-                timestamp, _ = self._keeper[0]
-                if now - timestamp > self.retention_seconds:
-                    self._keeper.popleft()
-                else:
-                    break
-
-
-# Global tensor keeper instance
-_tensor_keeper = TensorKeeper()
-
-
-# ---------------------------------------------------------------------------
-# Tensor Serialization
-# ---------------------------------------------------------------------------
-
-def serialize_tensor(tensor: Any) -> Dict[str, Any]:
-    """
-    Serialize a PyTorch tensor for IPC transfer.
-
-    Uses zero-copy methods when possible:
-    - CUDA tensors: CUDA IPC handles
-    - CPU tensors: Shared memory (file_system strategy)
-
-    Args:
-        tensor: PyTorch tensor to serialize
-
-    Returns:
-        Dict with tensor metadata and IPC handle/path
-    """
-    import torch
-
-    if not isinstance(tensor, torch.Tensor):
-        raise TypeError(f"Expected torch.Tensor, got {type(tensor)}")
-
-    if tensor.is_cuda:
-        return _serialize_cuda_tensor(tensor)
-    else:
-        return _serialize_cpu_tensor(tensor)
-
-
-def _serialize_cpu_tensor(tensor: Any) -> Dict[str, Any]:
-    """
-    Serialize CPU tensor using shared memory (file_system strategy).
-
-    The tensor is moved to shared memory, and we return the path/key
-    that the receiver can use to access it.
-    """
-    import torch
-    import torch.multiprocessing.reductions as reductions
-
-    # Keep tensor alive until receiver opens shared memory
-    _tensor_keeper.keep(tensor)
-
-    # Move to shared memory if not already
-    if not tensor.is_shared():
-        tensor.share_memory_()
-
-    # Get storage reduction info
-    storage = tensor.untyped_storage()
-    sfunc, sargs = reductions.reduce_storage(storage)
-
-    if sfunc.__name__ == 'rebuild_storage_filename':
-        # file_system strategy - sargs: (cls, manager_path, storage_key, size)
-        return {
-            "_type": "tensor_ipc",
-            "device": "cpu",
-            "strategy": "file_system",
-            "manager_path": sargs[1].decode('utf-8') if isinstance(sargs[1], bytes) else sargs[1],
-            "storage_key": sargs[2].decode('utf-8') if isinstance(sargs[2], bytes) else sargs[2],
-            "storage_size": sargs[3],
-            "dtype": str(tensor.dtype),
-            "shape": list(tensor.shape),
-            "stride": list(tensor.stride()),
-            "offset": tensor.storage_offset(),
-            "requires_grad": tensor.requires_grad,
-        }
-    elif sfunc.__name__ == 'rebuild_storage_fd':
-        # Force file_system strategy for compatibility
-        import torch.multiprocessing as mp
-        mp.set_sharing_strategy('file_system')
-        tensor.share_memory_()
-        return _serialize_cpu_tensor(tensor)
-    else:
-        # Fallback: pickle the tensor data (slow path)
-        logger.warning(f"Unknown storage reduction: {sfunc.__name__}, using pickle fallback")
-        return _serialize_tensor_fallback(tensor)
-
-
-def _serialize_cuda_tensor(tensor: Any) -> Dict[str, Any]:
-    """
-    Serialize CUDA tensor using CUDA IPC handles.
-
-    This is zero-copy - we only transfer the IPC handle, not the data.
-    The receiver uses the handle to map the same GPU memory.
-    """
-    import torch
-    import torch.multiprocessing.reductions as reductions
-
-    try:
-        func, args = reductions.reduce_tensor(tensor)
-    except RuntimeError as e:
-        if "received from another process" in str(e):
-            # Tensor was received via IPC and can't be re-shared
-            # Need to clone it (expensive but necessary)
-            tensor_size_mb = tensor.numel() * tensor.element_size() / (1024 * 1024)
-            if tensor_size_mb > 100:
-                logger.warning(
-                    f"Cloning large CUDA tensor ({tensor_size_mb:.1f}MB) - "
-                    "consider avoiding returning unmodified input tensors"
-                )
-            tensor = tensor.clone()
-            func, args = reductions.reduce_tensor(tensor)
-        else:
-            raise
-
-    # Keep tensor alive until receiver maps it
-    _tensor_keeper.keep(tensor)
-
-    # args structure for CUDA tensor:
-    # (cls, size, stride, offset, storage_type, dtype, device_idx, handle,
-    #  storage_size, storage_offset, requires_grad, ref_counter_handle,
-    #  ref_counter_offset, event_handle, event_sync_required)
-    return {
-        "_type": "tensor_ipc",
-        "device": "cuda",
-        "device_idx": args[6],
-        "shape": list(args[1]),
-        "stride": list(args[2]),
-        "offset": args[3],
-        "dtype": str(args[5]),
-        "handle": base64.b64encode(args[7]).decode('ascii'),
-        "storage_size": args[8],
-        "storage_offset": args[9],
-        "requires_grad": args[10],
-        "ref_counter_handle": base64.b64encode(args[11]).decode('ascii'),
-        "ref_counter_offset": args[12],
-        "event_handle": base64.b64encode(args[13]).decode('ascii') if args[13] else None,
-        "event_sync_required": args[14],
-    }
-
-
-def _serialize_tensor_fallback(tensor: Any) -> Dict[str, Any]:
-    """
-    Fallback serialization using pickle (slow, copies data).
-
-    Used when zero-copy methods aren't available.
-    """
-    import pickle
-
-    arr = tensor.cpu().numpy()
-    return {
-        "_type": "tensor_pickle",
-        "dtype": str(tensor.dtype),
-        "shape": list(tensor.shape),
-        "device": str(tensor.device),
-        "data": base64.b64encode(pickle.dumps(arr)).decode('ascii'),
-    }
-
-
-# ---------------------------------------------------------------------------
-# Tensor Deserialization
-# ---------------------------------------------------------------------------
-
-def deserialize_tensor(data: Dict[str, Any]) -> Any:
-    """
-    Deserialize a tensor from IPC format.
-
-    Args:
-        data: Dict with tensor metadata from serialize_tensor
-
-    Returns:
-        PyTorch tensor
-    """
-    import torch
-
-    # Already a tensor (shouldn't happen, but handle gracefully)
-    if isinstance(data, torch.Tensor):
-        return data
-
-    obj_type = data.get("_type")
-
-    if obj_type == "tensor_ipc":
-        device = data.get("device", "cpu")
-        if device == "cuda":
-            return _deserialize_cuda_tensor(data)
-        else:
-            return _deserialize_cpu_tensor(data)
-    elif obj_type == "tensor_pickle":
-        return _deserialize_tensor_fallback(data)
-    else:
-        raise ValueError(f"Unknown tensor type: {obj_type}")
-
-
-def _deserialize_cpu_tensor(data: Dict[str, Any]) -> Any:
-    """Deserialize CPU tensor from shared memory."""
-    import torch
-    import torch.multiprocessing.reductions as reductions
-
-    strategy = data.get("strategy")
-    if strategy != "file_system":
-        raise RuntimeError(f"Unsupported CPU tensor strategy: {strategy}")
-
-    dtype_str = data["dtype"]
-    dtype = getattr(torch, dtype_str.split(".")[-1])
-
-    manager_path = data["manager_path"]
-    storage_key = data["storage_key"]
-    storage_size = data["storage_size"]
-
-    # Convert to bytes if needed
-    if isinstance(manager_path, str):
-        manager_path = manager_path.encode('utf-8')
-    if isinstance(storage_key, str):
-        storage_key = storage_key.encode('utf-8')
-
-    # Rebuild storage
-    rebuilt_storage = reductions.rebuild_storage_filename(
-        torch.UntypedStorage, manager_path, storage_key, storage_size
-    )
-
-    # Wrap in typed storage
-    typed_storage = torch.storage.TypedStorage(
-        wrap_storage=rebuilt_storage, dtype=dtype, _internal=True
-    )
-
-    # Rebuild tensor
-    metadata = (
-        data["offset"],
-        tuple(data["shape"]),
-        tuple(data["stride"]),
-        data["requires_grad"],
-    )
-    tensor = reductions.rebuild_tensor(torch.Tensor, typed_storage, metadata)
-    return tensor
-
-
-def _deserialize_cuda_tensor(data: Dict[str, Any]) -> Any:
-    """Deserialize CUDA tensor from IPC handle."""
-    import torch
-    import torch.multiprocessing.reductions as reductions
-
-    dtype_str = data["dtype"]
-    dtype = getattr(torch, dtype_str.split(".")[-1])
-
-    handle = base64.b64decode(data["handle"])
-    ref_counter_handle = base64.b64decode(data["ref_counter_handle"])
-    event_handle = base64.b64decode(data["event_handle"]) if data.get("event_handle") else None
-    device_idx = data.get("device_idx", 0)
-
-    tensor = reductions.rebuild_cuda_tensor(
-        torch.Tensor,
-        tuple(data["shape"]),
-        tuple(data["stride"]),
-        data["offset"],
-        torch.storage.TypedStorage,
-        dtype,
-        device_idx,
-        handle,
-        data["storage_size"],
-        data["storage_offset"],
-        data["requires_grad"],
-        ref_counter_handle,
-        data["ref_counter_offset"],
-        event_handle,
-        data["event_sync_required"],
-    )
-    return tensor
-
-
-def _deserialize_tensor_fallback(data: Dict[str, Any]) -> Any:
-    """Deserialize tensor from pickle fallback format."""
-    import pickle
-    import torch
-
-    arr = pickle.loads(base64.b64decode(data["data"]))
-    tensor = torch.from_numpy(arr)
-
-    # Move to original device if CUDA
-    device = data.get("device", "cpu")
-    if "cuda" in device:
-        tensor = tensor.to(device)
-
-    return tensor
-
-
-# ---------------------------------------------------------------------------
-# Integration with protocol.py
-# ---------------------------------------------------------------------------
-
-def is_tensor(obj: Any) -> bool:
-    """Check if object is a PyTorch tensor."""
-    try:
-        import torch
-        return isinstance(obj, torch.Tensor)
-    except ImportError:
-        return False
-
-
-def can_use_ipc() -> bool:
-    """
-    Check if IPC tensor sharing is available.
-
-    Requires:
-    - PyTorch installed
-    - multiprocessing spawn context works
-    """
-    try:
-        import torch
-        import torch.multiprocessing as mp
-        return True
-    except ImportError:
-        return False