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

comfy_env/__init__.py

@@ -34,6 +34,8 @@ from .env.security import (
 from .ipc.bridge import WorkerBridge
 from .ipc.worker import BaseWorker, register
 from .decorator import isolated, shutdown_all_processes
+from .isolation import enable_isolation
+from .stub_imports import setup_isolated_imports, cleanup_stubs
 
 # New in-place installation API
 from .install import install, verify_installation
@@ -144,6 +146,11 @@ __all__ = [
     # Legacy Decorator API
     "isolated",
     "shutdown_all_processes",
+    # New: Enable isolation for entire node pack
+    "enable_isolation",
+    # Import stubbing for isolated packages
+    "setup_isolated_imports",
+    "cleanup_stubs",
 ]
 
 # Add torch-based IPC if available

comfy_env/decorator.py

@@ -41,7 +41,7 @@ import time
 from dataclasses import dataclass
 from functools import wraps
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 logger = logging.getLogger("comfy_env")
 
@@ -178,7 +178,10 @@ def _get_or_create_worker(config: WorkerConfig, log_fn: Callable):
             # 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)
+            worker = TorchMPWorker(
+                name=config.env_name,
+                sys_path=config.sys_path,
+            )
         else:
             # Different Python - use PersistentVenvWorker
             from .workers.venv import PersistentVenvWorker
@@ -444,3 +447,254 @@ def isolated(
         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

comfy_env/env/config.py

@@ -154,6 +154,8 @@ class IsolatedEnv:
     worker_script: Optional[str] = None   # e.g., "worker.py" -> worker.py
     # Conda configuration (uses pixi backend when present)
     conda: Optional["CondaConfig"] = None
+    # Runtime isolation - run node FUNCTION methods in isolated subprocess
+    isolated: bool = False
 
     def __post_init__(self):
         """Validate and normalize configuration."""

comfy_env/env/config_file.py

@@ -572,6 +572,9 @@ def _parse_single_env(name: str, env_data: Dict[str, Any], base_dir: Path) -> Is
         elif isinstance(darwin_section, list):
             darwin_reqs = darwin_section
 
+    # Parse isolated flag for runtime process isolation
+    isolated = env_data.get("isolated", False)
+
     return IsolatedEnv(
         name=name,
         python=python,
@@ -583,6 +586,7 @@ def _parse_single_env(name: str, env_data: Dict[str, Any], base_dir: Path) -> Is
         linux_requirements=linux_reqs,
         darwin_requirements=darwin_reqs,
         conda=conda_config,
+        isolated=isolated,
     )
 
 

comfy_env/install.py

@@ -1,21 +1,18 @@
 """
 Installation API for comfy-env.
 
-This module provides the main `install()` function that handles both:
-- In-place installation (CUDA wheels into current environment)
-- Isolated installation (create separate venv with dependencies)
+This module provides the main `install()` function that handles:
+- Named environments [envname] → pixi (isolated Python environment)
+- Local packages [local] → uv/pip (in-place to current Python)
 
 Example:
     from comfy_env import install
 
-    # In-place install (auto-discovers config)
+    # Auto-discovers config and installs
     install()
 
-    # In-place with explicit config
+    # With explicit config path
     install(config="comfy-env.toml")
-
-    # Isolated environment
-    install(config="comfy-env.toml", mode="isolated")
 """
 
 import inspect
@@ -27,7 +24,6 @@ from typing import Callable, Dict, List, Optional, Set, Union
 
 from .env.config import IsolatedEnv, LocalConfig, NodeReq, SystemConfig
 from .env.config_file import load_config, discover_config
-from .env.manager import IsolatedEnvManager
 from .errors import CUDANotFoundError, InstallError
 from .pixi import pixi_install
 from .registry import PACKAGE_REGISTRY, get_cuda_short2
@@ -139,31 +135,36 @@ def _install_node_dependencies(
 
 
 def install(
+    config: Optional[Union[str, Path]] = None,
+    node_dir: Optional[Path] = None,
     log_callback: Optional[Callable[[str], None]] = None,
     dry_run: bool = False,
 ) -> bool:
     """
-    Install dependencies from comfy-env.toml, auto-discovered from caller's directory.
+    Install dependencies from comfy-env.toml.
 
     Example:
         from comfy_env import install
         install()
 
     Args:
+        config: Optional path to comfy-env.toml. Auto-discovered if not provided.
+        node_dir: Optional node directory. Auto-discovered from caller if not provided.
         log_callback: Optional callback for logging. Defaults to print.
         dry_run: If True, show what would be installed without installing.
 
     Returns:
         True if installation succeeded.
     """
-    # Auto-discover caller's directory
-    frame = inspect.stack()[1]
-    caller_file = frame.filename
-    node_dir = Path(caller_file).parent.resolve()
+    # Auto-discover caller's directory if not provided
+    if node_dir is None:
+        frame = inspect.stack()[1]
+        caller_file = frame.filename
+        node_dir = Path(caller_file).parent.resolve()
 
     log = log_callback or print
 
-    full_config = _load_full_config(None, node_dir)
+    full_config = _load_full_config(config, node_dir)
     if full_config is None:
         raise FileNotFoundError(
             f"No comfy-env.toml found in {node_dir}. "
@@ -177,28 +178,36 @@ def install(
         _install_system_packages(full_config.system, log, dry_run)
 
     env_config = full_config.default_env
-    if env_config is None and not full_config.has_local:
-        log("No packages to install")
-        return True
-
-    if env_config:
-        log(f"Found configuration: {env_config.name}")
-
-    if env_config and env_config.uses_conda:
-        log("Environment uses conda packages - using pixi backend")
-        return pixi_install(env_config, node_dir, log, dry_run)
 
     # Get user wheel_sources overrides
     user_wheel_sources = full_config.wheel_sources if hasattr(full_config, 'wheel_sources') else {}
 
     if env_config:
-        if env_config.python:
-            return _install_isolated(env_config, node_dir, log, dry_run)
-        else:
-            return _install_inplace(env_config, node_dir, log, dry_run, user_wheel_sources)
+        # Named environment → always pixi
+        log(f"Found environment: {env_config.name}")
+        python_ver = env_config.python or "3.11"  # Default to 3.11 if not specified
+        if not env_config.python:
+            log(f"  No Python version specified, defaulting to {python_ver}")
+            env_config = IsolatedEnv(
+                name=env_config.name,
+                python=python_ver,
+                cuda=env_config.cuda,
+                pytorch=env_config.pytorch,
+                requirements=env_config.requirements,
+                no_deps_requirements=env_config.no_deps_requirements,
+                linux_requirements=env_config.linux_requirements,
+                darwin_requirements=env_config.darwin_requirements,
+                windows_requirements=env_config.windows_requirements,
+                conda=env_config.conda,
+                isolated=env_config.isolated,
+            )
+        log(f"  Using pixi backend (Python {python_ver})")
+        return pixi_install(env_config, node_dir, log, dry_run)
     elif full_config.has_local:
+        # [local] section → uv in-place install
         return _install_local(full_config.local, node_dir, log, dry_run, user_wheel_sources)
     else:
+        log("No packages to install")
         return True
 
 
@@ -212,81 +221,6 @@ def _load_full_config(config: Optional[Union[str, Path]], node_dir: Path):
     return discover_config(node_dir)
 
 
-def _install_isolated(
-    env_config: IsolatedEnv,
-    node_dir: Path,
-    log: Callable[[str], None],
-    dry_run: bool,
-) -> bool:
-    """Install in isolated mode using IsolatedEnvManager."""
-    log(f"Installing in isolated mode: {env_config.name}")
-
-    if dry_run:
-        log("Dry run - would create isolated environment:")
-        log(f"  Python: {env_config.python}")
-        log(f"  CUDA: {env_config.cuda or 'auto-detect'}")
-        if env_config.requirements:
-            log(f"  Requirements: {len(env_config.requirements)} packages")
-        return True
-
-    manager = IsolatedEnvManager(base_dir=node_dir, log_callback=log)
-    env_dir = manager.setup(env_config)
-    log(f"Isolated environment ready: {env_dir}")
-    return True
-
-
-def _install_inplace(
-    env_config: IsolatedEnv,
-    node_dir: Path,
-    log: Callable[[str], None],
-    dry_run: bool,
-    user_wheel_sources: Dict[str, str],
-) -> bool:
-    """Install in-place into current environment."""
-    log("Installing in-place mode")
-
-    if sys.platform == "win32":
-        log("Installing MSVC runtime for Windows...")
-        if not dry_run:
-            _pip_install(["msvc-runtime"], no_deps=False, log=log)
-
-    env = RuntimeEnv.detect()
-    log(f"Detected environment: {env}")
-
-    if not env.cuda_version:
-        cuda_packages = env_config.no_deps_requirements or []
-        if cuda_packages:
-            raise CUDANotFoundError(package=", ".join(cuda_packages))
-
-    cuda_packages = env_config.no_deps_requirements or []
-    regular_packages = env_config.requirements or []
-
-    if dry_run:
-        log("\nDry run - would install:")
-        for req in cuda_packages:
-            package, version = parse_wheel_requirement(req)
-            url = _resolve_wheel_url(package, version, env, user_wheel_sources)
-            log(f"  {package}: {url[:80]}...")
-        if regular_packages:
-            log("  Regular packages:")
-            for pkg in regular_packages:
-                log(f"    {pkg}")
-        return True
-
-    if cuda_packages:
-        log(f"\nInstalling {len(cuda_packages)} CUDA packages...")
-        for req in cuda_packages:
-            package, version = parse_wheel_requirement(req)
-            _install_cuda_package(package, version, env, user_wheel_sources, log)
-
-    if regular_packages:
-        log(f"\nInstalling {len(regular_packages)} regular packages...")
-        _pip_install(regular_packages, no_deps=False, log=log)
-
-    log("\nInstallation complete!")
-    return True
-
-
 def _install_local(
     local_config: LocalConfig,
     node_dir: Path,