comfy-env 0.0.5__tar.gz → 0.0.6__tar.gz

{comfy_env-0.0.5 → comfy_env-0.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comfy-env
-Version: 0.0.5
+Version: 0.0.6
 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.0.5 → comfy_env-0.0.6}/pyproject.toml

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

{comfy_env-0.0.5 → comfy_env-0.0.6}/src/comfy_env/__init__.py

@@ -40,7 +40,7 @@ This package provides:
 The @isolated decorator and WorkerBridge are still available.
 """
 
-__version__ = "0.0.5"
+__version__ = "0.0.6"
 
 from .env.config import IsolatedEnv, EnvManagerConfig, LocalConfig, NodeReq
 from .env.config_file import (

{comfy_env-0.0.5 → comfy_env-0.0.6}/src/comfy_env/workers/venv.py

@@ -412,56 +412,60 @@ class VenvWorker(Worker):
         return f"<VenvWorker name={self.name!r} python={self.python}>"
 
 
-def _queue_worker_entrypoint(to_host_queue, from_host_queue, shm_dir, sys_paths):
-    """
-    Worker process entrypoint using queue-based IPC.
-
-    This function runs in the isolated venv subprocess and handles
-    requests from the host process.
+# Persistent worker script - runs as __main__ in the venv Python subprocess
+# Uses stdin/stdout JSON for IPC - avoids Windows multiprocessing spawn issues entirely
+_PERSISTENT_WORKER_SCRIPT = '''
+import sys
+import os
+import json
+import traceback
+from types import SimpleNamespace
 
-    Args:
-        to_host_queue: Queue for sending messages to host (worker -> host)
-        from_host_queue: Queue for receiving messages from host (host -> worker)
-        shm_dir: Shared memory directory for tensor files
-        sys_paths: List of paths to add to sys.path
-    """
-    import sys
-    import traceback
-    from pathlib import Path
-    from types import SimpleNamespace
+def _deserialize_isolated_objects(obj):
+    """Reconstruct objects serialized with __isolated_object__ marker."""
+    if isinstance(obj, dict):
+        if obj.get("__isolated_object__"):
+            attrs = {k: _deserialize_isolated_objects(v) for k, v in obj.get("__attrs__", {}).items()}
+            ns = SimpleNamespace(**attrs)
+            ns.__class_name__ = obj.get("__class_name__", "Unknown")
+            return ns
+        return {k: _deserialize_isolated_objects(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [_deserialize_isolated_objects(v) for v in obj]
+    elif isinstance(obj, tuple):
+        return tuple(_deserialize_isolated_objects(v) for v in obj)
+    return obj
 
-    # Setup paths
-    for p in sys_paths:
+def main():
+    # Read config from first line
+    config_line = sys.stdin.readline()
+    if not config_line:
+        return
+    config = json.loads(config_line)
+
+    # Setup sys.path
+    for p in config.get("sys_paths", []):
         if p not in sys.path:
             sys.path.insert(0, p)
 
-    def _deserialize_isolated_objects(obj):
-        """Reconstruct objects serialized with __isolated_object__ marker."""
-        if isinstance(obj, dict):
-            if obj.get("__isolated_object__"):
-                attrs = {k: _deserialize_isolated_objects(v) for k, v in obj.get("__attrs__", {}).items()}
-                ns = SimpleNamespace(**attrs)
-                ns.__class_name__ = obj.get("__class_name__", "Unknown")
-                return ns
-            return {k: _deserialize_isolated_objects(v) for k, v in obj.items()}
-        elif isinstance(obj, list):
-            return [_deserialize_isolated_objects(v) for v in obj]
-        elif isinstance(obj, tuple):
-            return tuple(_deserialize_isolated_objects(v) for v in obj)
-        return obj
+    # Import torch after path setup
+    import torch
 
     # Signal ready
-    to_host_queue.put({"status": "ready"})
-
-    import torch
+    sys.stdout.write(json.dumps({"status": "ready"}) + "\\n")
+    sys.stdout.flush()
 
+    # Process requests
     while True:
         try:
-            request = from_host_queue.get()
+            line = sys.stdin.readline()
+            if not line:
+                break
+            request = json.loads(line)
         except Exception:
             break
 
-        if request is None or request.get("method") == "shutdown":
+        if request.get("method") == "shutdown":
             break
 
         try:
@@ -481,7 +485,6 @@ def _queue_worker_entrypoint(to_host_queue, from_host_queue, shm_dir, sys_paths)
             module = __import__(module_name, fromlist=[""])
 
             if request_type == "call_method":
-                # Call a method on a class instance
                 class_name = request["class_name"]
                 method_name = request["method_name"]
                 self_state = request.get("self_state")
@@ -493,7 +496,6 @@ def _queue_worker_entrypoint(to_host_queue, from_host_queue, shm_dir, sys_paths)
                 method = getattr(instance, method_name)
                 result = method(**inputs)
             else:
-                # Call a module-level function
                 func_name = request["func"]
                 func = getattr(module, func_name)
                 result = func(**inputs)
@@ -502,25 +504,37 @@ def _queue_worker_entrypoint(to_host_queue, from_host_queue, shm_dir, sys_paths)
             if outputs_path:
                 torch.save(result, outputs_path)
 
-            to_host_queue.put({"status": "ok"})
+            sys.stdout.write(json.dumps({"status": "ok"}) + "\\n")
+            sys.stdout.flush()
 
         except Exception as e:
-            to_host_queue.put({
+            sys.stdout.write(json.dumps({
                 "status": "error",
                 "error": str(e),
                 "traceback": traceback.format_exc(),
-            })
+            }) + "\\n")
+            sys.stdout.flush()
+
+if __name__ == "__main__":
+    main()
+'''
 
 
 class PersistentVenvWorker(Worker):
     """
     Persistent version of VenvWorker that keeps subprocess alive.
 
-    This reduces per-call overhead by ~200-400ms by avoiding subprocess spawn.
-    Uses multiprocessing.Queue for cross-platform communication with zero-copy
-    tensor transfer support.
+    Uses subprocess.Popen with stdin/stdout JSON IPC instead of multiprocessing.
+    This avoids Windows multiprocessing spawn issues where the child process
+    tries to reimport __main__ (which fails when using a different Python).
+
+    Benefits:
+    - Works on Windows with different venv Python (full isolation)
+    - Compiled CUDA extensions load correctly in the venv
+    - ~50-100ms per call (vs ~300-500ms for VenvWorker per-call spawn)
+    - Tensor transfer via shared memory files
 
-    Use this for high-frequency calls to the same venv.
+    Use this for high-frequency calls to isolated venvs.
     """
 
     def __init__(
@@ -530,7 +544,7 @@ class PersistentVenvWorker(Worker):
         sys_path: Optional[List[str]] = None,
         env: Optional[Dict[str, str]] = None,
         name: Optional[str] = None,
-        share_torch: bool = True,
+        share_torch: bool = True,  # Kept for API compatibility
     ):
         """
         Initialize persistent worker.
@@ -541,27 +555,26 @@ class PersistentVenvWorker(Worker):
             sys_path: Additional paths to add to sys.path.
             env: Additional environment variables.
             name: Optional name for logging.
-            share_torch: Use torch.multiprocessing for zero-copy tensor transfer.
+            share_torch: Ignored (kept for API compatibility).
         """
         self.python = Path(python)
         self.working_dir = Path(working_dir) if working_dir else Path.cwd()
         self.sys_path = sys_path or []
         self.extra_env = env or {}
         self.name = name or f"PersistentVenvWorker({self.python.parent.parent.name})"
