225 lines
7.4 KiB
Python
225 lines
7.4 KiB
Python
"""RTX Spark native-stack runtime provisioning (Windows on ARM, cu134).
|
|
|
|
Goal: a fresh Spark machine runs run_windows.bat and gets as close to
|
|
zero-manual-setup as licensing allows. The native wheels (torch etc.) do not
|
|
bundle CUDA / cuDNN / BLAS DLLs, and triton's launcher JIT wants MSVC. Policy:
|
|
NVIDIA components are never redistributed by us.
|
|
|
|
- CUDA 13.4 toolkit (developer preview): MANUAL install — the preview EULA
|
|
requires NVIDIA's own click-through, so the manager only detects it and
|
|
prints instructions when missing. This is the single manual step.
|
|
- cuDNN (arm64): auto-downloaded from NVIDIA's own official installer URL and
|
|
installed silently — fetched directly from NVIDIA, not redistributed.
|
|
- Arm Performance Libraries: auto-install via winget (official Arm package).
|
|
- MSVC Build Tools (triton torch.compile JIT only): auto-install via winget;
|
|
failure downgrades gracefully (training works, no torch.compile).
|
|
- VC redistributable (arm64): auto-install via winget when msvcp140 missing.
|
|
|
|
Everything is best-effort with warnings; the training stack itself only hard-
|
|
requires the CUDA + cuDNN + APL DLL dirs.
|
|
"""
|
|
|
|
import glob
|
|
import os
|
|
import subprocess
|
|
|
|
from .util import download, info, ok, warn, which
|
|
|
|
CUDA_DOWNLOAD_PAGE = (
|
|
"https://developer.nvidia.com/cuda-13-4-0-download-archive"
|
|
"?target_os=Windows&target_arch=arm64"
|
|
)
|
|
# NVIDIA's official public installer for cuDNN on Windows arm64. Downloaded
|
|
# straight from NVIDIA at install time (we do not redistribute it). Update
|
|
# together with the wheel set when moving to a newer cuDNN.
|
|
CUDNN_INSTALLER_URL = (
|
|
"https://developer.download.nvidia.com/compute/cudnn/9.25.0/"
|
|
"local_installers/cudnn_9.25.0_windows_arm64.exe"
|
|
)
|
|
|
|
# System install roots, newest version preferred (globs, not pinned versions)
|
|
_CUDA_BIN_GLOB = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v*\bin\arm64"
|
|
_CUDNN_BIN_GLOB = r"C:\Program Files\NVIDIA\CUDNN\v*\bin\*\arm64"
|
|
_ARMPL_BIN_GLOB = r"C:\Program Files\Arm Performance Libraries\armpl_*\bin"
|
|
|
|
_VS_BUILDTOOLS = (
|
|
r"C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools"
|
|
)
|
|
|
|
|
|
def _newest(pattern):
|
|
matches = sorted(glob.glob(pattern))
|
|
return matches[-1] if matches else None
|
|
|
|
|
|
def cuda_bin_dir():
|
|
return _newest(_CUDA_BIN_GLOB)
|
|
|
|
|
|
def cuda_root():
|
|
d = cuda_bin_dir()
|
|
# <root>\bin\arm64 -> <root>
|
|
return os.path.dirname(os.path.dirname(d)) if d else None
|
|
|
|
|
|
def cudnn_bin_dir():
|
|
return _newest(_CUDNN_BIN_GLOB)
|
|
|
|
|
|
def armpl_bin_dir():
|
|
return _newest(_ARMPL_BIN_GLOB)
|
|
|
|
|
|
def resolve_dll_dirs():
|
|
"""All runtime DLL dirs for the native stack (existing ones only)."""
|
|
return [d for d in (cuda_bin_dir(), cudnn_bin_dir(), armpl_bin_dir()) if d]
|
|
|
|
|
|
def runtime_complete():
|
|
return bool(cuda_bin_dir() and cudnn_bin_dir() and armpl_bin_dir())
|
|
|
|
|
|
def triton_tool_env():
|
|
"""TRITON_*_PATH env for ptxas etc. from the system CUDA install.
|
|
|
|
Our triton wheel deliberately does NOT bundle NVIDIA's compiler tools
|
|
(developer-preview licensing); resolve them from the user's toolkit.
|
|
"""
|
|
root = cuda_root()
|
|
if not root:
|
|
return {}
|
|
env = {}
|
|
for var, exe in (
|
|
("TRITON_PTXAS_PATH", "ptxas.exe"),
|
|
("TRITON_PTXAS_BLACKWELL_PATH", "ptxas.exe"),
|
|
("TRITON_CUOBJDUMP_PATH", "cuobjdump.exe"),
|
|
("TRITON_NVDISASM_PATH", "nvdisasm.exe"),
|
|
):
|
|
path = os.path.join(root, "bin", exe)
|
|
if os.path.isfile(path):
|
|
env[var] = path
|
|
return env
|
|
|
|
|
|
def check_cuda():
|
|
"""CUDA toolkit is the one manual install (preview EULA). Detect + guide."""
|
|
if cuda_bin_dir():
|
|
return True
|
|
warn(
|
|
"The CUDA 13.4 toolkit (arm64) is not installed. NVIDIA's developer "
|
|
"preview license requires installing it manually:\n"
|
|
" 1. Download from %s\n"
|
|
" 2. Install with default settings, then re-run this setup.\n"
|
|
"The RTX Spark developer driver (R616+) is required as well."
|
|
% CUDA_DOWNLOAD_PAGE
|
|
)
|
|
return False
|
|
|
|
|
|
def ensure_cudnn(dry_run=False):
|
|
"""Fetch + silently run NVIDIA's official cuDNN installer if missing."""
|
|
if cudnn_bin_dir():
|
|
return True
|
|
if dry_run:
|
|
info("[dry-run] would download and install cuDNN from NVIDIA")
|
|
return False
|
|
import tempfile
|
|
|
|
tmp = tempfile.mkdtemp(prefix="aitk_cudnn_")
|
|
try:
|
|
exe = os.path.join(tmp, os.path.basename(CUDNN_INSTALLER_URL))
|
|
download(CUDNN_INSTALLER_URL, exe, label="cuDNN (from NVIDIA)")
|
|
info("Installing cuDNN (silent)...")
|
|
code = subprocess.call([exe, "-s"])
|
|
if code != 0:
|
|
warn("cuDNN installer exited with %d." % code)
|
|
return cudnn_bin_dir() is not None
|
|
finally:
|
|
import shutil
|
|
|
|
shutil.rmtree(tmp, ignore_errors=True)
|
|
|
|
|
|
def have_msvc():
|
|
return bool(
|
|
glob.glob(os.path.join(_VS_BUILDTOOLS, "VC", "Tools", "MSVC", "*",
|
|
"bin", "Hostarm64", "arm64", "cl.exe"))
|
|
)
|
|
|
|
|
|
def _winget_install(args, label, dry_run=False):
|
|
winget = which("winget")
|
|
if not winget:
|
|
warn("winget not available — cannot auto-install %s." % label)
|
|
return False
|
|
if dry_run:
|
|
info("[dry-run] would winget install %s" % label)
|
|
return False
|
|
info("Installing %s (one-time, may take several minutes)..." % label)
|
|
code = subprocess.call(
|
|
[winget, "install", "--exact", "--source", "winget",
|
|
"--accept-source-agreements", "--accept-package-agreements"] + args,
|
|
stdout=subprocess.DEVNULL,
|
|
)
|
|
if code != 0:
|
|
warn("%s install failed (winget exit %d)." % (label, code))
|
|
return code == 0
|
|
|
|
|
|
def ensure_armpl(dry_run=False):
|
|
if armpl_bin_dir():
|
|
return True
|
|
return _winget_install(
|
|
["--id", "Arm.ArmPerformanceLibraries"],
|
|
"Arm Performance Libraries",
|
|
dry_run=dry_run,
|
|
)
|
|
|
|
|
|
def ensure_vcredist(dry_run=False):
|
|
"""VC runtime (msvcp140 etc.) — required by the native wheels."""
|
|
sysdir = os.path.join(os.environ.get("SystemRoot", r"C:\Windows"), "System32")
|
|
if os.path.isfile(os.path.join(sysdir, "msvcp140.dll")):
|
|
return True
|
|
return _winget_install(
|
|
["--id", "Microsoft.VCRedist.2015+.arm64"],
|
|
"Visual C++ Redistributable (arm64)",
|
|
dry_run=dry_run,
|
|
)
|
|
|
|
|
|
def ensure_msvc(dry_run=False):
|
|
"""MSVC Build Tools — only needed for triton's runtime kernel launchers.
|
|
|
|
Best effort: without it, training still works; torch.compile / triton
|
|
JIT is unavailable until the user installs Build Tools.
|
|
"""
|
|
if have_msvc():
|
|
return True
|
|
done = _winget_install(
|
|
["--id", "Microsoft.VisualStudio.2022.BuildTools", "--override",
|
|
"--quiet --wait --norestart "
|
|
"--add Microsoft.VisualStudio.Workload.VCTools "
|
|
"--add Microsoft.VisualStudio.Component.VC.Tools.ARM64 "
|
|
"--add Microsoft.VisualStudio.Component.Windows11SDK.26100"],
|
|
"MSVC Build Tools (for torch.compile/triton)",
|
|
dry_run=dry_run,
|
|
)
|
|
if not done and not dry_run:
|
|
warn(
|
|
"torch.compile/triton kernel JIT will be unavailable until MSVC "
|
|
"Build Tools are installed; training itself is unaffected."
|
|
)
|
|
return done
|
|
|
|
|
|
def ensure_spark_runtime(dry_run=False):
|
|
"""Full best-effort provisioning for the native Spark stack."""
|
|
check_cuda()
|
|
ensure_vcredist(dry_run=dry_run)
|
|
ensure_cudnn(dry_run=dry_run)
|
|
ensure_armpl(dry_run=dry_run)
|
|
ensure_msvc(dry_run=dry_run)
|
|
if runtime_complete():
|
|
ok("Spark native runtime present (CUDA + cuDNN + Arm PL).")
|