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

comfy_env/decorator.py

@@ -1,700 +0,0 @@
-"""
-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, Tuple, 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
-
-
-def _find_node_package_dir(source_file: Path) -> Path:
-    """
-    Find the node package root directory by searching for comfy-env.toml.
-
-    Walks up from the source file's directory until it finds a config file,
-    or falls back to heuristics if not found.
-    """
-    from .env.config_file import CONFIG_FILE_NAMES
-
-    current = source_file.parent
-
-    # Walk up the directory tree looking for config file
-    while current != current.parent:  # Stop at filesystem root
-        for config_name in CONFIG_FILE_NAMES:
-            if (current / config_name).exists():
-                return current
-        current = current.parent
-
-    # Fallback: use old heuristic if no config found
-    node_dir = source_file.parent
-    if node_dir.name == "nodes":
-        return node_dir.parent
-    return node_dir
-
-
-# ---------------------------------------------------------------------------
-# 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,
-                sys_path=config.sys_path,
-            )
-        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
-        node_package_dir = _find_node_package_dir(source_file)
-
-        # 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
-                # Note: ComfyUI uses full filesystem paths as module names for custom nodes.
-                # The worker's _execute_method_call handles this by using file-based imports.
-                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"[FAIL] {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
-
-
-# ---------------------------------------------------------------------------
-# The @auto_isolate Decorator (Function-level)
-# ---------------------------------------------------------------------------
-
-def _parse_import_error(e: ImportError) -> Optional[str]:
-    """Extract the module name from an ImportError."""
-    # Python's ImportError has a 'name' attribute with the module name
-    if hasattr(e, 'name') and e.name:
-        return e.name
-
-    # Fallback: parse from message "No module named 'xxx'"
-    msg = str(e)
-    if "No module named" in msg:
-        # Extract 'xxx' from "No module named 'xxx'" or "No module named 'xxx.yyy'"
-        import re
-        match = re.search(r"No module named ['\"]([^'\"\.]+)", msg)
-        if match:
-            return match.group(1)
-
-    return None
-
-
-def _find_env_for_module(
-    module_name: str,
-    source_file: Path,
-) -> Optional[Tuple[str, Path, Path]]:
-    """
-    Find which isolated environment contains the given module.
-
-    Searches comfy-env.toml configs starting from the source file's directory,
-    looking for the module in cuda packages, requirements, etc.
-
-    Args:
-        module_name: The module that failed to import (e.g., "cumesh")
-        source_file: Path to the source file containing the function
-
-    Returns:
-        Tuple of (env_name, python_path, node_dir) or None if not found
-    """
-    from .env.config_file import discover_config, CONFIG_FILE_NAMES
-
-    # Normalize module name (cumesh, pytorch3d, etc.)
-    module_lower = module_name.lower().replace("-", "_").replace(".", "_")
-
-    # Search for config file starting from source file's directory
-    node_dir = source_file.parent
-    while node_dir != node_dir.parent:
-        for config_name in CONFIG_FILE_NAMES:
-            config_path = node_dir / config_name
-            if config_path.exists():
-                # Found a config, check if it has our module
-                config = discover_config(node_dir)
-                if config is None:
-                    continue
-
-                # Check all environments in the config
-                for env_name, env_config in config.envs.items():
-                    # Check cuda/no_deps_requirements
-                    if env_config.no_deps_requirements:
-                        for req in env_config.no_deps_requirements:
-                            req_name = req.split("==")[0].split(">=")[0].split("<")[0].strip()
-                            req_lower = req_name.lower().replace("-", "_")
-                            if req_lower == module_lower:
-                                # Found it! Get the python path
-                                env_path = node_dir / f"_env_{env_name}"
-                                if not env_path.exists():
-                                    # Try pixi path
-                                    env_path = node_dir / ".pixi" / "envs" / "default"
-
-                                if env_path.exists():
-                                    python_path = env_path / "bin" / "python"
-                                    if not python_path.exists():
-                                        python_path = env_path / "Scripts" / "python.exe"
-                                    if python_path.exists():
-                                        return (env_name, python_path, node_dir)
-
-                    # Check regular requirements too
-                    if env_config.requirements:
-                        for req in env_config.requirements:
-                            req_name = req.split("==")[0].split(">=")[0].split("<")[0].split("[")[0].strip()
-                            req_lower = req_name.lower().replace("-", "_")
-                            if req_lower == module_lower:
-                                env_path = node_dir / f"_env_{env_name}"
-                                if not env_path.exists():
-                                    env_path = node_dir / ".pixi" / "envs" / "default"
-
-                                if env_path.exists():
-                                    python_path = env_path / "bin" / "python"
-                                    if not python_path.exists():
-                                        python_path = env_path / "Scripts" / "python.exe"
-                                    if python_path.exists():
-                                        return (env_name, python_path, node_dir)
-
-                # Config found but module not in it, stop searching
-                break
-
-        node_dir = node_dir.parent
-
-    return None
-
-
-# Cache for auto_isolate workers
-_auto_isolate_workers: Dict[str, Any] = {}
-_auto_isolate_lock = threading.Lock()
-
-
-def _get_auto_isolate_worker(env_name: str, python_path: Path, node_dir: Path):
-    """Get or create a worker for auto_isolate."""
-    cache_key = f"{env_name}:{python_path}"
-
-    with _auto_isolate_lock:
-        if cache_key in _auto_isolate_workers:
-            worker = _auto_isolate_workers[cache_key]
-            if worker.is_alive():
-                return worker
-
-        # Create new PersistentVenvWorker
-        from .workers.venv import PersistentVenvWorker
-
-        worker = PersistentVenvWorker(
-            python=str(python_path),
-            working_dir=node_dir,
-            sys_path=[str(node_dir)],
-            name=f"auto-{env_name}",
-        )
-
-        _auto_isolate_workers[cache_key] = worker
-        return worker
-
-
-def auto_isolate(func: Callable) -> Callable:
-    """
-    Decorator that automatically runs a function in an isolated environment
-    when an ImportError occurs for a package that exists in the isolated env.
-
-    This provides seamless isolation - just write normal code with imports,
-    and if the import fails in the host environment but the package is
-    configured in comfy-env.toml, the function automatically retries in
-    the isolated environment.
-
-    Example:
-        from comfy_env import auto_isolate
-
-        @auto_isolate
-        def process_with_cumesh(mesh, target_faces):
-            import cumesh  # If this fails, function retries in isolated env
-            import torch
-
-            v = torch.tensor(mesh.vertices).cuda()
-            f = torch.tensor(mesh.faces).cuda()
-
-            cm = cumesh.CuMesh()
-            cm.init(v, f)
-            cm.simplify(target_faces)
-
-            result_v, result_f = cm.read()
-            return result_v.cpu().numpy(), result_f.cpu().numpy()
-
-    How it works:
-        1. Function runs normally in the host environment
-        2. If ImportError occurs, decorator catches it
-        3. Extracts the module name from the error (e.g., "cumesh")
-        4. Searches comfy-env.toml for which env has that module
-        5. Re-runs the entire function in that isolated environment
-        6. Returns the result as if nothing happened
-
-    Benefits:
-        - Zero overhead when imports succeed (fast path)
-        - Auto-detects which environment to use from the failed import
-        - Function is the isolation boundary (clean, debuggable)
-        - Works with any import pattern (top of function, conditional, etc.)
-
-    Note:
-        Arguments and return values are serialized via torch.save/load,
-        so they should be tensors, numpy arrays, or pickle-able objects.
-    """
-    # Get source file for environment detection
-    source_file = Path(inspect.getfile(func))
-
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        try:
-            # Fast path: try running in host environment
-            return func(*args, **kwargs)
-
-        except ImportError as e:
-            # Extract module name from error
-            module_name = _parse_import_error(e)
-            if module_name is None:
-                # Can't determine module, re-raise
-                raise
-
-            # Find which env has this module
-            env_info = _find_env_for_module(module_name, source_file)
-            if env_info is None:
-                # Module not in any known isolated env, re-raise
-                raise
-
-            env_name, python_path, node_dir = env_info
-
-            _log(env_name, f"Import '{module_name}' failed in host, retrying in isolated env...")
-            _log(env_name, f"  Python: {python_path}")
-
-            # Get or create worker
-            worker = _get_auto_isolate_worker(env_name, python_path, node_dir)
-
-            # Prepare arguments - convert numpy arrays to lists for IPC
-            import numpy as np
-
-            def convert_for_ipc(obj):
-                if isinstance(obj, np.ndarray):
-                    return obj.tolist()
-                elif hasattr(obj, 'vertices') and hasattr(obj, 'faces'):
-                    # Trimesh-like object - convert to dict
-                    return {
-                        '__trimesh__': True,
-                        'vertices': obj.vertices.tolist() if hasattr(obj.vertices, 'tolist') else list(obj.vertices),
-                        'faces': obj.faces.tolist() if hasattr(obj.faces, 'tolist') else list(obj.faces),
-                    }
-                elif isinstance(obj, (list, tuple)):
-                    converted = [convert_for_ipc(x) for x in obj]
-                    return type(obj)(converted) if isinstance(obj, tuple) else converted
-                elif isinstance(obj, dict):
-                    return {k: convert_for_ipc(v) for k, v in obj.items()}
-                return obj
-
-            converted_args = [convert_for_ipc(arg) for arg in args]
-            converted_kwargs = {k: convert_for_ipc(v) for k, v in kwargs.items()}
-
-            # Call via worker
-            start_time = time.time()
-
-            result = worker.call_module(
-                module=source_file.stem,
-                func=func.__name__,
-                *converted_args,
-                **converted_kwargs,
-            )
-
-            elapsed = time.time() - start_time
-            _log(env_name, f"<- {func.__name__} completed in isolated env [{elapsed:.2f}s]")
-
-            return result
-
-    # Mark the function as auto-isolate enabled
-    wrapper._auto_isolate = True
-    wrapper._source_file = source_file
-
-    return wrapper