Buckets:

Quantumbraid
/

trainingdata

Files

xet

Quantumbraid/trainingdata / sparkdeps.py

Quantumbraid

about 1 month ago

download

raw

7.62 kB

"""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).")

Xet Storage Details

Size:: 7.62 kB
Xet hash:: 99d3f13a73c243d4dcfd9c3b83befbbf94efc7c3a7f481415d45f9cb406892ce

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""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).")