fix(rocm): guard c10d stub, fix TorchIndexFamily for 7.1, clean dead code + comments

- worker.py: wrap c10d stub injection in `if _c10d_key not in sys.modules` so
  Windows NVIDIA users with a real torch.distributed are never affected
- install.ps1: fix Get-TauriTorchIndexFamily receiving hardcoded "rocm7.2"
  even when ROCm 7.1 wheels are installed; now branches on $ROCmVersion
- main.py: remove dead `import ctypes as _ctypes` (ctypes is never called)
- hardware.py, install_python_stack.py, worker.py, install.ps1: shorten
  verbose multi-line comment blocks throughout
- tests: update 4 stale assertions that expected rocm7.2 to be absent/capped
This commit is contained in:
LeoBorcherding 2026-05-11 00:08:14 -05:00
commit 85841b5cad
6 changed files with 49 additions and 87 deletions

View file

@ -78,8 +78,6 @@ _ROCM_WINDOWS_WHEEL_BASE = (
or "https://repo.radeon.com/rocm/windows"
).rstrip("/")
# Maps (major, minor) → (release_folder, [wheel_filename, ...])
# Includes rocm_sdk_core and rocm_sdk_libraries_custom because the torch
# wheels declare them as hard dependencies (rocm[libraries]==<ver>).
_ROCM_WINDOWS_RELEASES: dict[tuple[int, int], tuple[str, list[str]]] = {
(7, 2): (
"rocm-rel-7.2.1",
@ -334,9 +332,7 @@ def _ensure_rocm_torch() -> None:
Uses pip_install() to respect uv, constraints, and --python targeting.
"""
global _rocm_windows_torch_installed
# setup.ps1 sets this env var when it successfully installs AMD wheels
# before calling install_python_stack.py, so we can skip the subprocess
# probe and avoid reinstalling what was just installed.
# setup.ps1 sets this when it already installed AMD wheels; skip the probe.
if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
_rocm_windows_torch_installed = True
return