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

comfy_env/isolation.py

@@ -1,310 +0,0 @@
-"""
-Process isolation for ComfyUI node packs.
-
-This module provides enable_isolation() which wraps all node classes
-to run their FUNCTION methods in an isolated Python environment.
-
-Usage:
-    # In your node pack's __init__.py:
-    from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
-    from comfy_env import enable_isolation
-
-    enable_isolation(NODE_CLASS_MAPPINGS)  # That's it!
-
-This requires `isolated = true` in comfy-env.toml:
-
-    [myenv]
-    python = "3.11"
-    isolated = true
-
-    [myenv.packages]
-    requirements = ["my-package"]
-"""
-
-import atexit
-import inspect
-import os
-import sys
-import threading
-from functools import wraps
-from pathlib import Path
-from typing import Any, Dict, Optional
-
-# Debug logging (set COMFY_ENV_DEBUG=1 to enable)
-_DEBUG = os.environ.get("COMFY_ENV_DEBUG", "").lower() in ("1", "true", "yes")
-
-# Global worker cache (one per isolated environment)
-_workers: Dict[str, Any] = {}
-_workers_lock = threading.Lock()
-
-
-def _get_worker(
-    env_name: str,
-    python_path: Path,
-    working_dir: Path,
-    sys_path: list[str],
-):
-    """Get or create a persistent worker for the isolated environment."""
-    from .workers.venv import PersistentVenvWorker
-
-    cache_key = str(python_path)
-
-    with _workers_lock:
-        if cache_key in _workers:
-            worker = _workers[cache_key]
-            if worker.is_alive():
-                return worker
-            # Worker died, will recreate
-
-        print(f"[comfy-env] Starting isolated worker: {env_name}")
-        print(f"[comfy-env]   Python: {python_path}")
-
-        worker = PersistentVenvWorker(
-            python=str(python_path),
-            working_dir=working_dir,
-            sys_path=sys_path,
-            name=env_name,
-        )
-        _workers[cache_key] = worker
-        return worker
-
-
-def _shutdown_workers():
-    """Shutdown all cached workers. Called at exit."""
-    with _workers_lock:
-        for name, worker in _workers.items():
-            try:
-                worker.shutdown()
-            except Exception:
-                pass
-        _workers.clear()
-
-
-atexit.register(_shutdown_workers)
-
-
-def _find_python_path(node_dir: Path, env_name: str) -> Optional[Path]:
-    """
-    Find the Python executable for the isolated environment.
-
-    Priority:
-    1. .pixi/envs/default/bin/python (pixi/conda environment)
-    2. _env_{name}/bin/python (uv venv)
-    3. _env_{name}/Scripts/python.exe (Windows uv venv)
-    """
-    # Check pixi environment first
-    if sys.platform == "win32":
-        pixi_python = node_dir / ".pixi" / "envs" / "default" / "python.exe"
-    else:
-        pixi_python = node_dir / ".pixi" / "envs" / "default" / "bin" / "python"
-
-    if pixi_python.exists():
-        return pixi_python
-
-    # Check _env_* directory (uv venv)
-    env_dir = node_dir / f"_env_{env_name}"
-    if sys.platform == "win32":
-        env_python = env_dir / "Scripts" / "python.exe"
-    else:
-        env_python = env_dir / "bin" / "python"
-
-    if env_python.exists():
-        return env_python
-
-    return None
-
-
-def _wrap_node_class(
-    cls: type,
-    env_name: str,
-    python_path: Path,
-    working_dir: Path,
-    sys_path: list[str],
-) -> type:
-    """
-    Wrap a node class so its FUNCTION method runs in the isolated environment.
-
-    Args:
-        cls: The node class to wrap
-        env_name: Name of the isolated environment
-        python_path: Path to the isolated Python executable
-        working_dir: Working directory for the worker
-        sys_path: Additional paths to add to sys.path in the worker
-
-    Returns:
-        The wrapped class (modified in place)
-    """
-    func_name = getattr(cls, "FUNCTION", None)
-    if not func_name:
-        return cls  # Not a valid ComfyUI node class
-
-    original_method = getattr(cls, func_name, None)
-    if original_method is None:
-        return cls
-
-    # Get source file for the class
-    try:
-        source_file = Path(inspect.getfile(cls)).resolve()
-    except (TypeError, OSError):
-        # Can't get source file, skip wrapping
-        return cls
-
-    # Compute relative module path from working_dir
-    # e.g., /path/to/nodes/io/load_mesh.py -> nodes.io.load_mesh
-    try:
-        relative_path = source_file.relative_to(working_dir)
-        # Convert path to module: nodes/io/load_mesh.py -> nodes.io.load_mesh
-        module_name = str(relative_path.with_suffix("")).replace("/", ".").replace("\\", ".")
-    except ValueError:
-        # File not under working_dir, use stem as fallback
-        module_name = source_file.stem
-
-    @wraps(original_method)
-    def proxy(self, **kwargs):
-        if _DEBUG:
-            print(f"[comfy-env] PROXY CALLED: {cls.__name__}.{func_name}", flush=True)
-            print(f"[comfy-env]   kwargs keys: {list(kwargs.keys())}", flush=True)
-
-        worker = _get_worker(env_name, python_path, working_dir, sys_path)
-        if _DEBUG:
-            print(f"[comfy-env]   worker alive: {worker.is_alive()}", flush=True)
-
-        # Clone tensors for IPC if needed
-        try:
-            from .decorator import _clone_tensor_if_needed
-
-            kwargs = {k: _clone_tensor_if_needed(v) for k, v in kwargs.items()}
-        except ImportError:
-            pass  # No torch available, skip cloning
-
-        if _DEBUG:
-            print(f"[comfy-env]   calling worker.call_method...", flush=True)
-        result = worker.call_method(
-            module_name=module_name,
-            class_name=cls.__name__,
-            method_name=func_name,
-            self_state=self.__dict__.copy() if hasattr(self, "__dict__") else None,
-            kwargs=kwargs,
-            timeout=600.0,
-        )
-        if _DEBUG:
-            print(f"[comfy-env]   call_method returned", flush=True)
-
-        # Clone result tensors
-        try:
-            from .decorator import _clone_tensor_if_needed
-
-            result = _clone_tensor_if_needed(result)
-        except ImportError:
-            pass
-
-        return result
-
-    # Replace the method
-    setattr(cls, func_name, proxy)
-
-    # Mark as isolated for debugging
-    cls._comfy_env_isolated = True
-    cls._comfy_env_name = env_name
-
-    return cls
-
-
-def enable_isolation(node_class_mappings: Dict[str, type]) -> None:
-    """
-    Enable process isolation for all nodes in a node pack.
-
-    Call this AFTER importing NODE_CLASS_MAPPINGS. It wraps all node classes
-    so their FUNCTION methods run in the isolated Python environment specified
-    in comfy-env.toml.
-
-    Requires `isolated = true` in comfy-env.toml:
-
-        [myenv]
-        python = "3.11"
-        isolated = true
-
-    Args:
-        node_class_mappings: The NODE_CLASS_MAPPINGS dict from the node pack.
-
-    Example:
-        from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
-        from comfy_env import enable_isolation
-
-        enable_isolation(NODE_CLASS_MAPPINGS)
-
-        __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
-    """
-    # Skip if running inside worker subprocess
-    if os.environ.get("COMFYUI_ISOLATION_WORKER") == "1":
-        return
-
-    # Find the calling module's directory (node pack root)
-    frame = inspect.currentframe()
-    if frame is None:
-        print("[comfy-env] Warning: Could not get current frame")
-        return
-
-    caller_frame = frame.f_back
-    if caller_frame is None:
-        print("[comfy-env] Warning: Could not get caller frame")
-        return
-
-    caller_file = caller_frame.f_globals.get("__file__")
-    if not caller_file:
-        print("[comfy-env] Warning: Could not determine caller location")
-        return
-
-    node_dir = Path(caller_file).resolve().parent
-
-    # Load config
-    from .env.config_file import discover_config
-
-    config = discover_config(node_dir)
-    if not config:
-        print(f"[comfy-env] No comfy-env.toml found in {node_dir}")
-        return
-
-    # Find isolated environment
-    isolated_env = None
-    env_name = None
-
-    for name, env in config.envs.items():
-        if getattr(env, "isolated", False):
-            isolated_env = env
-            env_name = name
-            break
-
-    if not isolated_env or not env_name:
-        # No isolated env configured, silently return
-        return
-
-    # Find Python executable
-    python_path = _find_python_path(node_dir, env_name)
-
-    if not python_path:
-        print(f"[comfy-env] Warning: Isolated environment not found for '{env_name}'")
-        print(f"[comfy-env] Expected: .pixi/envs/default/bin/python or _env_{env_name}/bin/python")
-        print(f"[comfy-env] Run 'comfy-env install' to create the environment")
-        return
-
-    # Build sys.path for the worker
-    sys_path = [str(node_dir)]
-
-    # Add nodes directory if it exists
-    nodes_dir = node_dir / "nodes"
-    if nodes_dir.exists():
-        sys_path.append(str(nodes_dir))
-
-    print(f"[comfy-env] Enabling isolation for {len(node_class_mappings)} nodes")
-    print(f"[comfy-env]   Environment: {env_name}")
-    print(f"[comfy-env]   Python: {python_path}")
-
-    # Wrap all node classes
-    wrapped_count = 0
-    for node_name, node_cls in node_class_mappings.items():
-        if hasattr(node_cls, "FUNCTION"):
-            _wrap_node_class(node_cls, env_name, python_path, node_dir, sys_path)
-            wrapped_count += 1
-
-    print(f"[comfy-env] Wrapped {wrapped_count} node classes for isolation")