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

comfy_env/isolation/wrap.py

@@ -0,0 +1,351 @@
+"""
+Process isolation for ComfyUI node packs.
+
+This module provides wrap_isolated_nodes() which wraps node classes
+to run their FUNCTION methods in an isolated Python environment.
+
+Usage:
+    # In your node pack's __init__.py:
+    from pathlib import Path
+    from comfy_env import wrap_isolated_nodes
+
+    NODE_CLASS_MAPPINGS = {}
+
+    # Main nodes (no isolation)
+    from .nodes.main import NODE_CLASS_MAPPINGS as main_nodes
+    NODE_CLASS_MAPPINGS.update(main_nodes)
+
+    # Isolated nodes (has comfy-env.toml in that directory)
+    from .nodes.isolated import NODE_CLASS_MAPPINGS as isolated_nodes
+    NODE_CLASS_MAPPINGS.update(
+        wrap_isolated_nodes(isolated_nodes, Path(__file__).parent / "nodes/isolated")
+    )
+"""
+
+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_isolated_python_version(env_dir: Path) -> Optional[str]:
+    """Get Python version from isolated environment."""
+    if sys.platform == "win32":
+        python_path = env_dir / "python.exe"
+    else:
+        python_path = env_dir / "bin" / "python"
+
+    if not python_path.exists():
+        return None
+
+    import subprocess
+    try:
+        result = subprocess.run(
+            [str(python_path), "-c", "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')"],
+            capture_output=True, text=True, timeout=5
+        )
+        if result.returncode == 0:
+            return result.stdout.strip()
+    except Exception:
+        pass
+    return None
+
+
+def _get_worker(
+    env_dir: Path,
+    working_dir: Path,
+    sys_path: list[str],
+    lib_path: Optional[str] = None,
+):
+    """Get or create a persistent worker for the isolated environment."""
+    cache_key = str(env_dir)
+
+    with _workers_lock:
+        if cache_key in _workers:
+            worker = _workers[cache_key]
+            if worker.is_alive():
+                return worker
+            # Worker died, will recreate
+
+        # Check if Python versions match
+        host_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+        isolated_version = _get_isolated_python_version(env_dir)
+
+        if isolated_version and isolated_version != host_version:
+            # Different Python version - must use SubprocessWorker
+            from ..workers.subprocess import SubprocessWorker
+
+            if sys.platform == "win32":
+                python_path = env_dir / "python.exe"
+            else:
+                python_path = env_dir / "bin" / "python"
+
+            print(f"[comfy-env] Starting isolated worker (SubprocessWorker)")
+            print(f"[comfy-env]   Python: {python_path} ({isolated_version} vs host {host_version})")
+
+            worker = SubprocessWorker(
+                python=str(python_path),
+                working_dir=working_dir,
+                sys_path=sys_path,
+                name=working_dir.name,
+            )
+        else:
+            # Same Python version - use MPWorker (faster)
+            from ..workers.mp import MPWorker
+
+            print(f"[comfy-env] Starting isolated worker (MPWorker)")
+            print(f"[comfy-env]   Env: {env_dir}")
+
+            worker = MPWorker(
+                name=working_dir.name,
+                sys_path=sys_path,
+                lib_path=lib_path,
+            )
+
+        _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_env_paths(node_dir: Path) -> tuple[Optional[Path], Optional[Path]]:
+    """
+    Find site-packages and lib directories for the isolated environment.
+
+    Returns:
+        (site_packages, lib_dir) - lib_dir is for LD_LIBRARY_PATH
+    """
+    import glob
+
+    # Check pixi environment first
+    pixi_env = node_dir / ".pixi" / "envs" / "default"
+    if pixi_env.exists():
+        # Find site-packages (pythonX.Y varies)
+        if sys.platform == "win32":
+            site_packages = pixi_env / "Lib" / "site-packages"
+            lib_dir = pixi_env / "Library" / "bin"  # Windows DLLs
+        else:
+            pattern = str(pixi_env / "lib" / "python*" / "site-packages")
+            matches = glob.glob(pattern)
+            site_packages = Path(matches[0]) if matches else None
+            lib_dir = pixi_env / "lib"
+        if site_packages and site_packages.exists():
+            return site_packages, lib_dir if lib_dir.exists() else None
+
+    # Check .venv directory
+    venv_dir = node_dir / ".venv"
+    if venv_dir.exists():
+        if sys.platform == "win32":
+            site_packages = venv_dir / "Lib" / "site-packages"
+        else:
+            pattern = str(venv_dir / "lib" / "python*" / "site-packages")
+            matches = glob.glob(pattern)
+            site_packages = Path(matches[0]) if matches else None
+        if site_packages and site_packages.exists():
+            return site_packages, None  # venvs don't have separate lib
+
+    return None, None
+
+
+def _find_env_dir(node_dir: Path) -> Optional[Path]:
+    """Find the environment directory (for cache key)."""
+    pixi_env = node_dir / ".pixi" / "envs" / "default"
+    if pixi_env.exists():
+        return pixi_env
+    venv_dir = node_dir / ".venv"
+    if venv_dir.exists():
+        return venv_dir
+    return None
+
+
+def _wrap_node_class(
+    cls: type,
+    env_dir: Path,
+    working_dir: Path,
+    sys_path: list[str],
+    lib_path: Optional[str] = None,
+) -> type:
+    """
+    Wrap a node class so its FUNCTION method runs in the isolated environment.
+
+    Args:
+        cls: The node class to wrap
+        env_dir: Path to the isolated environment directory
+        working_dir: Working directory for the worker
+        sys_path: Additional paths to add to sys.path in the worker
+        lib_path: Path to add to LD_LIBRARY_PATH for conda libraries
+
+    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_dir, working_dir, sys_path, lib_path)
+        if _DEBUG:
+            print(f"[comfy-env]   worker alive: {worker.is_alive()}", flush=True)
+
+        # Clone tensors for IPC if needed
+        try:
+            from ..workers.tensor_utils import prepare_for_ipc_recursive
+
+            kwargs = {k: prepare_for_ipc_recursive(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 ..workers.tensor_utils import prepare_for_ipc_recursive
+
+            result = prepare_for_ipc_recursive(result)
+        except ImportError:
+            pass
+
+        return result
+
+    # Replace the method
+    setattr(cls, func_name, proxy)
+
+    # Mark as isolated for debugging
+    cls._comfy_env_isolated = True
+
+    return cls
+
+
+def wrap_isolated_nodes(
+    node_class_mappings: Dict[str, type],
+    nodes_dir: Path,
+) -> Dict[str, type]:
+    """
+    Wrap nodes from a directory that has a comfy-env.toml.
+
+    This is the directory-based isolation API. Call it for each subdirectory
+    of nodes/ that has a comfy-env.toml.
+
+    Args:
+        node_class_mappings: The NODE_CLASS_MAPPINGS dict from the nodes in this dir.
+        nodes_dir: The directory containing comfy-env.toml and the node files.
+
+    Returns:
+        The same dict with node classes wrapped for isolation.
+
+    Example:
+        # __init__.py
+        from comfy_env import wrap_isolated_nodes
+        from pathlib import Path
+
+        NODE_CLASS_MAPPINGS = {}
+
+        # Native nodes (no isolation)
+        from .nodes.main import NODE_CLASS_MAPPINGS as main_nodes
+        NODE_CLASS_MAPPINGS.update(main_nodes)
+
+        # Isolated nodes (has comfy-env.toml)
+        from .nodes.cgal import NODE_CLASS_MAPPINGS as cgal_nodes
+        NODE_CLASS_MAPPINGS.update(
+            wrap_isolated_nodes(cgal_nodes, Path(__file__).parent / "nodes/cgal")
+        )
+    """
+    # Skip if running inside worker subprocess
+    if os.environ.get("COMFYUI_ISOLATION_WORKER") == "1":
+        return node_class_mappings
+
+    nodes_dir = Path(nodes_dir).resolve()
+
+    # Check for comfy-env.toml
+    config_file = nodes_dir / "comfy-env.toml"
+    if not config_file.exists():
+        print(f"[comfy-env] Warning: No comfy-env.toml in {nodes_dir}")
+        return node_class_mappings
+
+    # Find environment directory and paths
+    env_dir = _find_env_dir(nodes_dir)
+    site_packages, lib_dir = _find_env_paths(nodes_dir)
+
+    if not env_dir or not site_packages:
+        print(f"[comfy-env] Warning: Isolated environment not found")
+        print(f"[comfy-env] Expected: .pixi/envs/default or .venv")
+        print(f"[comfy-env] Run 'comfy-env install' in {nodes_dir}")
+        return node_class_mappings
+
+    # Build sys.path for the worker - site-packages first, then node dir
+    sys_path = [str(site_packages), str(nodes_dir)]
+
+    # lib_dir for LD_LIBRARY_PATH (conda libraries)
+    lib_path = str(lib_dir) if lib_dir else None
+
+    print(f"[comfy-env] Wrapping {len(node_class_mappings)} nodes from {nodes_dir.name}")
+    print(f"[comfy-env]   site-packages: {site_packages}")
+    if lib_path:
+        print(f"[comfy-env]   lib: {lib_path}")
+
+    # Wrap all node classes
+    for node_name, node_cls in node_class_mappings.items():
+        if hasattr(node_cls, "FUNCTION"):
+            _wrap_node_class(node_cls, env_dir, nodes_dir, sys_path, lib_path)
+
+    return node_class_mappings

comfy_env/nodes.py

@@ -16,7 +16,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Callable, List, Set
 
 if TYPE_CHECKING:
-    from .env.config import NodeReq
+    from .config.types import NodeReq
 
 
 def normalize_repo_url(repo: str) -> str:
@@ -111,7 +111,7 @@ def install_node_deps(
         log: Logging callback
         visited: Set of already-processed node names (for cycle detection)
     """