-        self.share_torch = share_torch
 
         if not self.python.exists():
             raise FileNotFoundError(f"Python not found: {self.python}")
 
         self._temp_dir = Path(tempfile.mkdtemp(prefix='comfyui_pvenv_'))
         self._shm_dir = _get_shm_dir()
-        self._process = None  # Will be multiprocessing.Process
-        self._to_worker = None  # Queue: host -> worker
-        self._from_worker = None  # Queue: worker -> host
-        self._transport = None
+        self._process: Optional[subprocess.Popen] = None
         self._shutdown = False
         self._lock = threading.Lock()
-        self._mp = None  # multiprocessing module reference
+
+        # Write worker script to temp file
+        self._worker_script = self._temp_dir / "persistent_worker.py"
+        self._worker_script.write_text(_PERSISTENT_WORKER_SCRIPT)
 
     def _find_comfyui_base(self) -> Optional[Path]:
         """Find ComfyUI base directory by walking up from working_dir."""
@@ -573,20 +586,13 @@ class PersistentVenvWorker(Worker):
         return None
 
     def _ensure_started(self):
-        """Start persistent worker process if not running."""
+        """Start persistent worker subprocess if not running."""
         if self._shutdown:
             raise RuntimeError(f"{self.name}: Worker has been shut down")
 
