comfy-env 0.0.8__py3-none-any.whl

comfy_env/decorator.py

@@ -0,0 +1,422 @@
+"""
+Decorator-based API for easy subprocess isolation.
+
+This module provides the @isolated decorator that makes it simple to run
+ComfyUI node methods in isolated subprocess environments.
+
+Architecture:
+    The decorator wraps the node's FUNCTION method. When called in the HOST
+    process, it forwards the call to an isolated worker (TorchMPWorker for
+    same-venv, PersistentVenvWorker for different venv).
+
+    When imported in the WORKER subprocess (COMFYUI_ISOLATION_WORKER=1),
+    the decorator is a transparent no-op.
+
+Example:
+    from comfy_env import isolated
+
+    @isolated(env="myenv")
+    class MyNode:
+        FUNCTION = "process"
+        RETURN_TYPES = ("IMAGE",)
+
+        def process(self, image):
+            # This code runs in isolated subprocess
+            import heavy_package
+            return (heavy_package.run(image),)
+
+Implementation:
+    This decorator is thin sugar over the workers module. Internally it uses:
+    - TorchMPWorker: Same Python, zero-copy tensor transfer via torch.mp.Queue
+    - PersistentVenvWorker: Different venv, tensor transfer via torch.save/load
+"""
+
+import os
+import sys
+import atexit
+import inspect
+import logging
+import threading
+import time
+from dataclasses import dataclass
+from functools import wraps
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Union
+
+logger = logging.getLogger("comfy_env")
+
+# Enable verbose logging by default (can be disabled)
+VERBOSE_LOGGING = os.environ.get("COMFYUI_ISOLATION_QUIET", "0") != "1"
+
+
+def _log(env_name: str, msg: str):
+    """Log with environment prefix."""
+    if VERBOSE_LOGGING:
+        print(f"[{env_name}] {msg}")
+
+
+def _is_worker_mode() -> bool:
+    """Check if we're running inside the worker subprocess."""
+    return os.environ.get("COMFYUI_ISOLATION_WORKER") == "1"
+
+
+def _describe_tensor(t) -> str:
+    """Get human-readable tensor description."""
+    try:
+        import torch
+        if isinstance(t, torch.Tensor):
+            size_mb = t.numel() * t.element_size() / (1024 * 1024)
+            return f"Tensor({list(t.shape)}, {t.dtype}, {t.device}, {size_mb:.1f}MB)"
+    except:
+        pass
+    return str(type(t).__name__)
+
+
+def _describe_args(args: dict) -> str:
+    """Describe arguments for logging."""
+    parts = []
+    for k, v in args.items():
+        parts.append(f"{k}={_describe_tensor(v)}")
+    return ", ".join(parts) if parts else "(no args)"
+
+
+def _clone_tensor_if_needed(obj: Any, smart_clone: bool = True) -> Any:
+    """
+    Defensively clone tensors to prevent mutation/re-share bugs.
+
+    This handles:
+    1. Input tensors that might be mutated in worker
+    2. Output tensors received via IPC that can't be re-shared
+
+    Args:
+        obj: Object to process (tensor or nested structure)
+        smart_clone: If True, use smart CUDA IPC detection (only clone
+                    when necessary). If False, always clone.
+    """
+    if smart_clone:
+        # Use smart detection - only clones CUDA tensors that can't be re-shared
+        from .workers.tensor_utils import prepare_for_ipc_recursive
+        return prepare_for_ipc_recursive(obj)
+
+    # Fallback: always clone (original behavior)
+    try:
+        import torch
+        if isinstance(obj, torch.Tensor):
+            return obj.clone()
+        elif isinstance(obj, (list, tuple)):
+            cloned = [_clone_tensor_if_needed(x, smart_clone=False) for x in obj]
+            return type(obj)(cloned)
+        elif isinstance(obj, dict):
+            return {k: _clone_tensor_if_needed(v, smart_clone=False) for k, v in obj.items()}
+    except ImportError:
+        pass
+    return obj
+
+
+# ---------------------------------------------------------------------------
+# Worker Management
+# ---------------------------------------------------------------------------
+
+@dataclass
+class WorkerConfig:
+    """Configuration for an isolated worker."""
+    env_name: str
+    python: Optional[str] = None  # None = same Python (TorchMPWorker)
+    working_dir: Optional[Path] = None
+    sys_path: Optional[List[str]] = None
+    timeout: float = 600.0
+
+
+# Global worker cache
+_workers: Dict[str, Any] = {}
+_workers_lock = threading.Lock()
+
+
+def _get_or_create_worker(config: WorkerConfig, log_fn: Callable):
+    """Get or create a worker for the given configuration.
+
+    Thread-safe: worker creation happens inside the lock to prevent
+    race conditions where multiple threads create duplicate workers.
+    """
+    cache_key = f"{config.env_name}:{config.python or 'same'}"
+
+    with _workers_lock:
+        if cache_key in _workers:
+            worker = _workers[cache_key]
+            if worker.is_alive():
+                return worker
+            # Worker died, recreate
+            log_fn(f"Worker died, recreating...")
+
+        # Create new worker INSIDE the lock (fixes race condition)
+        if config.python is None:
+            # Same Python - use TorchMPWorker (fast, zero-copy)
+            from .workers import TorchMPWorker
+            log_fn(f"Creating TorchMPWorker (same Python, zero-copy tensors)")
+            worker = TorchMPWorker(name=config.env_name)
+        else:
+            # Different Python - use PersistentVenvWorker
+            from .workers.venv import PersistentVenvWorker
+            log_fn(f"Creating PersistentVenvWorker (python={config.python})")
+            worker = PersistentVenvWorker(
+                python=config.python,
+                working_dir=config.working_dir,
+                sys_path=config.sys_path,
+                name=config.env_name,
+            )
+
+        _workers[cache_key] = worker
+        return worker
+
+
+def shutdown_all_processes():
+    """Shutdown all cached workers. Called at exit."""
+    with _workers_lock:
+        for name, worker in _workers.items():
+            try:
+                worker.shutdown()
+            except Exception as e:
+                logger.debug(f"Error shutting down {name}: {e}")
+        _workers.clear()
+
+
+atexit.register(shutdown_all_processes)
+
+
+# ---------------------------------------------------------------------------
+# The @isolated Decorator
+# ---------------------------------------------------------------------------
+
+def isolated(
+    env: str,
+    requirements: Optional[List[str]] = None,
+    config: Optional[str] = None,
+    python: Optional[str] = None,
+    cuda: Optional[str] = "auto",
+    timeout: float = 600.0,
+    log_callback: Optional[Callable[[str], None]] = None,
+    import_paths: Optional[List[str]] = None,
+    clone_tensors: bool = True,
+    same_venv: bool = False,
+):
+    """
+    Class decorator that runs node methods in isolated subprocess.
+
+    The decorated class's FUNCTION method will be executed in an isolated
+    Python environment. Tensors are transferred efficiently via PyTorch's
+    native IPC mechanisms (CUDA IPC for GPU, shared memory for CPU).
+
+    By default, auto-discovers config file (comfy_env_reqs.toml) and
+    uses full venv isolation with PersistentVenvWorker. Use same_venv=True
+    for lightweight same-venv isolation with TorchMPWorker.
+
+    Args:
+        env: Name of the isolated environment (used for logging/caching)
+        requirements: [DEPRECATED] Use config file instead
+        config: Path to TOML config file. If None, auto-discovers in node directory.
+        python: Path to Python executable (overrides config-based detection)
+        cuda: [DEPRECATED] Detected automatically
+        timeout: Timeout for calls in seconds (default: 10 minutes)
+        log_callback: Optional callback for logging
+        import_paths: Paths to add to sys.path in worker
+        clone_tensors: Clone tensors at boundary to prevent mutation bugs (default: True)
+        same_venv: If True, use TorchMPWorker (same venv, just process isolation).
+                   If False (default), use full venv isolation with auto-discovered config.
+
+    Example:
+        # Full venv isolation (default) - auto-discovers comfy_env_reqs.toml
+        @isolated(env="sam3d")
+        class MyNode:
+            FUNCTION = "process"
+
+            def process(self, image):
+                import heavy_lib
+                return heavy_lib.run(image)
+
+        # Lightweight same-venv isolation (opt-in)
+        @isolated(env="sam3d", same_venv=True)
+        class MyLightNode:
+            FUNCTION = "process"
+            ...
+    """
+    def decorator(cls):
+        # In worker mode, decorator is a no-op
+        if _is_worker_mode():
+            return cls
+
+        # --- HOST MODE: Wrap the FUNCTION method ---
+
+        func_name = getattr(cls, 'FUNCTION', None)
+        if not func_name:
+            raise ValueError(
+                f"Node class {cls.__name__} must have FUNCTION attribute."
+            )
+
+        original_method = getattr(cls, func_name, None)
+        if original_method is None:
+            raise ValueError(
+                f"Node class {cls.__name__} has FUNCTION='{func_name}' but "
+                f"no method with that name."
+            )
+
+        # Get source file info for sys.path setup
+        source_file = Path(inspect.getfile(cls))
+        node_dir = source_file.parent
+        if node_dir.name == "nodes":
+            node_package_dir = node_dir.parent
+        else:
+            node_package_dir = node_dir
+
+        # Build sys.path for worker
+        sys_path_additions = [str(node_dir)]
+        if import_paths:
+            for p in import_paths:
+                full_path = node_dir / p
+                sys_path_additions.append(str(full_path.resolve()))
+
+        # Resolve python path for venv isolation
+        resolved_python = python
+        env_config = None
+
+        # If same_venv=True, skip venv isolation entirely
+        if same_venv:
+            _log(env, "Using same-venv isolation (TorchMPWorker)")
+            resolved_python = None
+
+        # Otherwise, try to get a venv python path
+        elif python:
+            # Explicit python path provided
+            resolved_python = python
+
+        else:
+            # Auto-discover or use explicit config
+            if config:
+                # Explicit config file specified
+                config_file = node_package_dir / config
+                if config_file.exists():
+                    from .env.config_file import load_env_from_file
+                    env_config = load_env_from_file(config_file, node_package_dir)
+                else:
+                    _log(env, f"Warning: Config file not found: {config_file}")
+            else:
+                # Auto-discover config file - try v2 API first
+                from .env.config_file import discover_config, discover_env_config
+                v2_config = discover_config(node_package_dir)
+                if v2_config and env in v2_config.envs:
+                    # v2 schema: get the named environment
+                    env_config = v2_config.envs[env]
+                    _log(env, f"Auto-discovered v2 config: {env_config.name}")
+                else:
+                    # Fall back to v1 API
+                    env_config = discover_env_config(node_package_dir)
+                    if env_config:
+                        _log(env, f"Auto-discovered config: {env_config.name}")
+
+            # If we have a config, set up the venv
+            if env_config:
+                from .env.manager import IsolatedEnvManager
+                manager = IsolatedEnvManager(base_dir=node_package_dir)
+
+                if not manager.is_ready(env_config):
+                    _log(env, f"Setting up isolated environment...")
+                    manager.setup(env_config)
+
+                resolved_python = str(manager.get_python(env_config))
+            else:
+                # No config found - fall back to same-venv isolation
+                _log(env, "No config found, using same-venv isolation (TorchMPWorker)")
+                resolved_python = None
+
+        # Create worker config
+        worker_config = WorkerConfig(
+            env_name=env,
+            python=resolved_python,
+            working_dir=node_dir,
+            sys_path=sys_path_additions,
+            timeout=timeout,
+        )
+
+        # Setup logging
+        log_fn = log_callback or (lambda msg: _log(env, msg))
+
+        # Create the proxy method
+        @wraps(original_method)
+        def proxy(self, *args, **kwargs):
+            # Get or create worker
+            worker = _get_or_create_worker(worker_config, log_fn)
+
+            # Bind arguments to get kwargs dict
+            sig = inspect.signature(original_method)
+            try:
+                bound = sig.bind(self, *args, **kwargs)
+                bound.apply_defaults()
+                call_kwargs = dict(bound.arguments)
+                del call_kwargs['self']
+            except TypeError:
+                call_kwargs = kwargs
+
+            # Log entry with argument descriptions
+            if VERBOSE_LOGGING:
+                log_fn(f"→ {cls.__name__}.{func_name}({_describe_args(call_kwargs)})")
+
+            start_time = time.time()
+
+            try:
+                # Clone tensors defensively if enabled
+                if clone_tensors:
+                    call_kwargs = {k: _clone_tensor_if_needed(v) for k, v in call_kwargs.items()}
+
+                # Get module name for import in worker
+                module_name = cls.__module__
+
+                # Call worker using appropriate method
+                if worker_config.python is None:
+                    # TorchMPWorker - use call_method protocol (avoids pickle issues)
+                    result = worker.call_method(
+                        module_name=module_name,
+                        class_name=cls.__name__,
+                        method_name=func_name,
+                        self_state=self.__dict__.copy(),
+                        kwargs=call_kwargs,
+                        timeout=timeout,
+                    )
+                else:
+                    # PersistentVenvWorker - call by module/class/method path
+                    result = worker.call_method(
+                        module_name=source_file.stem,
+                        class_name=cls.__name__,
+                        method_name=func_name,
+                        self_state=self.__dict__.copy() if hasattr(self, '__dict__') else None,
+                        kwargs=call_kwargs,
+                        timeout=timeout,
+                    )
+
+                # Clone result tensors defensively
+                if clone_tensors:
+                    result = _clone_tensor_if_needed(result)
+
+                elapsed = time.time() - start_time
+                if VERBOSE_LOGGING:
+                    result_desc = _describe_tensor(result) if not isinstance(result, tuple) else f"tuple({len(result)} items)"
+                    log_fn(f"← {cls.__name__}.{func_name} returned {result_desc} [{elapsed:.2f}s]")
+
+                return result
+
+            except Exception as e:
+                elapsed = time.time() - start_time
+                log_fn(f"✗ {cls.__name__}.{func_name} failed after {elapsed:.2f}s: {e}")
+                raise
+
+        # Store original method before replacing (for worker to access)
+        cls._isolated_original_method = original_method
+
+        # Replace method with proxy
+        setattr(cls, func_name, proxy)
+
+        # Store metadata
+        cls._isolated_env = env
+        cls._isolated_node_dir = node_dir
+
+        return cls
+
+    return decorator

