comfy-env 0.1.19__tar.gz → 0.1.21__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (38) hide show
  1. {comfy_env-0.1.19 → comfy_env-0.1.21}/PKG-INFO +1 -1
  2. {comfy_env-0.1.19 → comfy_env-0.1.21}/pyproject.toml +1 -1
  3. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/__init__.py +0 -2
  4. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/isolation/__init__.py +0 -2
  5. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/isolation/workers/__init__.py +1 -3
  6. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/isolation/workers/base.py +1 -1
  7. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/isolation/wrap.py +8 -12
  8. comfy_env-0.1.19/src/comfy_env/isolation/workers/mp.py +0 -864
  9. {comfy_env-0.1.19 → comfy_env-0.1.21}/.github/workflows/ci.yml +0 -0
  10. {comfy_env-0.1.19 → comfy_env-0.1.21}/.github/workflows/publish.yml +0 -0
  11. {comfy_env-0.1.19 → comfy_env-0.1.21}/.gitignore +0 -0
  12. {comfy_env-0.1.19 → comfy_env-0.1.21}/LICENSE +0 -0
  13. {comfy_env-0.1.19 → comfy_env-0.1.21}/README.md +0 -0
  14. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/cli.py +0 -0
  15. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/config/__init__.py +0 -0
  16. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/config/parser.py +0 -0
  17. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/config/types.py +0 -0
  18. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/detection/__init__.py +0 -0
  19. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/detection/cuda.py +0 -0
  20. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/detection/gpu.py +0 -0
  21. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/detection/platform.py +0 -0
  22. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/detection/runtime.py +0 -0
  23. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/environment/__init__.py +0 -0
  24. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/environment/cache.py +0 -0
  25. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/environment/libomp.py +0 -0
  26. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/environment/paths.py +0 -0
  27. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/environment/setup.py +0 -0
  28. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/install.py +0 -0
  29. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/isolation/tensor_utils.py +0 -0
  30. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/isolation/workers/subprocess.py +0 -0
  31. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/packages/__init__.py +0 -0
  32. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/packages/apt.py +0 -0
  33. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/packages/cuda_wheels.py +0 -0
  34. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/packages/node_dependencies.py +0 -0
  35. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/packages/pixi.py +0 -0
  36. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/packages/toml_generator.py +0 -0
  37. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/templates/comfy-env-instructions.txt +0 -0
  38. {comfy_env-0.1.19 → comfy_env-0.1.21}/src/comfy_env/templates/comfy-env.toml +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: comfy-env
3
- Version: 0.1.19
3
+ Version: 0.1.21
4
4
  Summary: Environment management for ComfyUI custom nodes - CUDA wheel resolution and process isolation
5
5
  Project-URL: Homepage, https://github.com/PozzettiAndrea/comfy-env
6
6
  Project-URL: Repository, https://github.com/PozzettiAndrea/comfy-env
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "comfy-env"
3
- version = "0.1.19"
3
+ version = "0.1.21"
4
4
  description = "Environment management for ComfyUI custom nodes - CUDA wheel resolution and process isolation"
5
5
  readme = "README.md"
6
6
  license = {text = "MIT"}