-    from .env.config_file import discover_config
+    from .config.parser import discover_config
 
     for req in node_reqs:
         # Skip if already visited (cycle detection)

comfy_env/pixi/__init__.py

@@ -0,0 +1,48 @@
+"""
+Pixi integration for comfy-env.
+
+All dependencies go through pixi for unified management.
+"""
+
+from .core import (
+    ensure_pixi,
+    get_pixi_path,
+    get_pixi_python,
+    pixi_run,
+    pixi_install,
+    clean_pixi_artifacts,
+    CUDA_WHEELS_INDEX,
+)
+from .registry import PACKAGE_REGISTRY
+from .cuda_detection import (
+    detect_cuda_version,
+    detect_cuda_environment,
+    detect_gpu_info,
+    detect_gpus,
+    get_gpu_summary,
+    get_recommended_cuda_version,
+    GPUInfo,
+    CUDAEnvironment,
+)
+
+__all__ = [
+    # Core pixi functions
+    "ensure_pixi",
+    "get_pixi_path",
+    "get_pixi_python",
+    "pixi_run",
+    "pixi_install",
+    "clean_pixi_artifacts",
+    "CUDA_WHEELS_INDEX",
+    # Registry
+    "PACKAGE_REGISTRY",
+    # CUDA detection
+    "detect_cuda_version",
+    "detect_cuda_environment",
+    "detect_gpu_info",
+    "detect_gpus",
+    "get_gpu_summary",
+    "get_recommended_cuda_version",
+    "GPUInfo",
+    "CUDAEnvironment",
+]