-        if self._process is not None and self._process.is_alive():
+        if self._process is not None and self._process.poll() is None:
             return  # Already running
 
-        from ..ipc.transport import create_queue_pair, QueueTransport
-
-        # Create queue pair for bidirectional communication
-        self._to_worker, self._from_worker, self._mp = create_queue_pair(
-            share_torch=self.share_torch
-        )
-
         # Set up environment
         env = os.environ.copy()
         env.update(self.extra_env)
@@ -601,35 +607,57 @@ class PersistentVenvWorker(Worker):
         stubs_dir = Path(__file__).parent.parent / "stubs"
         all_sys_path = [str(stubs_dir), str(self.working_dir)] + self.sys_path
 
-        # On Windows, set VIRTUAL_ENV for proper venv activation
-        if os.name == "nt":
-            env["VIRTUAL_ENV"] = str(self.python.parent.parent)
-
-        # Set the Python executable for multiprocessing
-        self._mp.set_executable(str(self.python))
-
-        # Launch worker process
-        self._process = self._mp.Process(
-            target=_queue_worker_entrypoint,
-            args=(
-                self._to_worker,
-                self._from_worker,
-                str(self._shm_dir),
-                all_sys_path,
-            ),
+        # Launch subprocess with the venv Python
+        # This runs _PERSISTENT_WORKER_SCRIPT as __main__ - no reimport issues!
+        self._process = subprocess.Popen(
+            [str(self.python), str(self._worker_script)],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            cwd=str(self.working_dir),
+            env=env,
+            bufsize=1,  # Line buffered
+            text=True,  # Text mode for JSON
         )
-        self._process.start()
 
-        # Create transport wrapper
-        self._transport = QueueTransport(self._to_worker, self._from_worker)
+        # Send config
+        config = {"sys_paths": all_sys_path}
+        self._process.stdin.write(json.dumps(config) + "\n")
+        self._process.stdin.flush()
+
+        # Wait for ready signal with timeout
+        import select
+        if sys.platform == "win32":
+            # Windows: can't use select on pipes, use thread with timeout
+            ready_line = [None]
+            def read_ready():
+                try:
+                    ready_line[0] = self._process.stdout.readline()
+                except:
+                    pass
+            t = threading.Thread(target=read_ready, daemon=True)
+            t.start()
+            t.join(timeout=60)
+            line = ready_line[0]
+        else:
+            # Unix: use select for timeout
+            import select
+            ready, _, _ = select.select([self._process.stdout], [], [], 60)
+            line = self._process.stdout.readline() if ready else None
+
+        if not line:
+            stderr = ""
+            try:
+                self._process.kill()
+                _, stderr = self._process.communicate(timeout=5)
+            except:
+                pass
+            raise RuntimeError(f"{self.name}: Worker failed to start (timeout). stderr: {stderr}")
 
-        # Wait for ready signal
         try:
-            msg = self._transport.recv(timeout=30)
-        except Exception:
-            if self._process.is_alive():
-                self._process.terminate()
-            raise RuntimeError(f"{self.name}: Worker failed to start (timeout)")
+            msg = json.loads(line)
+        except json.JSONDecodeError as e:
+            raise RuntimeError(f"{self.name}: Invalid ready message: {line!r}") from e
 
         if msg.get("status") != "ready":
             raise RuntimeError(f"{self.name}: Unexpected ready message: {msg}")
@@ -646,6 +674,45 @@ class PersistentVenvWorker(Worker):
             f"{self.name}: Use call_module(module='...', func='...') instead."
         )
 
