comfy-env 0.0.12__tar.gz → 0.0.14__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 (51) hide show
  1. {comfy_env-0.0.12 → comfy_env-0.0.14}/.gitignore +3 -0
  2. {comfy_env-0.0.12 → comfy_env-0.0.14}/PKG-INFO +2 -1
  3. {comfy_env-0.0.12 → comfy_env-0.0.14}/pyproject.toml +5 -1
  4. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/cli.py +103 -5
  5. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/env/config_file.py +35 -32
  6. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/env/detection.py +0 -16
  7. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/env/manager.py +2 -2
  8. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/errors.py +0 -32
  9. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/install.py +1 -0
  10. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/ipc/bridge.py +0 -36
  11. comfy_env-0.0.14/src/comfy_env/registry.py +91 -0
  12. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/resolver.py +6 -17
  13. comfy_env-0.0.14/src/comfy_env/wheel_sources.yml +141 -0
  14. comfy_env-0.0.14/untitled.txt +0 -0
  15. comfy_env-0.0.12/src/comfy_env/registry.py +0 -252
  16. comfy_env-0.0.12/src/comfy_env/runner.py +0 -273
  17. {comfy_env-0.0.12 → comfy_env-0.0.14}/.github/workflows/publish.yml +0 -0
  18. {comfy_env-0.0.12 → comfy_env-0.0.14}/CLAUDE.md +0 -0
  19. {comfy_env-0.0.12 → comfy_env-0.0.14}/CRITICISM.md +0 -0
  20. {comfy_env-0.0.12 → comfy_env-0.0.14}/LICENSE +0 -0
  21. {comfy_env-0.0.12 → comfy_env-0.0.14}/README.md +0 -0
  22. {comfy_env-0.0.12 → comfy_env-0.0.14}/examples/basic_node/__init__.py +0 -0
  23. {comfy_env-0.0.12 → comfy_env-0.0.14}/examples/basic_node/comfy-env.toml +0 -0
  24. {comfy_env-0.0.12 → comfy_env-0.0.14}/examples/basic_node/nodes.py +0 -0
  25. {comfy_env-0.0.12 → comfy_env-0.0.14}/examples/basic_node/worker.py +0 -0
  26. {comfy_env-0.0.12 → comfy_env-0.0.14}/examples/decorator_node/__init__.py +0 -0
  27. {comfy_env-0.0.12 → comfy_env-0.0.14}/examples/decorator_node/nodes.py +0 -0
  28. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/__init__.py +0 -0
  29. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/decorator.py +0 -0
  30. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/env/__init__.py +0 -0
  31. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/env/config.py +0 -0
  32. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/env/platform/__init__.py +0 -0
  33. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/env/platform/base.py +0 -0
  34. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/env/platform/darwin.py +0 -0
  35. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/env/platform/linux.py +0 -0
  36. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/env/platform/windows.py +0 -0
  37. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/env/security.py +0 -0
  38. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/ipc/__init__.py +0 -0
  39. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/ipc/protocol.py +0 -0
  40. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/ipc/tensor.py +0 -0
  41. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/ipc/torch_bridge.py +0 -0
  42. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/ipc/transport.py +0 -0
  43. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/ipc/worker.py +0 -0
  44. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/stubs/__init__.py +0 -0
  45. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/stubs/folder_paths.py +0 -0
  46. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/workers/__init__.py +0 -0
  47. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/workers/base.py +0 -0
  48. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/workers/pool.py +0 -0
  49. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/workers/tensor_utils.py +0 -0
  50. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/workers/torch_mp.py +0 -0
  51. {comfy_env-0.0.12 → comfy_env-0.0.14}/src/comfy_env/workers/venv.py +0 -0
@@ -19,6 +19,9 @@ ENV/
19
19
  *.swp
20
20
  *.swo
21
21
 
22
+ # Jupyter
23
+ .ipynb_checkpoints/
24
+
22
25
  # Testing
23
26
  .pytest_cache/
24
27
  .coverage
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: comfy-env
3
- Version: 0.0.12
3
+ Version: 0.0.14
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.12"
3
+ version = "0.0.14"
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
- url = resolver.resolve(package, version, env, verify=args.verify)
275
- status = "OK" if args.verify else "resolved"
276
- print(f" {package}=={version}: {status}")
277
- print(f" {url}")
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
- print(f" {e}")
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" # Legacy: Pascal GPUs
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 legacy config formats:
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
- Legacy:
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 = [...] > legacy [packages] section
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 legacy format: [packages] with requirements/no_deps lists
261
+ # vs old format: [packages] with requirements/no_deps lists
262
262
 
263
- has_legacy_keys = any(k in packages_section for k in ["requirements", "no_deps", "requirements_file"])
263
+ has_old_keys = any(k in packages_section for k in ["requirements", "no_deps", "requirements_file"])
264
264
 