comfy_env/env/__init__.py

@@ -0,0 +1,30 @@
+"""Environment management for comfyui-isolation."""
+
+from .config import IsolatedEnv
+from .manager import IsolatedEnvManager
+from .detection import detect_cuda_version, detect_gpu_info, get_gpu_summary
+from .platform import get_platform, PlatformProvider, PlatformPaths
+from .security import (
+    normalize_env_name,
+    validate_dependency,
+    validate_dependencies,
+    validate_path_within_root,
+    validate_wheel_url,
+)
+
+__all__ = [
+    "IsolatedEnv",
+    "IsolatedEnvManager",
+    "detect_cuda_version",
+    "detect_gpu_info",
+    "get_gpu_summary",
+    "get_platform",
+    "PlatformProvider",
+    "PlatformPaths",
+    # Security
+    "normalize_env_name",
+    "validate_dependency",
+    "validate_dependencies",
+    "validate_path_within_root",
+    "validate_wheel_url",
+]

comfy_env/env/config.py

@@ -0,0 +1,144 @@
+"""Configuration for isolated environments."""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional
+
+
+@dataclass
+class LocalConfig:
+    """Configuration for local (host environment) installs.
+
+    These packages are installed into ComfyUI's main environment,
+    not into an isolated venv.
+    """
+    cuda_packages: Dict[str, str] = field(default_factory=dict)  # package -> version
+    requirements: List[str] = field(default_factory=list)
+
+
+@dataclass
+class NodeReq:
+    """A node dependency (another ComfyUI node pack)."""
+    name: str
+    repo: str  # GitHub repo path, e.g., "Kosinkadink/ComfyUI-VideoHelperSuite"
+
+
+@dataclass
+class EnvManagerConfig:
+    """
+    Full configuration parsed from comfyui_env.toml.
+
+    Supports the v2 schema:
+        [local.cuda]        - CUDA packages for host environment
+        [local.packages]    - Regular packages for host environment
+        [envname]           - Isolated env definition
+        [envname.cuda]      - CUDA packages for isolated env
+        [envname.packages]  - Regular packages for isolated env
+        [node_reqs]         - Node dependencies
+    """
+    local: LocalConfig = field(default_factory=LocalConfig)
+    envs: Dict[str, "IsolatedEnv"] = field(default_factory=dict)
+    node_reqs: List[NodeReq] = field(default_factory=list)
+
+    @property
+    def has_local(self) -> bool:
+        """Check if there are local packages to install."""
+        return bool(self.local.cuda_packages or self.local.requirements)
+
+    @property
+    def has_envs(self) -> bool:
+        """Check if there are isolated environments defined."""
+        return bool(self.envs)
+
+    def get_env(self, name: str) -> Optional["IsolatedEnv"]:
+        """Get an isolated environment by name."""
+        return self.envs.get(name)
+
+    @property
+    def default_env(self) -> Optional["IsolatedEnv"]:
+        """Get the first/only isolated environment, or None."""
+        if self.envs:
+            return next(iter(self.envs.values()))
+        return None
+
+
+@dataclass
+class IsolatedEnv:
+    """
+    Configuration for an isolated Python environment.
+
+    This defines what Python version, CUDA version, and dependencies
+    should be installed in the isolated environment.
+
+    Args:
+        name: Unique name for this environment (used for caching)
+        python: Python version (e.g., "3.10", "3.11")
+        cuda: CUDA version (e.g., "12.4", "12.8") or None for CPU-only
+        requirements: List of pip requirements (e.g., ["torch==2.8.0", "numpy"])
+        requirements_file: Path to requirements.txt file
+        wheel_sources: List of URLs for --find-links (custom wheel repos)
+        index_urls: List of URLs for --extra-index-url
+        env_dir: Custom directory for the venv (default: auto-generated)
+        pytorch_version: Specific PyTorch version (auto-detected if None)
+        worker_package: Worker package directory (e.g., "worker" -> worker/__main__.py)
+        worker_script: Worker script file (e.g., "worker.py")
+
+    Example:
+        env = IsolatedEnv(
+            name="my-node",
+            python="3.10",
+            cuda="12.8",
+            requirements=["torch==2.8.0", "nvdiffrast"],
+            wheel_sources=["https://my-wheels.github.io/"],
+        )
+    """
+
+    name: str
+    python: str = "3.10"
+    cuda: Optional[str] = None
+    requirements: list[str] = field(default_factory=list)
+    no_deps_requirements: list[str] = field(default_factory=list)  # Install with --no-deps
+    requirements_file: Optional[Path] = None
+    wheel_sources: list[str] = field(default_factory=list)
+    index_urls: list[str] = field(default_factory=list)
+    env_dir: Optional[Path] = None
+    pytorch_version: Optional[str] = None
+    # Worker configuration
+    worker_package: Optional[str] = None  # e.g., "worker" -> worker/__main__.py
+    worker_script: Optional[str] = None   # e.g., "worker.py" -> worker.py
+
+    def __post_init__(self):
+        """Validate and normalize configuration."""
+        # Normalize paths
+        if self.requirements_file is not None:
+            self.requirements_file = Path(self.requirements_file)
+        if self.env_dir is not None:
+            self.env_dir = Path(self.env_dir)
+
+        # Validate Python version
+        if not self.python.replace(".", "").isdigit():
+            raise ValueError(f"Invalid Python version: {self.python}")
+
+        # Validate CUDA version if specified
+        if self.cuda is not None:
+            cuda_clean = self.cuda.replace(".", "")
+            if not cuda_clean.isdigit():
+                raise ValueError(f"Invalid CUDA version: {self.cuda}")
+
+    @property
+    def cuda_short(self) -> Optional[str]:
+        """Get CUDA version without dots (e.g., '128' for '12.8')."""
+        if self.cuda is None:
+            return None
+        return self.cuda.replace(".", "")
+
+    @property
+    def python_short(self) -> str:
+        """Get Python version without dots (e.g., '310' for '3.10')."""
+        return self.python.replace(".", "")
+
+    def get_default_env_dir(self, base_dir: Path) -> Path:
+        """Get the default environment directory path."""
+        if self.env_dir is not None:
+            return self.env_dir
+        return base_dir / f"_env_{self.name}"