comfy-env 0.1.3__tar.gz → 0.1.5__tar.gz

{comfy_env-0.1.3 → comfy_env-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comfy-env
-Version: 0.1.3
+Version: 0.1.5
 Summary: Environment management for ComfyUI custom nodes - CUDA wheel resolution and process isolation
 Project-URL: Homepage, https://github.com/PozzettiAndrea/comfy-env
 Project-URL: Repository, https://github.com/PozzettiAndrea/comfy-env

{comfy_env-0.1.3 → comfy_env-0.1.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "comfy-env"
-version = "0.1.3"
+version = "0.1.5"
 description = "Environment management for ComfyUI custom nodes - CUDA wheel resolution and process isolation"
 readme = "README.md"
 license = {text = "MIT"}

{comfy_env-0.1.3 → comfy_env-0.1.5}/src/comfy_env/pixi/core.py

@@ -9,10 +9,12 @@ See: https://pixi.sh/
 
 import copy
 import platform
+import re
 import shutil
 import stat
 import subprocess
 import sys
+import urllib.request
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional
 
@@ -28,15 +30,9 @@ PIXI_URLS = {
     ("Windows", "AMD64"): "https://github.com/prefix-dev/pixi/releases/latest/download/pixi-x86_64-pc-windows-msvc.exe",
 }
 
-# CUDA wheels index
+# CUDA wheels index (includes flash-attn, PyG packages, and custom wheels)
 CUDA_WHEELS_INDEX = "https://pozzettiandrea.github.io/cuda-wheels/"
 
-# Flash attention pre-built wheels (mjun0812)
-FLASH_ATTN_INDEX = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/expanded_assets/v0.7.16"
-
-# PyTorch Geometric wheels index
-PYG_WHEELS_INDEX = "https://data.pyg.org/whl/"
-
 # CUDA version -> PyTorch version mapping
 CUDA_TORCH_MAP = {
     "12.8": "2.8",
@@ -44,19 +40,87 @@ CUDA_TORCH_MAP = {
     "12.1": "2.4",
 }
 
+def find_matching_wheel(package: str, torch_version: str, cuda_version: str) -> Optional[str]:
+    """
+    Query cuda-wheels index to find a wheel matching the CUDA/torch version.
+    Returns the full version spec (e.g., "flash-attn===2.8.3+cu128torch2.8") or None.
+    """
+    cuda_short = cuda_version.replace(".", "")[:3]  # "12.8" -> "128"
+    torch_short = torch_version.replace(".", "")[:2]  # "2.8" -> "28"
+
+    # Try different directory name variants
+    pkg_variants = [package, package.replace("-", "_"), package.replace("_", "-")]
+
+    for pkg_dir in pkg_variants:
+        url = f"{CUDA_WHEELS_INDEX}{pkg_dir}/"
+        try:
+            with urllib.request.urlopen(url, timeout=10) as resp:
+                html = resp.read().decode("utf-8")
+        except Exception:
+            continue
+
+        # Parse wheel filenames from href attributes
+        # Pattern: package_name-version+localversion-cpXX-cpXX-platform.whl
+        wheel_pattern = re.compile(
+            r'href="[^"]*?([^"/]+\.whl)"',
+            re.IGNORECASE
+        )
+
+        # Local version patterns to match:
+        # flash-attn style: +cu128torch2.8
+        # PyG style: +pt28cu128
+        local_patterns = [
+            f"+cu{cuda_short}torch{torch_version}",  # flash-attn style
+            f"+pt{torch_short}cu{cuda_short}",       # PyG style
+        ]
+
+        best_match = None
+        best_version = None
+
+        for match in wheel_pattern.finditer(html):
+            wheel_name = match.group(1)
+            # URL decode
+            wheel_name = wheel_name.replace("%2B", "+")
+
+            # Check if wheel matches our CUDA/torch version
+            for local_pattern in local_patterns:
+                if local_pattern in wheel_name:
+                    # Extract version from wheel name
+                    # Format: name-version+local-cpXX-cpXX-platform.whl
+                    parts = wheel_name.split("-")
+                    if len(parts) >= 2:
+                        version_part = parts[1]  # e.g., "2.8.3+cu128torch2.8"
+                        if best_version is None or version_part > best_version:
+                            best_version = version_part
+                            best_match = f"{package}==={version_part}"
+                    break
+
+        if best_match:
+            return best_match
+
+    return None
+
+
+def get_package_spec(package: str, torch_version: str, cuda_version: str) -> str:
+    """
+    Get package spec with local version for CUDA/torch compatibility.
+    Queries the index to find matching wheels dynamically.
+    """
+    spec = find_matching_wheel(package, torch_version, cuda_version)
+    return spec if spec else package
+
 
 def get_all_find_links(package: str, torch_version: str, cuda_version: str) -> list:
     """Get all find-links URLs for a CUDA package."""
-    cuda_short = cuda_version.replace(".", "")[:3]
-
-    return [
-        # cuda-wheels index (package-specific page)
-        f"{CUDA_WHEELS_INDEX}{package}/",
-        # flash-attn GitHub releases
-        FLASH_ATTN_INDEX,
-        # PyG index (torch version specific)
-        f"{PYG_WHEELS_INDEX}torch-{torch_version}.0+cu{cuda_short}.html",
-    ]
+    # Try both underscore and hyphen variants since directory naming is inconsistent
+    pkg_underscore = package.replace("-", "_")
+    pkg_hyphen = package.replace("_", "-")
+    urls = [f"{CUDA_WHEELS_INDEX}{package}/"]
+    if pkg_underscore != package:
+        urls.append(f"{CUDA_WHEELS_INDEX}{pkg_underscore}/")
+    if pkg_hyphen != package:
+        urls.append(f"{CUDA_WHEELS_INDEX}{pkg_hyphen}/")
+    return urls
 
 
 def get_current_platform() -> str:
@@ -340,6 +404,7 @@ def pixi_install(
         for package in cfg.cuda_packages:
             pip_cmd = [
                 str(python_path), "-m", "pip", "install",
+                "--no-index",  # Only use --find-links, don't query PyPI
                 "--no-deps",
                 "--no-cache-dir",
             ]
@@ -348,8 +413,10 @@ def pixi_install(
             for url in get_all_find_links(package, torch_version, cuda_version):
                 pip_cmd.extend(["--find-links", url])
 
-            log(f"  Installing {package}")
-            pip_cmd.append(package)
+            # Get package spec with local version for CUDA/torch compatibility
+            pkg_spec = get_package_spec(package, torch_version, cuda_version)
+            log(f"  Installing {pkg_spec}")
+            pip_cmd.append(pkg_spec)
             result = subprocess.run(pip_cmd, capture_output=True, text=True)
             if result.returncode != 0:
                 log(f"CUDA package install failed for {package}:\n{result.stderr}")

{comfy_env-0.1.3 → comfy_env-0.1.5}/src/comfy_env/workers/mp.py

@@ -40,7 +40,44 @@ _SHUTDOWN = object()
 _CALL_METHOD = "call_method"
 
 
-def _worker_loop(queue_in, queue_out, sys_path_additions=None, lib_path=None, env_vars=None):
+def _dump_worker_env(worker_name: str = "unknown"):
+    """Dump worker environment to .comfy-env/logs/ for debugging."""
+    import json
+    import os
+    import platform
+    import sys
+    from datetime import datetime
+    from pathlib import Path
+
+    log_dir = Path.cwd() / ".comfy-env" / "logs"
+    log_dir.mkdir(parents=True, exist_ok=True)
+
+    debug_info = {
+        "timestamp": datetime.now().isoformat(),
+        "worker_name": worker_name,
+        "pid": os.getpid(),
+        "cwd": os.getcwd(),
+        "python": {
+            "executable": sys.executable,
+            "version": sys.version,
+            "prefix": sys.prefix,
+        },
+        "platform": {
+            "system": platform.system(),
+            "machine": platform.machine(),
+            "release": platform.release(),
+        },
+        "env_vars": dict(os.environ),
+        "sys_path": sys.path,
+        "modules_loaded": sorted(sys.modules.keys()),
+    }
+
+    log_file = log_dir / f"worker_{worker_name}_{os.getpid()}.json"
+    log_file.write_text(json.dumps(debug_info, indent=2, default=str))
+    print(f"[comfy-env] Worker env dumped to: {log_file}")
+
+
+def _worker_loop(queue_in, queue_out, sys_path_additions=None, lib_path=None, env_vars=None, worker_name=None):
     """
     Worker process main loop.
 
@@ -57,6 +94,7 @@ def _worker_loop(queue_in, queue_out, sys_path_additions=None, lib_path=None, en
         sys_path_additions: Paths to add to sys.path
         lib_path: Path to add to LD_LIBRARY_PATH (for conda libraries)
         env_vars: Environment variables to set (from comfy-env.toml)
+        worker_name: Name of the worker (for logging)
     """
     import os
     import sys
@@ -69,6 +107,10 @@ def _worker_loop(queue_in, queue_out, sys_path_additions=None, lib_path=None, en
     # Set worker mode env var
     os.environ["COMFYUI_ISOLATION_WORKER"] = "1"
 
+    # Debug logging if enabled
+    if os.environ.get("COMFY_ENV_DEBUG", "").lower() in ("1", "true", "yes"):
+        _dump_worker_env(worker_name or "unknown")
+
     # DLL/library isolation - match SubprocessWorker's isolation level
     # Filter out conflicting paths from conda/mamba/etc and use proper DLL registration
     path_sep = ";" if sys.platform == "win32" else ":"
@@ -511,7 +553,7 @@ class MPWorker(Worker):
             self._queue_out = ctx.Queue()
             self._process = ctx.Process(
                 target=_worker_loop,
-                args=(self._queue_in, self._queue_out, self._sys_path, self._lib_path, self._env_vars),
+                args=(self._queue_in, self._queue_out, self._sys_path, self._lib_path, self._env_vars, self.name),
                 daemon=True,
             )
             self._process.start()