comfy-env 0.0.12__tar.gz → 0.0.13__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.
- {comfy_env-0.0.12 → comfy_env-0.0.13}/.gitignore +3 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/PKG-INFO +2 -1
- {comfy_env-0.0.12 → comfy_env-0.0.13}/pyproject.toml +5 -1
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/cli.py +103 -5
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/env/config_file.py +35 -32
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/env/detection.py +0 -16
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/errors.py +0 -32
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/install.py +1 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/ipc/bridge.py +0 -36
- comfy_env-0.0.13/src/comfy_env/registry.py +91 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/resolver.py +3 -17
- comfy_env-0.0.13/src/comfy_env/wheel_sources.yml +141 -0
- comfy_env-0.0.12/src/comfy_env/registry.py +0 -252
- comfy_env-0.0.12/src/comfy_env/runner.py +0 -273
- {comfy_env-0.0.12 → comfy_env-0.0.13}/.github/workflows/publish.yml +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/CLAUDE.md +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/CRITICISM.md +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/LICENSE +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/README.md +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/examples/basic_node/__init__.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/examples/basic_node/comfy-env.toml +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/examples/basic_node/nodes.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/examples/basic_node/worker.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/examples/decorator_node/__init__.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/examples/decorator_node/nodes.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/__init__.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/decorator.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/env/__init__.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/env/config.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/env/manager.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/env/platform/__init__.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/env/platform/base.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/env/platform/darwin.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/env/platform/linux.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/env/platform/windows.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/env/security.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/ipc/__init__.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/ipc/protocol.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/ipc/tensor.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/ipc/torch_bridge.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/ipc/transport.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/ipc/worker.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/stubs/__init__.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/stubs/folder_paths.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/workers/__init__.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/workers/base.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/workers/pool.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/workers/tensor_utils.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/workers/torch_mp.py +0 -0
- {comfy_env-0.0.12 → comfy_env-0.0.13}/src/comfy_env/workers/venv.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: comfy-env
|
|
3
|
-
Version: 0.0.
|
|
3
|
+
Version: 0.0.13
|
|
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
|
|
@@ -17,6 +17,7 @@ Classifier: Programming Language :: Python :: 3.11
|
|
|
17
17
|
Classifier: Programming Language :: Python :: 3.12
|
|
18
18
|
Classifier: Programming Language :: Python :: 3.13
|
|
19
19
|
Requires-Python: >=3.10
|
|
20
|
+
Requires-Dist: pyyaml>=6.0
|
|
20
21
|
Requires-Dist: tomli>=2.0.0; python_version < '3.11'
|
|
21
22
|
Requires-Dist: uv>=0.4.0
|
|
22
23
|
Provides-Extra: dev
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "comfy-env"
|
|
3
|
-
version = "0.0.
|
|
3
|
+
version = "0.0.13"
|
|
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"}
|
|
@@ -21,6 +21,7 @@ classifiers = [
|
|
|
21
21
|
dependencies = [
|
|
22
22
|
"tomli>=2.0.0; python_version < '3.11'", # TOML parsing (built-in tomllib for 3.11+)
|
|
23
23
|
"uv>=0.4.0", # Fast Python package installer and venv creator
|
|
24
|
+
"pyyaml>=6.0", # YAML parsing for wheel_sources.yml
|
|
24
25
|
]
|
|
25
26
|
|
|
26
27
|
[project.optional-dependencies]
|
|
@@ -41,6 +42,9 @@ build-backend = "hatchling.build"
|
|
|
41
42
|
[tool.hatch.build.targets.wheel]
|
|
42
43
|
packages = ["src/comfy_env"]
|
|
43
44
|
|
|
45
|
+
[tool.hatch.build.targets.wheel.force-include]
|
|
46
|
+
"src/comfy_env/wheel_sources.yml" = "comfy_env/wheel_sources.yml"
|
|
47
|
+
|
|
44
48
|
[tool.ruff]
|
|
45
49
|
line-length = 100
|
|
46
50
|
target-version = "py310"
|
|
@@ -234,6 +234,7 @@ def cmd_info(args) -> int:
|
|
|
234
234
|
def cmd_resolve(args) -> int:
|
|
235
235
|
"""Handle resolve command."""
|
|
236
236
|
from .resolver import RuntimeEnv, WheelResolver, parse_wheel_requirement
|
|
237
|
+
from .registry import PACKAGE_REGISTRY
|
|
237
238
|
from .env.config_file import discover_env_config, load_env_from_file
|
|
238
239
|
|
|
239
240
|
env = RuntimeEnv.detect()
|
|
@@ -270,19 +271,116 @@ def cmd_resolve(args) -> int:
|
|
|
270
271
|
print(f" {package}: No version specified, skipping")
|
|
271
272
|
continue
|
|
272
273
|
|
|
274
|
+
pkg_lower = package.lower()
|
|
273
275
|
try:
|
|
274
|
-
|
|
275
|
-
|
|
276
|
-
|
|
277
|
-
|
|
276
|
+
# Check if package is in registry with github_release method
|
|
277
|
+
if pkg_lower in PACKAGE_REGISTRY:
|
|
278
|
+
registry_config = PACKAGE_REGISTRY[pkg_lower]
|
|
279
|
+
method = registry_config.get("method")
|
|
280
|
+
|
|
281
|
+
if method == "github_release":
|
|
282
|
+
# Resolve URL from registry sources
|
|
283
|
+
url = _resolve_github_release_url(package, version, env, registry_config)
|
|
284
|
+
status = "OK" if args.verify else "resolved"
|
|
285
|
+
print(f" {package}=={version}: {status}")
|
|
286
|
+
print(f" {url}")
|
|
287
|
+
else:
|
|
288
|
+
# For other methods, just show what method will be used
|
|
289
|
+
print(f" {package}=={version}: uses {method} method")
|
|
290
|
+
if "index_url" in registry_config:
|
|
291
|
+
index_url = _substitute_template(registry_config["index_url"], env)
|
|
292
|
+
print(f" index: {index_url}")
|
|
293
|
+
elif "package_template" in registry_config:
|
|
294
|
+
pkg_name = _substitute_template(registry_config["package_template"], env)
|
|
295
|
+
print(f" installs as: {pkg_name}")
|
|
296
|
+
else:
|
|
297
|
+
# Fall back to WheelResolver
|
|
298
|
+
url = resolver.resolve(package, version, env, verify=args.verify)
|
|
299
|
+
status = "OK" if args.verify else "resolved"
|
|
300
|
+
print(f" {package}=={version}: {status}")
|
|
301
|
+
print(f" {url}")
|
|
278
302
|
except Exception as e:
|
|
279
303
|
print(f" {package}=={version}: FAILED")
|
|
280
|
-
|
|
304
|
+
_print_wheel_not_found_error(package, version, env, e)
|
|
281
305
|
all_ok = False
|
|
282
306
|
|
|
283
307
|
return 0 if all_ok else 1
|
|
284
308
|
|
|
285
309
|
|
|
310
|
+
def _substitute_template(template: str, env) -> str:
|
|
311
|
+
"""Substitute environment variables into a URL template."""
|
|
312
|
+
vars_dict = env.as_dict()
|
|
313
|
+
result = template
|
|
314
|
+
for key, value in vars_dict.items():
|
|
315
|
+
if value is not None:
|
|
316
|
+
result = result.replace(f"{{{key}}}", str(value))
|
|
317
|
+
return result
|
|
318
|
+
|
|
319
|
+
|
|
320
|
+
def _resolve_github_release_url(package: str, version: str, env, config: dict) -> str:
|
|
321
|
+
"""Resolve URL for github_release method packages."""
|
|
322
|
+
sources = config.get("sources", [])
|
|
323
|
+
if not sources:
|
|
324
|
+
raise ValueError(f"No sources configured for {package}")
|
|
325
|
+
|
|
326
|
+
# Build template variables
|
|
327
|
+
vars_dict = env.as_dict()
|
|
328
|
+
vars_dict["version"] = version
|
|
329
|
+
vars_dict["py_tag"] = f"cp{env.python_short}"
|
|
330
|
+
if env.cuda_version:
|
|
331
|
+
vars_dict["cuda_major"] = env.cuda_version.split(".")[0]
|
|
332
|
+
|
|
333
|
+
# Filter sources by platform
|
|
334
|
+
current_platform = env.platform_tag
|
|
335
|
+
compatible_sources = [
|
|
336
|
+
s for s in sources
|
|
337
|
+
if current_platform in s.get("platforms", [])
|
|
338
|
+
]
|
|
339
|
+
|
|
340
|
+
if not compatible_sources:
|
|
341
|
+
available = set()
|
|
342
|
+
for s in sources:
|
|
343
|
+
available.update(s.get("platforms", []))
|
|
344
|
+
raise ValueError(
|
|
345
|
+
f"No {package} wheels for platform {current_platform}. "
|
|
346
|
+
f"Available: {', '.join(sorted(available))}"
|
|
347
|
+
)
|
|
348
|
+
|
|
349
|
+
# Return URL from first compatible source
|
|
350
|
+
source = compatible_sources[0]
|
|
351
|
+
url_template = source.get("url_template", "")
|
|
352
|
+
url = url_template
|
|
353
|
+
for key, value in vars_dict.items():
|
|
354
|
+
if value is not None:
|
|
355
|
+
url = url.replace(f"{{{key}}}", str(value))
|
|
356
|
+
|
|
357
|
+
return url
|
|
358
|
+
|
|
359
|
+
|
|
360
|
+
def _print_wheel_not_found_error(package: str, version: str, env, error: Exception) -> None:
|
|
361
|
+
"""Print a formatted error message for wheel not found."""
|
|
362
|
+
from .errors import WheelNotFoundError
|
|
363
|
+
|
|
364
|
+
if isinstance(error, WheelNotFoundError):
|
|
365
|
+
print(f" CUDA wheel not found: {package}=={version}")
|
|
366
|
+
print()
|
|
367
|
+
print("+------------------------------------------------------------------+")
|
|
368
|
+
print("| CUDA Wheel Not Found |")
|
|
369
|
+
print("+------------------------------------------------------------------+")
|
|
370
|
+
print(f"| Package: {package}=={version:<46} |")
|
|
371
|
+
print(f"| Requested: cu{env.cuda_short}-torch{env.torch_mm}-{env.python_short}-{env.platform_tag:<17} |")
|
|
372
|
+
print("| |")
|
|
373
|
+
print(f"| Reason: {error.reason:<54} |")
|
|
374
|
+
print("| |")
|
|
375
|
+
print("| Suggestions: |")
|
|
376
|
+
print(f"| 1. Check if wheel exists: comfy-env resolve {package:<15} |")
|
|
377
|
+
print(f"| 2. Build wheel locally: comfy-env build {package:<18} |")
|
|
378
|
+
print("| |")
|
|
379
|
+
print("+------------------------------------------------------------------+")
|
|
380
|
+
else:
|
|
381
|
+
print(f" {error}")
|
|
382
|
+
|
|
383
|
+
|
|
286
384
|
def cmd_doctor(args) -> int:
|
|
287
385
|
"""Handle doctor command."""
|
|
288
386
|
from .install import verify_installation
|
|
@@ -168,7 +168,7 @@ def _get_default_pytorch_version(cuda_version: Optional[str]) -> str:
|
|
|
168
168
|
- CUDA 12.8 (Turing+): PyTorch 2.8.0
|
|
169
169
|
"""
|
|
170
170
|
if cuda_version == "12.4":
|
|
171
|
-
return "2.5.1" #
|
|
171
|
+
return "2.5.1" # Full: Pascal GPUs
|
|
172
172
|
return "2.8.0" # Modern: Turing through Blackwell
|
|
173
173
|
|
|
174
174
|
|
|
@@ -176,7 +176,7 @@ def _parse_config(data: Dict[str, Any], base_dir: Path) -> IsolatedEnv:
|
|
|
176
176
|
"""
|
|
177
177
|
Parse TOML data into IsolatedEnv.
|
|
178
178
|
|
|
179
|
-
Supports both simplified and
|
|
179
|
+
Supports both simplified and full config formats:
|
|
180
180
|
|
|
181
181
|
Simplified (CUDA packages only):
|
|
182
182
|
[packages]
|
|
@@ -186,7 +186,7 @@ def _parse_config(data: Dict[str, Any], base_dir: Path) -> IsolatedEnv:
|
|
|
186
186
|
Or as list:
|
|
187
187
|
packages = ["torch-scatter==2.1.2", "torch-cluster==1.6.3"]
|
|
188
188
|
|
|
189
|
-
|
|
189
|
+
Full:
|
|
190
190
|
[env]
|
|
191
191
|
name = "my-node"
|
|
192
192
|
[packages]
|
|
@@ -233,7 +233,7 @@ def _parse_config(data: Dict[str, Any], base_dir: Path) -> IsolatedEnv:
|
|
|
233
233
|
variables.setdefault("pytorch_mm", pytorch_mm)
|
|
234
234
|
|
|
235
235
|
# Parse CUDA packages - support multiple formats
|
|
236
|
-
# Priority: [cuda] section > cuda = [...] >
|
|
236
|
+
# Priority: [cuda] section > cuda = [...] > [packages] section
|
|
237
237
|
no_deps_requirements = []
|
|
238
238
|
requirements = []
|
|
239
239
|
|
|
@@ -258,12 +258,12 @@ def _parse_config(data: Dict[str, Any], base_dir: Path) -> IsolatedEnv:
|
|
|
258
258
|
|
|
259
259
|
elif isinstance(packages_section, dict):
|
|
260
260
|
# Check for simplified format: [packages] with key=value pairs
|
|
261
|
-
# vs
|
|
261
|
+
# vs old format: [packages] with requirements/no_deps lists
|
|
262
262
|
|
|
263
|
-
|
|
263
|
+
has_old_keys = any(k in packages_section for k in ["requirements", "no_deps", "requirements_file"])
|
|
264
264
|
|
|
265
|
-
if
|
|
266
|
-
#
|
|
265
|
+
if has_old_keys:
|
|
266
|
+
# Old format
|
|
267
267
|
raw_requirements = packages_section.get("requirements", [])
|
|
268
268
|
requirements = [_substitute_vars(req, variables) for req in raw_requirements]
|
|
269
269
|
|
|
@@ -340,9 +340,9 @@ def load_config(
|
|
|
340
340
|
base_dir: Optional[Path] = None,
|
|
341
341
|
) -> EnvManagerConfig:
|
|
342
342
|
"""
|
|
343
|
-
Load full EnvManagerConfig from a TOML file
|
|
343
|
+
Load full EnvManagerConfig from a TOML file.
|
|
344
344
|
|
|
345
|
-
Supports both
|
|
345
|
+
Supports both full schema (named envs) and simple format (auto-detected).
|
|
346
346
|
|
|
347
347
|
Args:
|
|
348
348
|
path: Path to the TOML config file
|
|
@@ -373,7 +373,7 @@ def load_config(
|
|
|
373
373
|
with open(path, "rb") as f:
|
|
374
374
|
data = tomllib.load(f)
|
|
375
375
|
|
|
376
|
-
return
|
|
376
|
+
return _parse_full_config(data, base_dir)
|
|
377
377
|
|
|
378
378
|
|
|
379
379
|
def discover_config(
|
|
@@ -404,9 +404,9 @@ def discover_config(
|
|
|
404
404
|
return None
|
|
405
405
|
|
|
406
406
|
|
|
407
|
-
def
|
|
407
|
+
def _parse_full_config(data: Dict[str, Any], base_dir: Path) -> EnvManagerConfig:
|
|
408
408
|
"""
|
|
409
|
-
Parse TOML data into EnvManagerConfig
|
|
409
|
+
Parse TOML data into EnvManagerConfig.
|
|
410
410
|
|
|
411
411
|
Schema:
|
|
412
412
|
[local.cuda] - CUDA packages for host
|
|
@@ -416,16 +416,16 @@ def _parse_config_v2(data: Dict[str, Any], base_dir: Path) -> EnvManagerConfig:
|
|
|
416
416
|
[envname.packages] - Packages for env
|
|
417
417
|
[node_reqs] - Node dependencies
|
|
418
418
|
|
|
419
|
-
Also supports
|
|
419
|
+
Also supports simple format ([env] + [packages]) for backward compatibility.
|
|
420
420
|
"""
|
|
421
|
-
# Detect if this is
|
|
422
|
-
|
|
421
|
+
# Detect if this is full schema or simple format
|
|
422
|
+
is_full = "local" in data or _has_named_env(data)
|
|
423
423
|
|
|
424
|
-
if not
|
|
425
|
-
#
|
|
426
|
-
return
|
|
424
|
+
if not is_full:
|
|
425
|
+
# Simple format - convert to full structure
|
|
426
|
+
return _convert_simple_to_full(data, base_dir)
|
|
427
427
|
|
|
428
|
-
# Parse
|
|
428
|
+
# Parse full schema
|
|
429
429
|
local = _parse_local_section(data.get("local", {}))
|
|
430
430
|
envs = _parse_env_sections(data, base_dir)
|
|
431
431
|
node_reqs = _parse_node_reqs(data.get("node_reqs", {}))
|
|
@@ -582,27 +582,30 @@ def _parse_node_reqs(node_reqs_data: Dict[str, Any]) -> List[NodeReq]:
|
|
|
582
582
|
return reqs
|
|
583
583
|
|
|
584
584
|
|
|
585
|
-
def
|
|
586
|
-
"""Convert
|
|
587
|
-
# Parse using legacy parser to get IsolatedEnv
|
|
588
|
-
legacy_env = _parse_config(data, base_dir)
|
|
585
|
+
def _convert_simple_to_full(data: Dict[str, Any], base_dir: Path) -> EnvManagerConfig:
|
|
586
|
+
"""Convert simple config format to full EnvManagerConfig.
|
|
589
587
|
|
|
590
|
-
|
|
591
|
-
|
|
588
|
+
Simple configs have [env] and [packages] sections but no named environments.
|
|
589
|
+
This converts them to the full format with a single named environment.
|
|
590
|
+
"""
|
|
591
|
+
# Parse using simple parser to get IsolatedEnv
|
|
592
|
+
simple_env = _parse_config(data, base_dir)
|
|
593
|
+
|
|
594
|
+
# Check if this has explicit env settings (isolated venv) vs just CUDA packages (local install)
|
|
592
595
|
env_section = data.get("env", {})
|
|
593
596
|
has_explicit_env = bool(env_section.get("name") or env_section.get("python"))
|
|
594
597
|
|
|
595
598
|
if has_explicit_env:
|
|
596
|
-
#
|
|
599
|
+
# Isolated venv config
|
|
597
600
|
return EnvManagerConfig(
|
|
598
601
|
local=LocalConfig(),
|
|
599
|
-
envs={
|
|
602
|
+
envs={simple_env.name: simple_env},
|
|
600
603
|
node_reqs=[],
|
|
601
604
|
)
|
|
602
605
|
else:
|
|
603
|
-
#
|
|
606
|
+
# Local CUDA packages only (no isolated venv)
|
|
604
607
|
cuda_packages = {}
|
|
605
|
-
for req in
|
|
608
|
+
for req in simple_env.no_deps_requirements:
|
|
606
609
|
if "==" in req:
|
|
607
610
|
pkg, ver = req.split("==", 1)
|
|
608
611
|
cuda_packages[pkg] = ver
|
|
@@ -612,8 +615,8 @@ def _convert_legacy_to_v2(data: Dict[str, Any], base_dir: Path) -> EnvManagerCon
|
|
|
612
615
|
return EnvManagerConfig(
|
|
613
616
|
local=LocalConfig(
|
|
614
617
|
cuda_packages=cuda_packages,
|
|
615
|
-
requirements=
|
|
618
|
+
requirements=simple_env.requirements,
|
|
616
619
|
),
|
|
617
620
|
envs={},
|
|
618
621
|
node_reqs=[],
|
|
619
|
-
)
|
|
622
|
+
)
|
|
@@ -92,22 +92,6 @@ def is_blackwell_gpu(name: str, compute_cap: str) -> bool:
|
|
|
92
92
|
return False
|
|
93
93
|
|
|
94
94
|
|
|
95
|
-
def needs_cuda_128() -> bool:
|
|
96
|
-
"""
|
|
97
|
-
Check if any detected GPU requires CUDA 12.8.
|
|
98
|
-
|
|
99
|
-
Returns:
|
|
100
|
-
True if Blackwell GPU detected, False otherwise.
|
|
101
|
-
"""
|
|
102
|
-
gpus = detect_gpu_info()
|
|
103
|
-
|
|
104
|
-
for gpu in gpus:
|
|
105
|
-
if is_blackwell_gpu(gpu["name"], gpu["compute_cap"]):
|
|
106
|
-
return True
|
|
107
|
-
|
|
108
|
-
return False
|
|
109
|
-
|
|
110
|
-
|
|
111
95
|
def is_legacy_gpu(compute_cap: str) -> bool:
|
|
112
96
|
"""
|
|
113
97
|
Check if GPU is Pascal or older (requires legacy CUDA/PyTorch).
|
|
@@ -291,35 +291,3 @@ class InstallError(EnvManagerError):
|
|
|
291
291
|
|
|
292
292
|
details = "\n".join(details_parts) if details_parts else None
|
|
293
293
|
super().__init__(message, details)
|
|
294
|
-
|
|
295
|
-
|
|
296
|
-
def format_environment_mismatch(
|
|
297
|
-
expected: "RuntimeEnv",
|
|
298
|
-
actual: "RuntimeEnv",
|
|
299
|
-
) -> str:
|
|
300
|
-
"""
|
|
301
|
-
Format a message explaining environment mismatch.
|
|
302
|
-
|
|
303
|
-
Used when the current environment doesn't match what's needed.
|
|
304
|
-
"""
|
|
305
|
-
mismatches = []
|
|
306
|
-
|
|
307
|
-
if expected.cuda_version != actual.cuda_version:
|
|
308
|
-
mismatches.append(
|
|
309
|
-
f" CUDA: expected {expected.cuda_version}, got {actual.cuda_version}"
|
|
310
|
-
)
|
|
311
|
-
|
|
312
|
-
if expected.torch_version != actual.torch_version:
|
|
313
|
-
mismatches.append(
|
|
314
|
-
f" PyTorch: expected {expected.torch_version}, got {actual.torch_version}"
|
|
315
|
-
)
|
|
316
|
-
|
|
317
|
-
if expected.python_version != actual.python_version:
|
|
318
|
-
mismatches.append(
|
|
319
|
-
f" Python: expected {expected.python_version}, got {actual.python_version}"
|
|
320
|
-
)
|
|
321
|
-
|
|
322
|
-
if not mismatches:
|
|
323
|
-
return "Environment matches expected configuration"
|
|
324
|
-
|
|
325
|
-
return "Environment mismatch:\n" + "\n".join(mismatches)
|
|
@@ -48,10 +48,6 @@ class WorkerBridge:
|
|
|
48
48
|
result = bridge.call("process", image=my_image)
|
|
49
49
|
"""
|
|
50
50
|
|
|
51
|
-
# Singleton instances by environment hash
|
|
52
|
-
_instances: Dict[str, "WorkerBridge"] = {}
|
|
53
|
-
_instances_lock = threading.Lock()
|
|
54
|
-
|
|
55
51
|
def __init__(
|
|
56
52
|
self,
|
|
57
53
|
env: IsolatedEnv,
|
|
@@ -81,38 +77,6 @@ class WorkerBridge:
|
|
|
81
77
|
self._process_lock = threading.Lock()
|
|
82
78
|
self._stderr_thread: Optional[threading.Thread] = None
|
|
83
79
|
|
|
84
|
-
@classmethod
|
|
85
|
-
def get_instance(
|
|
86
|
-
cls,
|
|
87
|
-
env: IsolatedEnv,
|
|
88
|
-
worker_script: Path,
|
|
89
|
-
base_dir: Optional[Path] = None,
|
|
90
|
-
log_callback: Optional[Callable[[str], None]] = None,
|
|
91
|
-
) -> "WorkerBridge":
|
|
92
|
-
"""
|
|
93
|
-
Get or create a singleton bridge instance for an environment.
|
|
94
|
-
|
|
95
|
-
Args:
|
|
96
|
-
env: Isolated environment configuration
|
|
97
|
-
worker_script: Path to the worker Python script
|
|
98
|
-
base_dir: Base directory for environments
|
|
99
|
-
log_callback: Optional callback for logging
|
|
100
|
-
|
|
101
|
-
Returns:
|
|
102
|
-
WorkerBridge instance (reused if same env hash)
|
|
103
|
-
"""
|
|
104
|
-
env_hash = env.get_env_hash()
|
|
105
|
-
|
|
106
|
-
with cls._instances_lock:
|
|
107
|
-
if env_hash not in cls._instances:
|
|
108
|
-
cls._instances[env_hash] = cls(
|
|
109
|
-
env=env,
|
|
110
|
-
worker_script=worker_script,
|
|
111
|
-
base_dir=base_dir,
|
|
112
|
-
log_callback=log_callback,
|
|
113
|
-
)
|
|
114
|
-
return cls._instances[env_hash]
|
|
115
|
-
|
|
116
80
|
@classmethod
|
|
117
81
|
def from_config_file(
|
|
118
82
|
cls,
|
|
@@ -0,0 +1,91 @@
|
|
|
1
|
+
"""Built-in registry of CUDA packages and their wheel sources.
|
|
2
|
+
|
|
3
|
+
This module loads package configurations from wheel_sources.yml and provides
|
|
4
|
+
lookup functions for the install module.
|
|
5
|
+
|
|
6
|
+
Install method types:
|
|
7
|
+
- "index": Use pip --extra-index-url (PEP 503 simple repository)
|
|
8
|
+
- "github_index": GitHub Pages index (--find-links)
|
|
9
|
+
- "find_links": Use pip --find-links (for PyG, etc.)
|
|
10
|
+
- "pypi_variant": Package name varies by CUDA version (e.g., spconv-cu124)
|
|
11
|
+
- "github_release": Direct wheel URL from GitHub releases with fallback sources
|
|
12
|
+
- "pypi": Standard PyPI install
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
from pathlib import Path
|
|
16
|
+
from typing import Any, Dict, Optional
|
|
17
|
+
|
|
18
|
+
import yaml
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def get_cuda_short2(cuda_version: str) -> str:
|
|
22
|
+
"""Convert CUDA version to 2-3 digit format for spconv.
|
|
23
|
+
|
|
24
|
+
spconv uses "cu124" not "cu1240" for CUDA 12.4.
|
|
25
|
+
|
|
26
|
+
Args:
|
|
27
|
+
cuda_version: CUDA version string (e.g., "12.4", "12.8")
|
|
28
|
+
|
|
29
|
+
Returns:
|
|
30
|
+
Short format string (e.g., "124", "128")
|
|
31
|
+
|
|
32
|
+
Examples:
|
|
33
|
+
>>> get_cuda_short2("12.4")
|
|
34
|
+
'124'
|
|
35
|
+
>>> get_cuda_short2("12.8")
|
|
36
|
+
'128'
|
|
37
|
+
>>> get_cuda_short2("11.8")
|
|
38
|
+
'118'
|
|
39
|
+
"""
|
|
40
|
+
parts = cuda_version.split(".")
|
|
41
|
+
major = parts[0]
|
|
42
|
+
minor = parts[1] if len(parts) > 1 else "0"
|
|
43
|
+
return f"{major}{minor}"
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def _load_wheel_sources() -> Dict[str, Dict[str, Any]]:
|
|
47
|
+
"""Load package registry from wheel_sources.yml."""
|
|
48
|
+
yml_path = Path(__file__).parent / "wheel_sources.yml"
|
|
49
|
+
with open(yml_path, "r") as f:
|
|
50
|
+
data = yaml.safe_load(f)
|
|
51
|
+
return data.get("packages", {})
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
# Load registry at module import time
|
|
55
|
+
PACKAGE_REGISTRY: Dict[str, Dict[str, Any]] = _load_wheel_sources()
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def get_package_info(package: str) -> Optional[Dict[str, Any]]:
|
|
59
|
+
"""Get registry info for a package.
|
|
60
|
+
|
|
61
|
+
Args:
|
|
62
|
+
package: Package name (case-insensitive)
|
|
63
|
+
|
|
64
|
+
Returns:
|
|
65
|
+
Registry entry dict or None if not found
|
|
66
|
+
"""
|
|
67
|
+
return PACKAGE_REGISTRY.get(package.lower())
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def list_packages() -> Dict[str, str]:
|
|
71
|
+
"""List all registered packages with their descriptions.
|
|
72
|
+
|
|
73
|
+
Returns:
|
|
74
|
+
Dict mapping package name to description
|
|
75
|
+
"""
|
|
76
|
+
return {
|
|
77
|
+
name: info.get("description", "No description")
|
|
78
|
+
for name, info in PACKAGE_REGISTRY.items()
|
|
79
|
+
}
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
def is_registered(package: str) -> bool:
|
|
83
|
+
"""Check if a package is in the registry.
|
|
84
|
+
|
|
85
|
+
Args:
|
|
86
|
+
package: Package name (case-insensitive)
|
|
87
|
+
|
|
88
|
+
Returns:
|
|
89
|
+
True if package is registered
|
|
90
|
+
"""
|
|
91
|
+
return package.lower() in PACKAGE_REGISTRY
|
|
@@ -218,21 +218,6 @@ class WheelSource:
|
|
|
218
218
|
return package.lower() in [p.lower() for p in self.packages]
|
|
219
219
|
|
|
220
220
|
|
|
221
|
-
# Default wheel sources for common CUDA packages
|
|
222
|
-
DEFAULT_WHEEL_SOURCES = [
|
|
223
|
-
WheelSource(
|
|
224
|
-
name="nvdiffrast-wheels",
|
|
225
|
-
url_template="https://github.com/PozzettiAndrea/nvdiffrast-full-wheels/releases/download/v{version}/nvdiffrast-{version}%2Bcu{cuda_short}torch{torch_mm}-cp{py_short}-cp{py_short}-{platform}.whl",
|
|
226
|
-
packages=["nvdiffrast"],
|
|
227
|
-
),
|
|
228
|
-
WheelSource(
|
|
229
|
-
name="cumesh-wheels",
|
|
230
|
-
url_template="https://github.com/PozzettiAndrea/cumesh-wheels/releases/download/v{version}/{package}-{version}%2Bcu{cuda_short}torch{torch_mm}-cp{py_short}-cp{py_short}-{platform}.whl",
|
|
231
|
-
packages=["pytorch3d", "torch-cluster", "torch-scatter", "torch-sparse"],
|
|
232
|
-
),
|
|
233
|
-
]
|
|
234
|
-
|
|
235
|
-
|
|
236
221
|
class WheelResolver:
|
|
237
222
|
"""
|
|
238
223
|
Resolves CUDA wheel URLs from package name and runtime environment.
|
|
@@ -255,10 +240,11 @@ class WheelResolver:
|
|
|
255
240
|
Initialize resolver.
|
|
256
241
|
|
|
257
242
|
Args:
|
|
258
|
-
sources: List of WheelSource configurations.
|
|
243
|
+
sources: List of WheelSource configurations. Defaults to empty
|
|
244
|
+
(use PACKAGE_REGISTRY in install.py for actual sources).
|
|
259
245
|
overrides: Package-specific URL overrides (package -> template).
|
|
260
246
|
"""
|
|
261
|
-
self.sources = sources or
|
|
247
|
+
self.sources = sources or []
|
|
262
248
|
self.overrides = overrides or {}
|
|
263
249
|
|
|
264
250
|
def resolve(
|