265
- if has_legacy_keys:
266
- # Legacy format
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 (v2 schema).
343
+ Load full EnvManagerConfig from a TOML file.
344
344
 
345
- Supports both v2 schema and legacy format (auto-detected).
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 _parse_config_v2(data, base_dir)
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 _parse_config_v2(data: Dict[str, Any], base_dir: Path) -> EnvManagerConfig:
407
+ def _parse_full_config(data: Dict[str, Any], base_dir: Path) -> EnvManagerConfig:
408
408
  """
409
- Parse TOML data into EnvManagerConfig (v2 schema).
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 legacy format for backward compatibility.
419
+ Also supports simple format ([env] + [packages]) for backward compatibility.
420
420
  """
421
- # Detect if this is v2 schema or legacy
422
- is_v2 = "local" in data or _has_named_env(data)
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 is_v2:
425
- # Legacy format - convert to v2 structure
426
- return _convert_legacy_to_v2(data, base_dir)
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 v2 schema
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 _convert_legacy_to_v2(data: Dict[str, Any], base_dir: Path) -> EnvManagerConfig:
586
- """Convert legacy config format to v2 EnvManagerConfig."""
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
- # Check if this is really just a Type 2 config (local CUDA only, no venv)
591
- # Type 2 indicators: no [env] section with name, or name matches directory
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
- # This is a Type 1 config (isolated venv)
599
+ # Isolated venv config
597
600
  return EnvManagerConfig(
598
601
  local=LocalConfig(),
599
- envs={legacy_env.name: legacy_env},
602
+ envs={simple_env.name: simple_env},
600
603
  node_reqs=[],
601
604
  )
602
605
  else:
603
- # This is a Type 2 config (local CUDA only)
606
+ # Local CUDA packages only (no isolated venv)
604
607
  cuda_packages = {}
605
- for req in legacy_env.no_deps_requirements:
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=legacy_env.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).
@@ -475,7 +475,7 @@ class IsolatedEnvManager:
475
475
  url_template = source["url_template"]
476
476
  url = self._substitute_template(url_template, vars_dict)
477
477
 
478
- self.log(f" Trying {source.get('name', 'unknown')}: {url[:80]}...")
478
+ self.log(f" Trying {source.get('name', 'unknown')}: {url}")
479
479
  result = subprocess.run(
480
480
  pip_args + ["--no-deps", url],
481
481
  capture_output=True, text=True,
@@ -484,7 +484,7 @@ class IsolatedEnvManager:
484
484
  if result.returncode == 0:
485
485
  return # Success!
486
486
 
487
- errors.append(f"{source.get('name', 'unknown')}: {result.stderr[:100]}")
487
+ errors.append(f"{source.get('name', 'unknown')}: {result.stderr.strip()}")
488
488
 
489
489
  # All sources failed
490
490
  raise RuntimeError(
@@ -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)
@@ -421,6 +421,7 @@ def _install_from_github_release(
421
421
  url = url.replace(f"{{{key}}}", str(value))
422
422
 
423
423
  log(f" Trying {source_name}: {package}=={version}...")
424
+ log(f" Resolved wheel to: {url}")
424
425
 
425
426
  try:
426
427
  pip_cmd = _get_pip_command()
@@ -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
@@ -137,11 +137,14 @@ class RuntimeEnv:
137
137
  "py_version": self.python_version,
138
138
  "py_short": self.python_short,
139
139
  "py_minor": py_minor,
140
+ "py_tag": f"cp{self.python_short}", # e.g., cp310, cp311
140
141
  }
141
142
 
142
143
  if self.cuda_version:
143
144
  result["cuda_version"] = self.cuda_version
144
145
  result["cuda_short"] = self.cuda_short
146
+ # cuda_major: just the major version (e.g., "12" from "12.8")
147
+ result["cuda_major"] = self.cuda_version.split(".")[0]
145
148
 
146
149
  if self.torch_version:
147
150
  result["torch_version"] = self.torch_version
@@ -218,21 +221,6 @@ class WheelSource:
218
221
  return package.lower() in [p.lower() for p in self.packages]
219
222
 
220
223
 
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
224
  class WheelResolver:
237
225
  """
238
226
  Resolves CUDA wheel URLs from package name and runtime environment.
@@ -255,10 +243,11 @@ class WheelResolver:
255
243
  Initialize resolver.
256
244
 
257
245
  Args:
258
- sources: List of WheelSource configurations.
246
+ sources: List of WheelSource configurations. Defaults to empty
247
+ (use PACKAGE_REGISTRY in install.py for actual sources).
259
248
  overrides: Package-specific URL overrides (package -> template).
260
249
  """
261
- self.sources = sources or DEFAULT_WHEEL_SOURCES
250
+ self.sources = sources or []
262
251
  self.overrides = overrides or {}
263
252
 
264
253
  def resolve(