@@ -108,7 +108,6 @@ from .isolation import (
108
108
  # Workers
109
109
  Worker,
110
110
  WorkerError,
111
- MPWorker,
112
111
  SubprocessWorker,
113
112
  # Tensor utilities
114
113
  TensorKeeper,
@@ -168,7 +167,6 @@ __all__ = [
168
167
  # Workers
169
168
  "Worker",
170
169
  "WorkerError",
171
- "MPWorker",
172
170
  "SubprocessWorker",
173
171
  "TensorKeeper",
174
172
  ]
@@ -11,7 +11,6 @@ from .wrap import (
11
11
  from .workers import (
12
12
  Worker,
13
13
  WorkerError,
14
- MPWorker,
15
14
  SubprocessWorker,
16
15
  )
17
16
  from .tensor_utils import (
@@ -29,7 +28,6 @@ __all__ = [
29
28
  # Workers
30
29
  "Worker",
31
30
  "WorkerError",
32
- "MPWorker",
33
31
  "SubprocessWorker",
34
32
  # Tensor utilities
35
33
  "TensorKeeper",
@@ -1,16 +1,14 @@
1
1
  """
2
2
  Workers - Process isolation implementations.
3
3
 
4
- Provides multiprocessing and subprocess-based workers for isolated execution.
4
+ Provides subprocess-based workers for isolated execution.
5
5
  """
6
6
 
7
7
  from .base import Worker, WorkerError
8
- from .mp import MPWorker
9
8
  from .subprocess import SubprocessWorker
10
9
 
11
10
  __all__ = [
12
11
  "Worker",
13
12
  "WorkerError",
14
- "MPWorker",
15
13
  "SubprocessWorker",
16
14
  ]
@@ -16,7 +16,7 @@ class Worker(ABC):
16
16
 
17
17
  Workers should be used as context managers when possible:
18
18
 
19
- with MPWorker() as worker:
19
+ with SubprocessWorker(python="/path/to/venv/bin/python") as worker:
20
20
  result = worker.call(my_func, arg1, arg2)
21
21
  """
22
22
 
@@ -72,18 +72,14 @@ def _get_worker(env_dir: Path, working_dir: Path, sys_path: list[str],
72
72
  if cache_key in _workers and _workers[cache_key].is_alive():
73
73
  return _workers[cache_key]
74
74
 
75
- host_ver = f"{sys.version_info.major}.{sys.version_info.minor}"
76
- iso_ver = _get_python_version(env_dir)
77
-
78
- if iso_ver and iso_ver != host_ver:
79
- from .workers.subprocess import SubprocessWorker
80
- python = env_dir / ("python.exe" if sys.platform == "win32" else "bin/python")
81
- print(f"[comfy-env] SubprocessWorker: {python} ({iso_ver} vs {host_ver})")
82
- worker = SubprocessWorker(python=str(python), working_dir=working_dir, sys_path=sys_path, name=working_dir.name)
83
- else:
84
- from .workers.mp import MPWorker
85
- print(f"[comfy-env] MPWorker: {env_dir}")
86
- worker = MPWorker(name=working_dir.name, sys_path=sys_path, lib_path=lib_path, env_vars=env_vars)
75
+ python = env_dir / ("python.exe" if sys.platform == "win32" else "bin/python")
76
+
77
+ # Always use SubprocessWorker - MPWorker's spawn mechanism tries to re-import
78
+ # the parent's __main__ (ComfyUI's main.py), which fails with import errors.
79
+ # SubprocessWorker uses a clean entry script that avoids this issue.
80
+ from .workers.subprocess import SubprocessWorker
81
+ print(f"[comfy-env] SubprocessWorker: {python}")
82
+ worker = SubprocessWorker(python=str(python), working_dir=working_dir, sys_path=sys_path, name=working_dir.name)
87
83
 
88
84
  _workers[cache_key] = worker
89
85
  return worker
@@ -1,864 +0,0 @@
1
- """
2
- MPWorker - Same-venv isolation using multiprocessing.
3
-
4
- This is the simplest and fastest worker type:
5
- - Uses multiprocessing.Queue for IPC
6
- - Zero-copy tensor transfer via shared memory (automatic)
7
- - Fresh CUDA context in subprocess
8
- - ~30ms overhead per call
9
-
10
- Use this when you need:
11
- - Memory isolation between nodes
12
- - Fresh CUDA context (automatic VRAM cleanup on worker death)
13
- - Same Python environment as host
14
-
15
- Example:
16
- worker = MPWorker()
17
-
18
- def gpu_work(image):
19
- import torch
20
- return image * 2
21
-
22
- result = worker.call(gpu_work, image=my_tensor)
23
- worker.shutdown()
24
- """
25
-
26
- import logging
27
- import traceback
28
- from queue import Empty as QueueEmpty
29
- from typing import Any, Callable, Optional
30
-
31
- from .base import Worker, WorkerError
32
- from ..tensor_utils import prepare_for_ipc_recursive, keep_tensors_recursive
33
-
34
- logger = logging.getLogger("comfy_env")
35
-
36
-
37
- # Sentinel value for shutdown
38
- _SHUTDOWN = object()
39
-
40
- # Message type for method calls (avoids pickling issues with functions)
41
- _CALL_METHOD = "call_method"
42
-
43
-
44
- def _can_use_cuda_ipc():
45
- """
46
- Check if CUDA IPC is available.
47
-
48
- CUDA IPC works with native allocator but breaks with cudaMallocAsync.
49
- If no backend is specified, CUDA IPC should work (PyTorch default is native).
50
- """
51
- import os
52
- conf = os.environ.get('PYTORCH_CUDA_ALLOC_CONF', '')
53
- return 'cudaMallocAsync' not in conf
54
-
55
-
56
- # ---------------------------------------------------------------------------
57
- # Tensor file transfer - fallback for cudaMallocAsync (CUDA IPC doesn't work)
58
- # ---------------------------------------------------------------------------
59
-
60
- def _save_tensors_to_files(obj, file_registry=None):
61
- """Recursively save torch tensors to temp files for IPC."""
62
- if file_registry is None:
63
- file_registry = []
64
-
65
- try:
66
- import torch
67
- if isinstance(obj, torch.Tensor):
68
- import tempfile
69
- f = tempfile.NamedTemporaryFile(suffix='.pt', delete=False)
70
- torch.save(obj.cpu(), f.name) # Always save as CPU tensor
71
- f.close()
72
- file_registry.append(f.name)
73
- return {"__tensor_file__": f.name, "dtype": str(obj.dtype), "device": str(obj.device)}
74
- except ImportError:
75
- pass
76
-
77
- if isinstance(obj, dict):
78
- return {k: _save_tensors_to_files(v, file_registry) for k, v in obj.items()}
79
- elif isinstance(obj, list):
80
- return [_save_tensors_to_files(v, file_registry) for v in obj]
81
- elif isinstance(obj, tuple):
82
- return tuple(_save_tensors_to_files(v, file_registry) for v in obj)
83
- return obj
84
-
85
-
86
- def _load_tensors_from_files(obj):
87
- """Recursively load torch tensors from temp files."""
88
- if isinstance(obj, dict):
89
- if "__tensor_file__" in obj:
90
- import os
91
- import torch
92
- tensor = torch.load(obj["__tensor_file__"], weights_only=True)
93
- os.unlink(obj["__tensor_file__"]) # Cleanup temp file
94
- return tensor
95
- return {k: _load_tensors_from_files(v) for k, v in obj.items()}
96
- elif isinstance(obj, list):
97
- return [_load_tensors_from_files(v) for v in obj]
98
- elif isinstance(obj, tuple):
99
- return tuple(_load_tensors_from_files(v) for v in obj)
100
- return obj
101
-
102
-
103
- def _dump_worker_env(worker_name: str = "unknown", print_to_terminal: bool = False):
104
- """Dump worker environment to .comfy-env/logs/ (always) and optionally print."""
105
- import json
106
- import os
107
- import platform
108
- import sys
109
- from datetime import datetime
110
- from pathlib import Path
111
-
112
- log_dir = Path.cwd() / ".comfy-env" / "logs"
113
- log_dir.mkdir(parents=True, exist_ok=True)
114
-
115
- debug_info = {
116
- "timestamp": datetime.now().isoformat(),
117
- "worker_name": worker_name,
118
- "pid": os.getpid(),
119
- "cwd": os.getcwd(),
120
- "python": {
121
- "executable": sys.executable,
122
- "version": sys.version,
123
- "prefix": sys.prefix,
124
- },
125
- "platform": {
126
- "system": platform.system(),
127
- "machine": platform.machine(),
128
- "release": platform.release(),
129
- },
130
- "env_vars": dict(os.environ),
131
- "sys_path": sys.path,
132
- "modules_loaded": sorted(sys.modules.keys()),
133
- }
134
-
135
- log_file = log_dir / f"worker_{worker_name}_{os.getpid()}.json"
136
- log_file.write_text(json.dumps(debug_info, indent=2, default=str))
137
-
138
- if print_to_terminal:
139
- print(f"[comfy-env] === WORKER ENV DEBUG: {worker_name} ===")
140
- print(f"[comfy-env] Python: {sys.executable}")
141
- print(f"[comfy-env] Version: {sys.version.split()[0]}")
142
- print(f"[comfy-env] PID: {os.getpid()}, CWD: {os.getcwd()}")
143
- for var in ['PATH', 'LD_LIBRARY_PATH', 'DYLD_LIBRARY_PATH', 'PYTHONPATH', 'OMP_NUM_THREADS', 'KMP_DUPLICATE_LIB_OK']:
144
- val = os.environ.get(var, '<unset>')
145
- if len(val) > 100:
146
- val = val[:100] + '...'
147
- print(f"[comfy-env] {var}={val}")
148
- print(f"[comfy-env] Env dumped to: {log_file}")
149
-
150
-
151
- def _worker_loop(queue_in, queue_out, sys_path_additions=None, lib_path=None, env_vars=None, worker_name=None):
152
- """
153
- Worker process main loop.
154
-
155
- Receives work items and executes them:
156
- - ("call_method", module_name, class_name, method_name, self_state, kwargs): Call a method on a class
157
- - (func, args, kwargs): Execute a function directly
158
- - _SHUTDOWN: Shutdown the worker
159
-
160
- Runs until receiving _SHUTDOWN sentinel.
161
-
162
- Args:
163
- queue_in: Input queue for receiving work items
164
- queue_out: Output queue for sending results
165
- sys_path_additions: Paths to add to sys.path
166
- lib_path: Path to add to LD_LIBRARY_PATH (for conda libraries)
167
- env_vars: Environment variables to set (from comfy-env.toml)
168
- worker_name: Name of the worker (for logging)
169
- """
170
- import os
171
- import sys
172
- from pathlib import Path
173
-
174
- # Apply env_vars FIRST (before any library imports that might check them)
175
- if env_vars:
176
- os.environ.update(env_vars)
177
-
178
- # Set worker mode env var
179
- os.environ["COMFYUI_ISOLATION_WORKER"] = "1"
180
-
181
- # Always dump env to file, print to terminal if debug enabled
182
- print_debug = os.environ.get("COMFY_ENV_DEBUG", "").lower() in ("1", "true", "yes")
183
- _dump_worker_env(worker_name or "unknown", print_to_terminal=print_debug)
184
-
185
- # DLL/library isolation - match SubprocessWorker's isolation level
186
- # Filter out conflicting paths from conda/mamba/etc and use proper DLL registration
187
- path_sep = ";" if sys.platform == "win32" else ":"
188
-
189
- if sys.platform == "win32":
190
- # Use os.add_dll_directory() for explicit DLL registration (Python 3.8+)
191
- if lib_path and hasattr(os, "add_dll_directory"):
192
- try:
193
- os.add_dll_directory(lib_path)
194
- except Exception:
195
- pass
196
-
197
- # Filter conflicting paths from PATH (matches subprocess.py:1203-1212)
198
- current_path = os.environ.get("PATH", "")
199
- clean_parts = [
200
- p for p in current_path.split(path_sep)
201
- if not any(x in p.lower() for x in (".ct-envs", "conda", "mamba", "miniforge", "miniconda", "anaconda", "mingw"))
202
- ]
203
- if lib_path:
204
- clean_parts.insert(0, lib_path)
205
- os.environ["PATH"] = path_sep.join(clean_parts)
206
- elif sys.platform == "darwin":
207
- # macOS: ONLY use the isolated lib_path, don't inherit
208
- if lib_path:
209
- os.environ["DYLD_LIBRARY_PATH"] = lib_path
210
- else:
211
- os.environ.pop("DYLD_LIBRARY_PATH", None)
212
- else:
213
- # Linux: Use LD_LIBRARY_PATH
214
- current = os.environ.get("LD_LIBRARY_PATH", "")
215
- clean_parts = [
216
- p for p in current.split(path_sep) if p
217
- and not any(x in p.lower() for x in (".ct-envs", "conda", "mamba", "miniforge", "miniconda", "anaconda"))
218
- ]
219
- if lib_path:
220
- clean_parts.insert(0, lib_path)
221
- os.environ["LD_LIBRARY_PATH"] = path_sep.join(clean_parts)
222
-
223
- # Find ComfyUI base and add to sys.path for real folder_paths/comfy modules
224
- # This works because comfy.options.args_parsing=False by default, so folder_paths
225
- # auto-detects its base directory from __file__ location
226
- def _find_comfyui_base():
227
- cwd = Path.cwd().resolve()
228
- # Check common child directories (for test environments)
229
- for base in [cwd, cwd.parent]:
230
- for child in [".comfy-test-env/ComfyUI", "ComfyUI"]:
231
- candidate = base / child
232
- if (candidate / "main.py").exists() and (candidate / "comfy").exists():
233
- return candidate
234
- # Walk up from cwd looking for ComfyUI
235
- current = cwd
236
- for _ in range(10):
237
- if (current / "main.py").exists() and (current / "comfy").exists():
238
- return current
239
- current = current.parent
240
- # Check COMFYUI_BASE env var as fallback
241
- if os.environ.get("COMFYUI_BASE"):
242
- return Path(os.environ["COMFYUI_BASE"])
243
- return None
244
-
245
- comfyui_base = _find_comfyui_base()
246
- if comfyui_base and str(comfyui_base) not in sys.path:
247
- sys.path.insert(0, str(comfyui_base))
248
-
249
- # Add custom paths to sys.path for module discovery
250
- if sys_path_additions:
251
- for path in sys_path_additions:
252
- if path not in sys.path:
253
- sys.path.insert(0, path)
254
-
255
- while True:
256
- try:
257
- item = queue_in.get()
258
-
259
- # Check for shutdown signal
260
- if item is _SHUTDOWN:
261
- queue_out.put(("shutdown", None))
262
- break
263
-
264
- try:
265
- # Handle method call protocol
266
- if isinstance(item, tuple) and len(item) == 6 and item[0] == _CALL_METHOD:
267
- _, module_name, class_name, method_name, self_state, kwargs = item
268
- # Load tensors from files if using file-based transfer
269
- if not _can_use_cuda_ipc():
270
- kwargs = _load_tensors_from_files(kwargs)
271
- result = _execute_method_call(
272
- module_name, class_name, method_name, self_state, kwargs
273
- )
274
- # Handle result based on allocator
275
- if _can_use_cuda_ipc():
276
- keep_tensors_recursive(result)
277
- else:
278
- result = _save_tensors_to_files(result)
279
- queue_out.put(("ok", result))
280
- else:
281
- # Direct function call (legacy)
282
- func, args, kwargs = item
283
- # Load tensors from files if using file-based transfer
284
- if not _can_use_cuda_ipc():
285
- args = tuple(_load_tensors_from_files(a) for a in args)
286
- kwargs = _load_tensors_from_files(kwargs)
287
- result = func(*args, **kwargs)
288
- # Handle result based on allocator
289
- if _can_use_cuda_ipc():
290
- keep_tensors_recursive(result)
291
- else:
292
- result = _save_tensors_to_files(result)
293
- queue_out.put(("ok", result))
294
-
295
- except Exception as e:
296
- tb = traceback.format_exc()
297
- queue_out.put(("error", (str(e), tb)))
298
-
299
- except Exception as e:
300
- # Queue error - try to report, then exit
301
- try:
302
- queue_out.put(("fatal", str(e)))
303
- except:
304
- pass
305
- break
306
-
307
-
308
- class PathBasedModuleFinder:
309
- """
310
- Meta path finder that handles ComfyUI's path-based module names.
311
-
312
- ComfyUI uses full filesystem paths as module names for custom nodes.
313
- This finder intercepts imports of such modules and loads them from disk.
314
- """
315
-
316
- def find_spec(self, fullname, path, target=None):
317
- import importlib.util
318
- import os
319
-
320
- # Only handle path-based module names (starting with /)
321
- if not fullname.startswith('/'):
322
- return None
323
-
324
- # Parse the module name to find base path and submodule parts
325
- parts = fullname.split('.')
326
- base_path = parts[0]
327
- submodule_parts = parts[1:] if len(parts) > 1 else []
328
-
329
- # Walk through parts to find where path ends and module begins
330
- for i, part in enumerate(submodule_parts):
331
- test_path = os.path.join(base_path, part)
332
- if os.path.exists(test_path):
333
- base_path = test_path
334
- else:
335
- # Remaining parts are module names
336
- submodule_parts = submodule_parts[i:]
337
- break
338
- else:
339
- # All parts were path components
340
- submodule_parts = []
341
-
342
- # Determine the file to load
343
- if submodule_parts:
344
- # We're importing a submodule
345
- current_path = base_path
346
- for part in submodule_parts[:-1]:
347
- current_path = os.path.join(current_path, part)
348
-
349
- submod = submodule_parts[-1]
350
- submod_file = os.path.join(current_path, submod + '.py')
351
- submod_pkg = os.path.join(current_path, submod, '__init__.py')
352
-
353
- if os.path.exists(submod_file):
354
- return importlib.util.spec_from_file_location(fullname, submod_file)
355
- elif os.path.exists(submod_pkg):
356
- return importlib.util.spec_from_file_location(
357
- fullname, submod_pkg,
358
- submodule_search_locations=[os.path.join(current_path, submod)]
359
- )
360
- else:
361
- # Top-level path-based module
362
- if os.path.isdir(base_path):
363
- init_path = os.path.join(base_path, "__init__.py")
364
- if os.path.exists(init_path):
365
- return importlib.util.spec_from_file_location(
366
- fullname, init_path,
367
- submodule_search_locations=[base_path]
368
- )
369
- elif os.path.isfile(base_path):
370
- return importlib.util.spec_from_file_location(fullname, base_path)
371
-
372
- return None
373
-
374
-
375
- # Global flag to track if we've installed the finder
376
- _path_finder_installed = False
377
-
378
-
379
- def _ensure_path_finder_installed():
380
- """Install the PathBasedModuleFinder if not already installed."""
381
- import sys
382
- global _path_finder_installed
383
- if not _path_finder_installed:
384
- sys.meta_path.insert(0, PathBasedModuleFinder())
385
- _path_finder_installed = True
386
- logger.debug("[comfy_env] Installed PathBasedModuleFinder for path-based module names")
387
-
388
-
389
- def _load_path_based_module(module_name: str):
390
- """
391
- Load a module that has a filesystem path as its name.
392
-
393
- ComfyUI uses full filesystem paths as module names for custom nodes.
394
- This function handles that case by using file-based imports.
395
- """
396
- import importlib.util
397
- import os
398
- import sys
399
-
400
- # Check if it's already in sys.modules
401
- if module_name in sys.modules:
402
- return sys.modules[module_name]
403
-
404
- # Check if module_name contains submodule parts (e.g., "/path/to/pkg.submod.subsubmod")
405
- # In this case, we need to load the parent packages first
406
- if '.' in module_name:
407
- parts = module_name.split('.')
408
- # Find where the path ends and module parts begin
409
- # The path part won't exist as a directory when combined with module parts
410
- base_path = parts[0]
411
- submodule_parts = []
412
-
413
- for i, part in enumerate(parts[1:], 1):
414
- test_path = os.path.join(base_path, part)
415
- if os.path.exists(test_path):
416
- base_path = test_path
417
- else:
418
- # This and remaining parts are module names, not path components
419
- submodule_parts = parts[i:]
420
- break
421
-
422
- if submodule_parts:
423
- # Load parent package first
424
- parent_module = _load_path_based_module(base_path)
425
-
426
- # Now load submodules
427
- current_module = parent_module
428
- current_name = base_path
429
- for submod in submodule_parts:
430
- current_name = f"{current_name}.{submod}"
431
- if current_name in sys.modules:
432
- current_module = sys.modules[current_name]
433
- else:
434
- # Try to import as attribute or load from file
435
- if hasattr(current_module, submod):
436
- current_module = getattr(current_module, submod)
437
- else:
438
- # Try to load the submodule file
439
- if hasattr(current_module, '__path__'):
440
- for parent_path in current_module.__path__:
441
- submod_file = os.path.join(parent_path, submod + '.py')
442
- submod_pkg = os.path.join(parent_path, submod, '__init__.py')
443
- if os.path.exists(submod_file):
444
- spec = importlib.util.spec_from_file_location(current_name, submod_file)
445
- current_module = importlib.util.module_from_spec(spec)
446
- current_module.__package__ = f"{base_path}.{'.'.join(submodule_parts[:-1])}" if len(submodule_parts) > 1 else base_path
447
- sys.modules[current_name] = current_module
448
- spec.loader.exec_module(current_module)
449
- break
450
- elif os.path.exists(submod_pkg):
451
- spec = importlib.util.spec_from_file_location(current_name, submod_pkg,
452
- submodule_search_locations=[os.path.dirname(submod_pkg)])
453
- current_module = importlib.util.module_from_spec(spec)
454
- sys.modules[current_name] = current_module
455
- spec.loader.exec_module(current_module)
456
- break
457
- else:
458
- raise ModuleNotFoundError(f"Cannot find submodule {submod} in {current_name}")
459
- return current_module
460
-
461
- # Simple path-based module (no submodule parts)
462
- if os.path.isdir(module_name):
463
- init_path = os.path.join(module_name, "__init__.py")
464
- submodule_search_locations = [module_name]
465
- else:
466
- init_path = module_name
467
- submodule_search_locations = None
468
-
469
- if not os.path.exists(init_path):
470
- raise ModuleNotFoundError(f"Cannot find module at path: {module_name}")
471
-
472
- spec = importlib.util.spec_from_file_location(
473
- module_name,
474
- init_path,
475
- submodule_search_locations=submodule_search_locations
476
- )
477
- module = importlib.util.module_from_spec(spec)
478
-
479
- # Set up package attributes for relative imports
480
- if os.path.isdir(module_name):
481
- module.__path__ = [module_name]
482
- module.__package__ = module_name
483
- else:
484
- module.__package__ = module_name.rsplit('.', 1)[0] if '.' in module_name else ''
485
-
486
- sys.modules[module_name] = module
487
- spec.loader.exec_module(module)
488
-
489
- return module
490
-
491
-
492
- def _execute_method_call(module_name: str, class_name: str, method_name: str,
493
- self_state: dict, kwargs: dict) -> Any:
494
- """
495
- Execute a method call in the worker process.
496
-
497
- This function imports the class fresh and calls the original (un-decorated) method.
498
- """
499
- import importlib
500
- import os
501
- import sys
502
-
503
- # Import the module
504
- logger.debug(f"Attempting to import module_name={module_name}")
505
-
506
- # Check if module_name is a filesystem path (ComfyUI uses paths as module names)
507
- # This happens because ComfyUI's load_custom_node uses the full path as sys_module_name
508
- if module_name.startswith('/') or (os.sep in module_name and not module_name.startswith('.')):
509
- # Check if the base path exists to confirm it's a path-based module
510
- base_path = module_name.split('.')[0] if '.' in module_name else module_name
511
- if os.path.exists(base_path):
512
- logger.debug(f"Detected path-based module name, using file-based import")
513
- # Install the meta path finder to handle relative imports within the package
514
- _ensure_path_finder_installed()
515
- module = _load_path_based_module(module_name)
516
- else:
517
- # Doesn't look like a valid path, try standard import
518
- module = importlib.import_module(module_name)
519
- else:
520
- # Standard module name - use importlib.import_module
521
- module = importlib.import_module(module_name)
522
- cls = getattr(module, class_name)
523
-
524
- # Create instance with proper __slots__ handling
525
- instance = object.__new__(cls)
526
-
527
- # Handle both __slots__ and __dict__ based classes
528
- if hasattr(cls, '__slots__'):
529
- # Class uses __slots__ - set attributes individually
530
- for slot in cls.__slots__:
531
- if slot in self_state:
532
- setattr(instance, slot, self_state[slot])
533
- # Also check for __dict__ slot (hybrid classes)
534
- if '__dict__' in cls.__slots__ or hasattr(instance, '__dict__'):
535
- for key, value in self_state.items():
536
- if key not in cls.__slots__:
537
- setattr(instance, key, value)
538
- else:
539
- # Standard class with __dict__
540
- instance.__dict__.update(self_state)
541
-
542
- # Get the ORIGINAL method stored by the decorator, not the proxy
543
- # This avoids the infinite recursion of proxy -> worker -> proxy
544
- original_method = getattr(cls, '_isolated_original_method', None)
545
- if original_method is None:
546
- # Fallback: class wasn't decorated, use the method directly
547
- original_method = getattr(cls, method_name)
548
- return original_method(instance, **kwargs)
549
-
550
- # Call the original method (it's an unbound function, pass instance)
551
- return original_method(instance, **kwargs)
552
-
553
-
554
- class MPWorker(Worker):
555
- """
556
- Worker using torch.multiprocessing for same-venv isolation.
557
-
558
- Features:
559
- - Zero-copy CUDA tensor transfer (via CUDA IPC handles)
560
- - Zero-copy CPU tensor transfer (via shared memory)
561
- - Fresh CUDA context (subprocess has independent GPU state)
562
- - Automatic cleanup on worker death
563
-
564
- The subprocess uses 'spawn' start method, ensuring a clean Python
565
- interpreter without inherited state from the parent.
566
- """
567
-
568
- def __init__(self, name: Optional[str] = None, sys_path: Optional[list] = None, lib_path: Optional[str] = None, env_vars: Optional[dict] = None):
569
- """
570
- Initialize the worker.
571
-
572
- Args:
573
- name: Optional name for logging/debugging.
574
- sys_path: Optional list of paths to add to sys.path in worker process.
575
- lib_path: Optional path to add to LD_LIBRARY_PATH (for conda libraries).
576
- env_vars: Optional environment variables to set in worker process.
577
- """
578
- self.name = name or "MPWorker"
579
- self._sys_path = sys_path or []
580
- self._lib_path = lib_path
581
- self._env_vars = env_vars or {}
582
- self._process = None
583
- self._queue_in = None
584
- self._queue_out = None
585
- self._started = False
586
- self._shutdown = False
587
-
588
- def _ensure_started(self):
589
- """Lazily start the worker process on first call."""
590
- if self._shutdown:
591
- raise RuntimeError(f"{self.name}: Worker has been shut down")
592
-
593
- if self._started:
594
- if not self._process.is_alive():
595
- raise RuntimeError(f"{self.name}: Worker process died unexpectedly")
596
- return
597
-
598
- # Import torch here to avoid import at module level
599
- import os
600
- import sys
601
-
602
- # Clear conda/pixi environment variables FIRST, before importing multiprocessing
603
- # These can cause the child process to pick up the wrong Python interpreter
604
- # or stdlib, leading to sys.version mismatch errors in platform module
605
- conda_env_vars = [
606
- 'CONDA_PREFIX',
607
- 'CONDA_DEFAULT_ENV',
608
- 'CONDA_PYTHON_EXE',
609
- 'CONDA_EXE',
610
- 'CONDA_SHLVL',
611
- 'PYTHONHOME',
612
- 'PYTHONPATH', # Also clear PYTHONPATH to prevent pixi paths
613
- '_CE_CONDA',
614
- '_CE_M',
615
- ]
616
- saved_env = {}
617
- for var in conda_env_vars:
618
- if var in os.environ:
619
- saved_env[var] = os.environ.pop(var)
620
-
621
- # Also remove pixi paths from LD_LIBRARY_PATH
622
- ld_lib = os.environ.get('LD_LIBRARY_PATH', '')
623
- if '.pixi' in ld_lib:
624
- saved_env['LD_LIBRARY_PATH'] = ld_lib
625
- # Filter out pixi paths
626
- new_ld_lib = ':'.join(p for p in ld_lib.split(':') if '.pixi' not in p)
627
- if new_ld_lib:
628
- os.environ['LD_LIBRARY_PATH'] = new_ld_lib
629
- else:
630
- os.environ.pop('LD_LIBRARY_PATH', None)
631
-
632
- import torch.multiprocessing as mp
633
-
634
- try:
635
- # Use spawn to get clean subprocess (no inherited CUDA context)
636
- ctx = mp.get_context('spawn')
637
-
638
- # Explicitly set the spawn executable to the current Python
639
- # This prevents pixi/conda from hijacking the spawn process
640
- import multiprocessing.spawn as mp_spawn
641
- original_exe = mp_spawn.get_executable()
642
- if original_exe != sys.executable.encode() and original_exe != sys.executable:
643
- print(f"[comfy-env] Warning: spawn executable was {original_exe}, forcing to {sys.executable}")
644
- mp_spawn.set_executable(sys.executable)
645
-
646
- self._queue_in = ctx.Queue()
647
- self._queue_out = ctx.Queue()
648
- self._process = ctx.Process(
649
- target=_worker_loop,
650
- args=(self._queue_in, self._queue_out, self._sys_path, self._lib_path, self._env_vars, self.name),
651
- daemon=True,
652
- )
653
- self._process.start()
654
- self._started = True
655
-
656
- # Restore original executable setting
657
- mp_spawn.set_executable(original_exe)
658
- finally:
659
- # Restore env vars in parent process
660
- os.environ.update(saved_env)
661
-
662
- def call(
663
- self,
664
- func: Callable,
665
- *args,
666
- timeout: Optional[float] = None,
667
- **kwargs
668
- ) -> Any:
669
- """
670
- Execute a function in the worker process.
671
-
672
- Args:
673
- func: Function to execute. Must be picklable (module-level or staticmethod).
674
- *args: Positional arguments.
675
- timeout: Timeout in seconds (None = no timeout, default).
676
- **kwargs: Keyword arguments.
677
-
678
- Returns:
679
- Return value of func(*args, **kwargs).
680
-
681
- Raises:
682
- WorkerError: If func raises an exception.
683
- TimeoutError: If execution exceeds timeout.
684
- RuntimeError: If worker process dies.
685
- """
686
- self._ensure_started()
687
-
688
- # Handle tensors based on allocator
689
- if _can_use_cuda_ipc():
690
- # CUDA IPC - zero copy (works with native allocator)
691
- kwargs = {k: prepare_for_ipc_recursive(v) for k, v in kwargs.items()}
692
- args = tuple(prepare_for_ipc_recursive(a) for a in args)
693
- else:
694
- # File-based transfer (fallback for cudaMallocAsync)
695
- kwargs = _save_tensors_to_files(kwargs)
696
- args = tuple(_save_tensors_to_files(a) for a in args)
697
-
698
- # Send work item
699
- self._queue_in.put((func, args, kwargs))
700
-
701
- return self._get_result(timeout)
702
-
703
- def call_method(
704
- self,
705
- module_name: str,
706
- class_name: str,
707
- method_name: str,
708
- self_state: dict,
709
- kwargs: dict,
710
- timeout: Optional[float] = None,
711
- ) -> Any:
712
- """
713
- Execute a class method in the worker process.
714
-
715
- This uses a string-based protocol to avoid pickle issues with decorated methods.
716
- The worker imports the module fresh and calls the original (un-decorated) method.
717
-
718
- Args:
719
- module_name: Full module path (e.g., 'my_package.nodes.my_node')
720
- class_name: Class name (e.g., 'MyNode')
721
- method_name: Method name (e.g., 'process')
722
- self_state: Instance __dict__ to restore
723
- kwargs: Method keyword arguments
724
- timeout: Timeout in seconds (None = no timeout, default).
725
-
726
- Returns:
727
- Return value of method.
728
-
729
- Raises:
730
- WorkerError: If method raises an exception.
731
- TimeoutError: If execution exceeds timeout.
732
- RuntimeError: If worker process dies.
733
- """
734
- self._ensure_started()
735
-
736
- # Handle tensors based on allocator
737
- if _can_use_cuda_ipc():
738
- # CUDA IPC - zero copy (works with native allocator)
739
- kwargs = prepare_for_ipc_recursive(kwargs)
740
- else:
741
- # File-based transfer (fallback for cudaMallocAsync)
742
- kwargs = _save_tensors_to_files(kwargs)
743
-
744
- # Send method call request using protocol
745
- self._queue_in.put((
746
- _CALL_METHOD,
747
- module_name,
748
- class_name,
749
- method_name,
750
- self_state,
751
- kwargs,
752
- ))
753
-
754
- return self._get_result(timeout)
755
-
756
- def _get_result(self, timeout: Optional[float]) -> Any:
757
- """Wait for and return result from worker."""
758
- try:
759
- status, result = self._queue_out.get(timeout=timeout)
760
- except QueueEmpty:
761
- # Timeout - use graceful escalation
762
- self._handle_timeout(timeout)
763
- # _handle_timeout always raises, but just in case:
764
- raise TimeoutError(f"{self.name}: Call timed out after {timeout}s")
765
- except Exception as e:
766
- raise RuntimeError(f"{self.name}: Failed to get result: {e}")
767
-
768
- # Handle response
769
- if status == "ok":
770
- # Load tensors from temp files if using file-based transfer
771
- if not _can_use_cuda_ipc():
772
- result = _load_tensors_from_files(result)
773
- return result
774
- elif status == "error":
775
- msg, tb = result
776
- raise WorkerError(msg, traceback=tb)
777
- elif status == "fatal":
778
- self._shutdown = True
779
- raise RuntimeError(f"{self.name}: Fatal worker error: {result}")
780
- else:
781
- raise RuntimeError(f"{self.name}: Unknown response status: {status}")
782
-
783
- def shutdown(self) -> None:
784
- """Shut down the worker process."""
785
- if self._shutdown or not self._started:
786
- return
787
-
788
- self._shutdown = True
789
-
790
- try:
791
- # Send shutdown signal
792
- self._queue_in.put(_SHUTDOWN)
793
-
794
- # Wait for acknowledgment
795
- try:
796
- self._queue_out.get(timeout=5.0)
797
- except:
798
- pass
799
-
800
- # Wait for process to exit
801
- self._process.join(timeout=5.0)
802
-
803
- if self._process.is_alive():
804
- self._process.kill()
805
- self._process.join(timeout=1.0)
806
-
807
- except Exception:
808
- # Force kill if anything goes wrong
809
- if self._process and self._process.is_alive():
810
- self._process.kill()
811
-
812
- def _handle_timeout(self, timeout: float) -> None:
813
- """
814
- Handle timeout with graceful escalation.
815
-
816
- Instead of immediately killing the worker (which can leak GPU memory),
817
- try graceful shutdown first, then escalate to SIGTERM, then SIGKILL.
818
-
819
- Inspired by pyisolate's timeout handling pattern.
820
- """
821
- logger.warning(f"{self.name}: Call timed out after {timeout}s, attempting graceful shutdown")
822
-
823
- # Stage 1: Send shutdown signal, wait 3s for graceful exit
824
- try:
825
- self._queue_in.put(_SHUTDOWN)
826
- self._queue_out.get(timeout=3.0)
827
- self._process.join(timeout=2.0)
828
- if not self._process.is_alive():
829
- self._shutdown = True
830
- raise TimeoutError(f"{self.name}: Graceful shutdown after timeout ({timeout}s)")
831
- except QueueEmpty:
832
- pass
833
- except TimeoutError:
834
- raise
835
- except Exception:
836
- pass
837
-
838
- # Stage 2: SIGTERM, wait 5s
839
- if self._process.is_alive():
840
- logger.warning(f"{self.name}: Graceful shutdown failed, sending SIGTERM")
841
- self._process.terminate()
842
- self._process.join(timeout=5.0)
843
-
844
- # Stage 3: SIGKILL as last resort
845
- if self._process.is_alive():
846
- logger.error(f"{self.name}: SIGTERM failed, force killing worker (may leak GPU memory)")
847
- self._process.kill()
848
- self._process.join(timeout=1.0)
849
-
850
- self._shutdown = True
851
- raise TimeoutError(f"{self.name}: Call timed out after {timeout}s")
852
-
853
- def is_alive(self) -> bool:
854
- """Check if worker process is running or can be started."""
855
- if self._shutdown:
856
- return False
857
- # Not started yet = can still be started = "alive"
858
- if not self._started:
859
- return True
860
- return self._process.is_alive()
861
-
862
- def __repr__(self):
863
- status = "alive" if self.is_alive() else "stopped"
864
- return f"<MPWorker name={self.name!r} status={status}>"
File without changes
File without changes
File without changes