Edit on GitHub

starplast.install_gpu

Installing the GPU stack for the machine this is actually running on.

Why this is a run-time command and not something the installer does: a wheel's dependencies are fixed when the wheel is BUILT. pip resolves them from static metadata, and the only conditionals available are environment markers -- sys_platform, python_version, platform_machine. There is no marker for "this machine has an NVIDIA card". Probing at install time would mean shipping an sdist with executable setup.py, which loses the wheel for the common case, breaks --only-binary, and makes one version resolve to different dependency trees on different machines, so lockfiles stop meaning anything. In CI or a container it would probe the build machine rather than the target.

There is a second reason, which is that finding a card does not tell you which CUDA to install. The driver does. Guess wrong and it is two gigabytes of wheels that cannot load.

So detection happens here, when the real machine is known, and the user is told what will be installed before anything is downloaded.

WHEELS = {12: ['cuml-cu12>=24.10', 'cupy-cuda12x>=13.0'], 13: ['cuml-cu13', 'cupy-cuda13x']}

def driver_cuda() -> tuple:

(major, detail) for the CUDA the driver serves, or (0, why not).

Read from nvidia-smi, which reports the highest CUDA a driver supports -- not from nvcc, which reports whatever toolkit happens to be installed and is frequently older or absent. The wheels are built against a CUDA major version and load against the driver, so the driver is the one that decides.

def torch_cuda() -> int:

The CUDA major version torch was built against, or 0 -- another vote on which set to install.

Worth asking because mixing CUDA majors in one environment is how two libraries end up loading two runtimes: this machine has torch built for cu12, and installing cuml for cu13 beside it is asking for it.

def plan(prefer: int = 0) -> dict:

What would be installed, and why -- decided before anything is downloaded.

The newest set the driver could run is not the right answer. NVIDIA drivers are backward compatible, so a driver serving CUDA 13 runs CUDA 12 wheels perfectly, while cu13 wheels need a CUDA 13 driver and are the newer, less-travelled build. So the rule is: take what torch in this environment was built against if it is installed and supported, else the lowest set the driver can run -- which is cu12 on everything current. --cuda N overrides, for a machine that wants the newer set on purpose.

def command(packages) -> list:

The exact pip command, so it can be shown before it is run.

def main(argv=None) -> int:

starplast-install-gpu: report the plan, then run it unless asked not to.