comfy-env 0.0.49__py3-none-any.whl → 0.0.50__py3-none-any.whl

comfy_env/isolation.py

@@ -0,0 +1,297 @@
+"""
+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
+
+# 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):
+        worker = _get_worker(env_name, python_path, working_dir, sys_path)
+
+        # 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
+
+        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,
+        )
+
+        # 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")

comfy_env/pixi.py

@@ -10,12 +10,13 @@ See: https://pixi.sh/
 
 import os
 import platform
+import re
 import shutil
 import stat
 import subprocess
 import sys
 from pathlib import Path
-from typing import Callable, List, Optional
+from typing import Callable, List, Optional, Tuple
 
 from .env.config import IsolatedEnv, CondaConfig
 
@@ -144,6 +145,42 @@ def ensure_pixi(
     return pixi_path
 
 
+def _parse_pypi_requirement(dep: str) -> Tuple[str, Optional[str], List[str]]:
+    """
+    Parse a pip requirement into (name, version_spec, extras).
+
+    Examples:
+        "trimesh[easy]>=4.0.0" -> ("trimesh", ">=4.0.0", ["easy"])
+        "numpy>=1.21.0" -> ("numpy", ">=1.21.0", [])
+        "torch" -> ("torch", None, [])
+        "pkg[a,b]" -> ("pkg", None, ["a", "b"])
+
+    Returns:
+        Tuple of (package_name, version_spec_or_None, list_of_extras)
+    """
+    dep = dep.strip()
+
+    # Match: name[extras]version_spec or name version_spec
+    # Package names can contain letters, numbers, underscores, hyphens, and dots
+    match = re.match(r'^([a-zA-Z0-9._-]+)(?:\[([^\]]+)\])?(.*)$', dep)
+    if not match:
+        return dep, None, []
+
+    name = match.group(1)
+    extras_str = match.group(2)
+    version_spec = match.group(3).strip() if match.group(3) else None
+
+    extras = []
+    if extras_str:
+        extras = [e.strip() for e in extras_str.split(',')]
+
+    # Return None instead of empty string for version_spec
+    if version_spec == "":
+        version_spec = None
+
+    return name, version_spec, extras
+
+
 def create_pixi_toml(
     env_config: IsolatedEnv,
     node_dir: Path,
@@ -227,7 +264,7 @@ def create_pixi_toml(
     special_deps = {}  # For dependencies that need special syntax (path, etc.)
 
     # Always include comfy-env for worker support
-    # Use local wheel if available (for ct test --local)
+    # Priority: 1. COMFY_LOCAL_WHEELS env var, 2. ~/utils/comfy-env, 3. PyPI
     local_wheels_dir = os.environ.get("COMFY_LOCAL_WHEELS")
     if local_wheels_dir:
         local_wheels = list(Path(local_wheels_dir).glob("comfy_env-*.whl"))
@@ -238,15 +275,78 @@ def create_pixi_toml(
         else:
             pypi_deps.append("comfy-env")
     else:
-        pypi_deps.append("comfy-env")
+        # Check for local editable comfy-env at ~/utils/comfy-env
+        local_comfy_env = Path.home() / "utils" / "comfy-env"
+        if local_comfy_env.exists() and (local_comfy_env / "pyproject.toml").exists():
+            special_deps["comfy-env"] = f'{{ path = "{local_comfy_env}", editable = true }}'
+        else:
+            pypi_deps.append("comfy-env")
 
     # Add regular requirements
     if env_config.requirements:
         pypi_deps.extend(env_config.requirements)
 
-    # NOTE: CUDA packages (no_deps_requirements) are NOT added here.
-    # They require special wheel resolution and are installed separately
-    # by the comfy-env resolver after pixi install completes.
+    # Add CUDA packages with resolved wheel URLs
+    if env_config.no_deps_requirements:
+        from .registry import PACKAGE_REGISTRY
+
+        # Use fixed CUDA 12.8 / PyTorch 2.8 for pixi environments (modern GPU default)
+        # This ensures wheels match what pixi will install, not what the host has
+        vars_dict = {
+            "cuda_version": "12.8",
+            "cuda_short": "128",
+            "cuda_short2": "128",
+            "cuda_major": "12",
+            "torch_version": "2.8.0",
+            "torch_short": "280",
+            "torch_mm": "28",
+            "torch_dotted_mm": "2.8",
+        }
+
+        # Platform detection
+        if sys.platform == "linux":
+            vars_dict["platform"] = "linux_x86_64"
+        elif sys.platform == "darwin":
+            vars_dict["platform"] = "macosx_arm64" if platform.machine() == "arm64" else "macosx_x86_64"
+        elif sys.platform == "win32":
+            vars_dict["platform"] = "win_amd64"
+
+        # Python version from pixi env config
+        if env_config.python:
+            py_parts = env_config.python.split(".")
+            py_major = py_parts[0]
+            py_minor = py_parts[1] if len(py_parts) > 1 else "0"
+            vars_dict["py_version"] = env_config.python
+            vars_dict["py_short"] = f"{py_major}{py_minor}"
+            vars_dict["py_minor"] = py_minor
+            vars_dict["py_tag"] = f"cp{py_major}{py_minor}"
+
+        for req in env_config.no_deps_requirements:
+            # Parse requirement (e.g., "cumesh" or "cumesh==0.0.1")
+            if "==" in req:
+                pkg_name, version = req.split("==", 1)
+            else:
+                pkg_name = req
+                version = None
+
+            pkg_lower = pkg_name.lower()
+            if pkg_lower in PACKAGE_REGISTRY:
+                config = PACKAGE_REGISTRY[pkg_lower]
+                template = config.get("wheel_template")
+                if template:
+                    # Use version from requirement or default
+                    v = version or config.get("default_version")
+                    if v:
+                        vars_dict["version"] = v
+
+                    # Resolve URL
+                    url = template
+                    for key, value in vars_dict.items():
+                        if value:
+                            url = url.replace(f"{{{key}}}", str(value))
+
+                    special_deps[pkg_name] = f'{{ url = "{url}" }}'
+                    log(f"  CUDA package {pkg_name}: resolved wheel URL")
 
     # Add platform-specific requirements
     if sys.platform == "linux" and env_config.linux_requirements:
@@ -265,25 +365,23 @@ def create_pixi_toml(
 
         for dep in pypi_deps:
             # Parse pip requirement format to pixi format
-            dep_clean = dep.strip()
-            if ">=" in dep_clean:
-                name, version = dep_clean.split(">=", 1)
-                # Handle complex version specs like ">=1.0,<2.0"
-                name = name.strip()
-                version = version.strip()
-                lines.append(f'{name} = ">={version}"')
-            elif "==" in dep_clean:
-                name, version = dep_clean.split("==", 1)
-                lines.append(f'{name.strip()} = "=={version.strip()}"')
-            elif ">" in dep_clean:
-                name, version = dep_clean.split(">", 1)
-                lines.append(f'{name.strip()} = ">{version.strip()}"')
-            elif "<" in dep_clean:
-                name, version = dep_clean.split("<", 1)
-                lines.append(f'{name.strip()} = "<{version.strip()}"')
+            # Handles extras like trimesh[easy]>=4.0.0
+            name, version_spec, extras = _parse_pypi_requirement(dep)
+
+            if extras:
+                # Use table syntax for packages with extras
+                # e.g., trimesh = { version = ">=4.0.0", extras = ["easy"] }
+                extras_json = "[" + ", ".join(f'"{e}"' for e in extras) + "]"
+                if version_spec:
+                    lines.append(f'{name} = {{ version = "{version_spec}", extras = {extras_json} }}')
+                else:
+                    lines.append(f'{name} = {{ version = "*", extras = {extras_json} }}')
             else:
-                # No version spec
-                lines.append(f'{dep_clean} = "*"')
+                # Simple syntax for packages without extras
+                if version_spec:
+                    lines.append(f'{name} = "{version_spec}"')
+                else:
+                    lines.append(f'{name} = "*"')
 
     content = "\n".join(lines) + "\n"
 
@@ -370,6 +468,8 @@ def pixi_install(
             log(f"  - Install {len(env_config.conda.packages)} conda packages")
         if env_config.requirements:
             log(f"  - Install {len(env_config.requirements)} pip packages")
+        if env_config.no_deps_requirements:
+            log(f"  - Install {len(env_config.no_deps_requirements)} CUDA packages: {', '.join(env_config.no_deps_requirements)}")
         return True
 
     # Clean previous pixi artifacts