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

comfy_env/env/manager.py

@@ -1,636 +0,0 @@
-"""
-IsolatedEnvManager - Creates and manages isolated Python environments.
-
-Uses uv for fast environment creation and package installation.
-"""
-
-import os
-import subprocess
-import shutil
-import sys
-from pathlib import Path
-from typing import Optional, Callable
-
-from .config import IsolatedEnv
-from .platform import get_platform, PlatformProvider
-from .cuda_gpu_detection import detect_cuda_version
-from .security import (
-    normalize_env_name,
-    validate_dependency,
-    validate_dependencies,
-    validate_path_within_root,
-    validate_wheel_url,
-)
-from ..registry import PACKAGE_REGISTRY, is_registered, get_cuda_short2
-from ..resolver import RuntimeEnv, parse_wheel_requirement
-
-
-class IsolatedEnvManager:
-    """
-    Manages isolated Python environments for ComfyUI nodes.
-
-    This class handles:
-    - Creating Python virtual environments using uv
-    - Installing PyTorch with correct CUDA version
-    - Installing packages from requirements or wheel sources
-    - Caching environments (same config = reuse existing)
-    - Platform-specific handling (Windows DLLs, etc.)
-
-    Example:
-        manager = IsolatedEnvManager(base_dir=Path("./"))
-
-        env = IsolatedEnv(
-            name="my-node",
-            python="3.10",
-            cuda="12.8",
-            requirements=["torch==2.8.0", "nvdiffrast"],
-        )
-
-        # Create environment and install dependencies
-        env_path = manager.setup(env)
-
-        # Get Python executable path
-        python_exe = manager.get_python(env)
-    """
-
-    def __init__(
-        self,
-        base_dir: Path,
-        log_callback: Optional[Callable[[str], None]] = None,
-    ):
-        """
-        Initialize environment manager.
-
-        Args:
-            base_dir: Directory where environments will be created
-            log_callback: Optional callback for logging (default: print)
-        """
-        self.base_dir = Path(base_dir)
-        self.platform: PlatformProvider = get_platform()
-        self.log = log_callback or print
-
-        # Check platform compatibility
-        is_compatible, error = self.platform.check_prerequisites()
-        if not is_compatible:
-            raise RuntimeError(f"Platform incompatible: {error}")
-
-    def _find_uv(self) -> Optional[Path]:
-        """Find uv executable in PATH or current Python environment."""
-        # First check system PATH
-        uv_path = shutil.which("uv")
-        if uv_path:
-            return Path(uv_path)
-
-        # Check current Python environment's Scripts/bin folder
-        import sys
-        if sys.platform == "win32":
-            candidates = [
-                Path(sys.prefix) / "Scripts" / "uv.exe",
-                Path(sys.base_prefix) / "Scripts" / "uv.exe",
-            ]
-        else:
-            candidates = [
-                Path(sys.prefix) / "bin" / "uv",
-                Path(sys.base_prefix) / "bin" / "uv",
-            ]
-
-        for candidate in candidates:
-            if candidate.exists():
-                return candidate
-
-        return None
-
-    def _run_uv(self, args: list, env_dir: Optional[Path] = None, **kwargs) -> subprocess.CompletedProcess:
-        """Run uv command."""
-        uv = self._find_uv()
-        if not uv:
-            raise RuntimeError(
-                "uv not found. Please install it:\n"
-                "  curl -LsSf https://astral.sh/uv/install.sh | sh\n"
-                "  # or on Windows: powershell -c \"irm https://astral.sh/uv/install.ps1 | iex\""
-            )
-
-        cmd = [str(uv)] + args
-        return subprocess.run(cmd, capture_output=True, text=True, **kwargs)
-
-    def get_env_dir(self, env: IsolatedEnv) -> Path:
-        """Get the environment directory path for a config."""
-        # Validate environment name to prevent directory traversal
-        safe_name = normalize_env_name(env.name)
-        env_dir = env.get_default_env_dir(self.base_dir)
-
-        # Ensure the resulting path is within base_dir
-        validate_path_within_root(env_dir, self.base_dir)
-
-        return env_dir
-
-    def get_python(self, env: IsolatedEnv) -> Path:
-        """Get the Python executable path for an environment."""
-        env_dir = self.get_env_dir(env)
-        return self.platform.get_env_paths(env_dir, env.python).python
-
-    def get_pip(self, env: IsolatedEnv) -> Path:
-        """Get the pip executable path for an environment."""
-        env_dir = self.get_env_dir(env)
-        return self.platform.get_env_paths(env_dir, env.python).pip
-
-    def exists(self, env: IsolatedEnv) -> bool:
-        """Check if environment already exists."""
-        python_exe = self.get_python(env)
-        return python_exe.exists()
-
-    def is_ready(self, env: IsolatedEnv, verify_packages: Optional[list] = None) -> bool:
-        """
-        Check if environment is ready to use.
-
-        Args:
-            env: Environment configuration
-            verify_packages: Optional list of packages to verify (e.g., ["torch", "numpy"])
-
-        Returns:
-            True if environment exists and packages are importable
-        """
-        python_exe = self.get_python(env)
-        if not python_exe.exists():
-            return False
-
-        if verify_packages:
-            imports = ", ".join(verify_packages)
-            try:
-                result = subprocess.run(
-                    [str(python_exe), "-c", f"import {imports}"],
-                    capture_output=True,
-                    text=True,
-                    timeout=30
-                )
-                return result.returncode == 0
-            except (subprocess.SubprocessError, OSError):
-                return False
-
-        return True
-
-    def create_venv(self, env: IsolatedEnv) -> Path:
-        """
-        Create a virtual environment using uv.
-
-        Args:
-            env: Environment configuration
-
-        Returns:
-            Path to the environment directory
-        """
-        env_dir = self.get_env_dir(env)
-
-        if env_dir.exists():
-            self.log(f"Environment already exists: {env_dir}")
-            return env_dir
-
-        self.log(f"Creating environment: {env_dir}")
-        self.log(f"  Python: {env.python}")
-
-        result = self._run_uv([
-            "venv",
-            str(env_dir),
-            "--python", env.python,
-        ])
-
-        if result.returncode != 0:
-            raise RuntimeError(f"Failed to create venv: {result.stderr}")
-
-        self.log(f"Environment created successfully")
-        return env_dir
-
-    def install_pytorch(self, env: IsolatedEnv) -> None:
-        """
-        Install PyTorch with appropriate CUDA version.
-
-        Args:
-            env: Environment configuration
-        """
-        cuda_version = env.cuda
-        if cuda_version is None:
-            # Auto-detect CUDA version
-            cuda_version = detect_cuda_version()
-            if cuda_version:
-                self.log(f"Auto-detected CUDA version: {cuda_version}")
-            else:
-                self.log("No GPU detected, installing CPU-only PyTorch")
-
-        python_exe = self.get_python(env)
-
-        # Determine PyTorch index URL
-        if cuda_version:
-            cuda_short = cuda_version.replace(".", "")
-            index_url = f"https://download.pytorch.org/whl/cu{cuda_short}"
-            self.log(f"Installing PyTorch with CUDA {cuda_version}")
-        else:
-            index_url = "https://download.pytorch.org/whl/cpu"
-            self.log("Installing PyTorch (CPU)")
-
-        # Build uv pip command
-        uv = self._find_uv()
-        pip_args = [
-            str(uv), "pip", "install",
-            "--python", str(python_exe),
-            "--index-url", index_url,
-        ]
-
-        # Install torch + torchvision together from same index to ensure ABI compatibility
-        if env.pytorch_version:
-            pip_args.append(f"torch=={env.pytorch_version}")
-        else:
-            pip_args.append("torch")
-
-        # Always install torchvision from the same index - uv will resolve matching version
-        pip_args.append("torchvision")
-
-        result = subprocess.run(pip_args, capture_output=True, text=True)
-
-        if result.returncode != 0:
-            raise RuntimeError(f"Failed to install PyTorch: {result.stderr}")
-
-        self.log("PyTorch + torchvision installed successfully")
-
-    def install_requirements(
-        self,
-        env: IsolatedEnv,
-        extra_args: Optional[list] = None,
-    ) -> None:
-        """
-        Install requirements for an environment.
-
-        Args:
-            env: Environment configuration
-            extra_args: Extra pip arguments
-        """
-        python_exe = self.get_python(env)
-
-        # Merge platform-specific requirements
-        if sys.platform == 'win32':
-            platform_reqs = env.windows_requirements
-            platform_name = 'Windows'
-        elif sys.platform == 'darwin':
-            platform_reqs = env.darwin_requirements
-            platform_name = 'macOS'
-        else:
-            platform_reqs = env.linux_requirements
-            platform_name = 'Linux'
-
-        all_requirements = list(env.requirements) + list(platform_reqs)
-
-        if platform_reqs:
-            self.log(f"Including {len(platform_reqs)} {platform_name}-specific packages")
-
-        if not all_requirements and not env.requirements_file:
-            self.log("No requirements to install")
-            return
-
-        # Validate requirements for security
-        if all_requirements:
-            validate_dependencies(all_requirements)
-
-        # Validate wheel sources
-        for wheel_source in env.wheel_sources:
-            validate_wheel_url(wheel_source)
-
-        # Validate index URLs
-        for index_url in env.index_urls:
-            validate_wheel_url(index_url)
-
-        uv = self._find_uv()
-        pip_args = [str(uv), "pip", "install", "--python", str(python_exe)]
-
-        # Add wheel sources as --find-links
-        for wheel_source in env.wheel_sources:
-            pip_args.extend(["--find-links", wheel_source])
-
-        # Add extra index URLs
-        for index_url in env.index_urls:
-            pip_args.extend(["--extra-index-url", index_url])
-
-        # Add extra args
-        if extra_args:
-            pip_args.extend(extra_args)
-
-        # Install from requirements file
-        if env.requirements_file and env.requirements_file.exists():
-            self.log(f"Installing from {env.requirements_file}")
-            result = subprocess.run(
-                pip_args + ["-r", str(env.requirements_file)],
-                capture_output=True,
-                text=True,
-            )
-            if result.returncode != 0:
-                raise RuntimeError(f"Failed to install requirements: {result.stderr}")
-
-        # Install no-deps requirements first (e.g., CUDA extensions with conflicting metadata)
-        # For packages in registry, resolve proper wheel URLs
-        if env.no_deps_requirements:
-            self.log(f"Installing {len(env.no_deps_requirements)} CUDA packages")
-            self._install_cuda_packages(env, pip_args)
-
-        # Install individual requirements (including platform-specific)
-        if all_requirements:
-            self.log(f"Installing {len(all_requirements)} packages")
-            result = subprocess.run(
-                pip_args + all_requirements,
-                capture_output=True,
-                text=True,
-            )
-            if result.returncode != 0:
-                raise RuntimeError(f"Failed to install packages: {result.stderr}")
-
-        self.log("Requirements installed successfully")
-
-    def _install_cuda_packages(self, env: IsolatedEnv, pip_args: list) -> None:
-        """Install CUDA packages using registry-based wheel resolution."""
-        import platform as plat
-
-        # Build RuntimeEnv from IsolatedEnv config
-        os_name = plat.system().lower()
-        if os_name == 'darwin':
-            platform_tag = f"macosx_{plat.mac_ver()[0].replace('.', '_')}_{plat.machine()}"
-        elif os_name == 'windows':
-            platform_tag = f"win_{plat.machine().lower()}"
-        else:
-            platform_tag = f"linux_{plat.machine()}"
-
-        cuda_short = env.cuda.replace(".", "") if env.cuda else None
-        torch_short = env.pytorch_version.replace(".", "") if env.pytorch_version else None
-        torch_mm = "".join(env.pytorch_version.split(".")[:2]) if env.pytorch_version else None
-
-        runtime_env = RuntimeEnv(
-            os_name=os_name,
-            platform_tag=platform_tag,
-            python_version=env.python,
-            python_short=env.python.replace(".", ""),
-            cuda_version=env.cuda,
-            cuda_short=cuda_short,
-            torch_version=env.pytorch_version,
-            torch_short=torch_short,
-            torch_mm=torch_mm,
-        )
-        vars_dict = runtime_env.as_dict()
-        if env.cuda:
-            vars_dict["cuda_short2"] = get_cuda_short2(env.cuda)
-        vars_dict["py_tag"] = f"cp{env.python.replace('.', '')}"
-
-        for req in env.no_deps_requirements:
-            package, version = parse_wheel_requirement(req)
-            pkg_lower = package.lower()
-
-            if pkg_lower in PACKAGE_REGISTRY:
-                config = PACKAGE_REGISTRY[pkg_lower]
-
-                if "wheel_template" in config:
-                    # Direct wheel URL from template
-                    template = config["wheel_template"]
-
-                    # Only require version if template uses {version}
-                    effective_version = None
-                    if "{version}" in template:
-                        effective_version = version or config.get("default_version")
-                        if not effective_version:
-                            raise RuntimeError(f"Package {package} requires version (no default in registry)")
-                        vars_dict["version"] = effective_version
-
-                    wheel_url = self._substitute_template(template, vars_dict)
-                    version_str = f"=={effective_version}" if effective_version else ""
-                    self.log(f"  Installing {package}{version_str}...")
-                    self.log(f"    URL: {wheel_url}")
-                    result = subprocess.run(
-                        pip_args + ["--no-deps", wheel_url],
-                        capture_output=True, text=True,
-                    )
-                    if result.returncode != 0:
-                        raise RuntimeError(f"Failed to install {package}: {result.stderr}")
-
-                elif "find_links" in config:
-                    # find_links: pip resolves wheel from index URL
-                    find_links_url = config["find_links"]
-                    effective_version = version or config.get("default_version")
-                    pkg_spec = f"{package}=={effective_version}" if effective_version else package
-                    self.log(f"  Installing {pkg_spec} from {find_links_url}...")
-                    result = subprocess.run(
-                        pip_args + ["--no-deps", "--find-links", find_links_url, pkg_spec],
-                        capture_output=True, text=True,
-                    )
-                    if result.returncode != 0:
-                        raise RuntimeError(f"Failed to install {package}: {result.stderr}")
-
-                elif "package_name" in config:
-                    # PyPI variant (e.g., spconv-cu124)
-                    pkg_name = self._substitute_template(config["package_name"], vars_dict)
-                    pkg_spec = f"{pkg_name}=={version}" if version else pkg_name
-                    self.log(f"  Installing {package} as {pkg_spec}...")
-                    result = subprocess.run(
-                        pip_args + [pkg_spec],
-                        capture_output=True, text=True,
-                    )
-                    if result.returncode != 0:
-                        raise RuntimeError(f"Failed to install {package}: {result.stderr}")
-
-                else:
-                    raise RuntimeError(f"Package {package} in registry but missing wheel_template, find_links, or package_name")
-
-            else:
-                # Not in registry - error
-                raise RuntimeError(
-                    f"Package {package} not found in registry. "
-                    f"Add wheel_template to [wheel_sources] in comfy-env.toml"
-                )
-
-    def _substitute_template(self, template: str, vars_dict: dict) -> str:
-        """Substitute template variables with environment values."""
-        result = template
-        for key, value in vars_dict.items():
-            if value is not None:
-                result = result.replace(f"{{{key}}}", str(value))
-        return result
-
-    def _install_comfy_env(self, env: IsolatedEnv) -> None:
-        """Install comfy-env package (needed for BaseWorker)."""
-        python_exe = self.get_python(env)
-        uv = self._find_uv()
-
-        self.log("Installing comfy-env (for worker support)...")
-
-        # Check for local wheel in COMFY_LOCAL_WHEELS directory (set by comfy-test --local)
-        install_target = "comfy-env @ git+https://github.com/PozzettiAndrea/comfy-env"
-        wheels_dir = os.environ.get("COMFY_LOCAL_WHEELS")
-        self.log(f"  COMFY_LOCAL_WHEELS={wheels_dir}")
-        if wheels_dir:
-            wheels_path = Path(wheels_dir)
-            self.log(f"  Wheels dir exists: {wheels_path.exists()}")
-            if wheels_path.exists():
-                local_wheels = list(wheels_path.glob("comfy_env-*.whl"))
-                self.log(f"  Found wheels: {[w.name for w in local_wheels]}")
-                if local_wheels:
-                    install_target = str(local_wheels[0])
-                    self.log(f"  Using local wheel: {local_wheels[0].name}")
-
-        result = subprocess.run(
-            [str(uv), "pip", "install", "--python", str(python_exe),
-             "--upgrade", "--no-cache", install_target],
-            capture_output=True,
-            text=True,
-        )
-        if result.returncode != 0:
-            # Non-fatal - might already be installed or network issues
-            self.log(f"Warning: Failed to install comfy-env: {result.stderr}")
-        else:
-            self.log("comfy-env installed")
-
-    def _setup_windows_deps(self, env: IsolatedEnv, env_dir: Path) -> None:
-        """
-        Set up Windows-specific dependencies (VC++ runtime, DLL paths).
-
-        This:
-        1. Installs msvc-runtime package (provides VC++ DLLs)
-        2. Sets up the Lib/x64/vc17/bin directory structure for opencv
-        """
-        python_exe = self.get_python(env)
-        uv = self._find_uv()
-
-        # Install msvc-runtime package to get VC++ DLLs
-        self.log("Installing VC++ runtime (msvc-runtime)...")
-        result = subprocess.run(
-            [str(uv), "pip", "install", "--python", str(python_exe), "msvc-runtime"],
-            capture_output=True,
-            text=True,
-        )
-        if result.returncode != 0:
-            self.log(f"Warning: Failed to install msvc-runtime: {result.stderr}")
-        else:
-            self.log("msvc-runtime installed")
-
-        # Set up opencv DLL paths (copies DLLs to Lib/x64/vc17/bin)
-        self.log("Setting up opencv DLL paths...")
-        success, msg = self.platform.setup_opencv_dll_paths(env_dir)
-        if success:
-            if msg:
-                self.log(f"  {msg}")
-        else:
-            self.log(f"Warning: {msg}")
-
-    def ensure_system_deps(self) -> bool:
-        """
-        Ensure system-level dependencies are installed (Windows only).
-
-        On Windows, this checks for Media Foundation and installs if missing.
-        Returns True if all deps are available, False otherwise.
-        """
-        if self.platform.name != 'windows':
-            return True
-
-        # Check and install Media Foundation if needed
-        success, error = self.platform.ensure_media_foundation(self.log)
-        if not success:
-            self.log(f"WARNING: {error}")
-            self.log("Some packages (like opencv) may not work correctly.")
-            return False
-
-        return True
-
-    def setup(
-        self,
-        env: IsolatedEnv,
-        install_pytorch: bool = True,
-        verify_packages: Optional[list] = None,
-    ) -> Path:
-        """
-        Create environment and install all dependencies.
-
-        This is the main entry point for setting up an isolated environment.
-
-        Args:
-            env: Environment configuration
-            install_pytorch: Whether to install PyTorch
-            verify_packages: Packages to verify after installation
-
-        Returns:
-            Path to the environment directory
-        """
-        self.log("=" * 50)
-        self.log(f"Setting up isolated environment: {env.name}")
-        self.log("=" * 50)
-
-        # Ensure system-level deps (Media Foundation on Windows)
-        self.ensure_system_deps()
-
-        # Check if already ready
-        if self.is_ready(env, verify_packages):
-            self.log("Environment already ready, skipping setup")
-            return self.get_env_dir(env)
-
-        # Create virtual environment
-        env_dir = self.create_venv(env)
-
-        # Install PyTorch
-        if install_pytorch and (env.cuda is not None or detect_cuda_version()):
-            self.install_pytorch(env)
-
-        # Install comfy-env (needed for BaseWorker in workers)
-        self._install_comfy_env(env)
-
-        # Install other requirements
-        self.install_requirements(env)
-
-        # Windows: Install VC++ runtime and set up DLL paths
-        if self.platform.name == 'windows':
-            self._setup_windows_deps(env, env_dir)
-
-        # Verify installation
-        if verify_packages:
-            if self.is_ready(env, verify_packages):
-                self.log("Verification passed!")
-            else:
-                raise RuntimeError(f"Verification failed: could not import {verify_packages}")
-
-        self.log("=" * 50)
-        self.log("Setup complete!")
-        self.log(f"Python: {self.get_python(env)}")
-        self.log("=" * 50)
-
-        return env_dir
-
-    def delete(self, env: IsolatedEnv) -> bool:
-        """
-        Delete an environment.
-
-        Args:
-            env: Environment configuration
-
-        Returns:
-            True if deleted, False if didn't exist
-        """
-        env_dir = self.get_env_dir(env)
-
-        if not env_dir.exists():
-            return False
-
-        self.log(f"Deleting environment: {env_dir}")
-        self.platform.rmtree_robust(env_dir)
-        self.log("Environment deleted")
-        return True
-
-    def repair(
-        self,
-        env: IsolatedEnv,
-        install_pytorch: bool = True,
-        verify_packages: Optional[list] = None,
-    ) -> Path:
-        """
-        Delete and recreate an environment.
-
-        Args:
-            env: Environment configuration
-            install_pytorch: Whether to install PyTorch
-            verify_packages: Packages to verify after installation
-
-        Returns:
-            Path to the new environment directory
-        """
-        self.log("Repairing environment (delete + recreate)")
-        self.delete(env)
-        return self.setup(env, install_pytorch, verify_packages)