+    def _send_request(self, request: dict, timeout: float) -> dict:
+        """Send request via stdin and read response from stdout with timeout."""
+        # Send request
+        self._process.stdin.write(json.dumps(request) + "\n")
+        self._process.stdin.flush()
+
+        # Read response with timeout
+        if sys.platform == "win32":
+            # Windows: use thread for timeout
+            response_line = [None]
+            def read_response():
+                try:
+                    response_line[0] = self._process.stdout.readline()
+                except:
+                    pass
+            t = threading.Thread(target=read_response, daemon=True)
+            t.start()
+            t.join(timeout=timeout)
+            line = response_line[0]
+        else:
+            # Unix: use select
+            import select
+            ready, _, _ = select.select([self._process.stdout], [], [], timeout)
+            line = self._process.stdout.readline() if ready else None
+
+        if not line:
+            # Timeout - kill process
+            try:
+                self._process.kill()
+            except:
+                pass
+            self._shutdown = True
+            raise TimeoutError(f"{self.name}: Call timed out after {timeout}s")
+
+        try:
+            return json.loads(line)
+        except json.JSONDecodeError as e:
+            raise WorkerError(f"Invalid response from worker: {line!r}") from e
+
     def call_method(
         self,
         module_name: str,
@@ -669,8 +736,6 @@ class PersistentVenvWorker(Worker):
         Returns:
             Return value of the method.
         """
-        import queue as queue_module
-
         with self._lock:
             self._ensure_started()
 
@@ -697,15 +762,7 @@ class PersistentVenvWorker(Worker):
                     "inputs_path": str(inputs_path) if kwargs else None,
                     "outputs_path": str(outputs_path),
                 }
-                self._transport.send(request)
-
-                # Wait for response with timeout
-                try:
-                    response = self._transport.recv(timeout=timeout)
-                except queue_module.Empty:
-                    self._process.terminate()
-                    self._shutdown = True
-                    raise TimeoutError(f"{self.name}: Call timed out")
+                response = self._send_request(request, timeout)
 
                 if response.get("status") == "error":
                     raise WorkerError(
@@ -732,8 +789,6 @@ class PersistentVenvWorker(Worker):
         **kwargs
     ) -> Any:
         """Call a function by module path."""
-        import queue as queue_module
-
         with self._lock:
             self._ensure_started()
 
@@ -758,15 +813,7 @@ class PersistentVenvWorker(Worker):
                     "inputs_path": str(inputs_path) if kwargs else None,
                     "outputs_path": str(outputs_path),
                 }
-                self._transport.send(request)
-
-                # Wait for response with timeout
-                try:
-                    response = self._transport.recv(timeout=timeout)
-                except queue_module.Empty:
-                    self._process.terminate()
-                    self._shutdown = True
-                    raise TimeoutError(f"{self.name}: Call timed out")
+                response = self._send_request(request, timeout)
 
                 if response.get("status") == "error":
                     raise WorkerError(
@@ -792,32 +839,21 @@ class PersistentVenvWorker(Worker):
             return
         self._shutdown = True
 
-        # Send shutdown signal via queue
-        if self._transport:
+        # Send shutdown signal via stdin
+        if self._process and self._process.poll() is None:
             try:
-                self._transport.send({"method": "shutdown"})
+                self._process.stdin.write(json.dumps({"method": "shutdown"}) + "\n")
+                self._process.stdin.flush()
+                self._process.stdin.close()
             except:
                 pass
-            self._transport.close()
-
-        # Wait for process to exit
-        if self._process:
-            self._process.join(timeout=5)
-            if self._process.is_alive():
-                self._process.terminate()
-                self._process.join(timeout=2)
 
-        # Clean up queues
-        if self._to_worker:
-            try:
-                self._to_worker.close()
-            except:
-                pass
-        if self._from_worker:
+            # Wait for process to exit
             try:
-                self._from_worker.close()
-            except:
-                pass
+                self._process.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                self._process.kill()
+                self._process.wait(timeout=2)
 
         shutil.rmtree(self._temp_dir, ignore_errors=True)
 
@@ -826,7 +862,7 @@ class PersistentVenvWorker(Worker):
             return False
         if self._process is None:
             return False
-        return self._process.is_alive()
+        return self._process.poll() is None
 
     def __repr__(self):
         status = "alive" if self.is_alive() else "stopped"