Align PyTorch CUDA wheel versions with Dockerfile CUDA 12.6 in pytorch+llmcompressor runtime

[coderabbitai]

Problem Description

The runtimes/pytorch+llmcompressor/ubi9-python-3.12/Dockerfile.cuda builds against CUDA 12.6, but the corresponding pylock.toml file references PyTorch wheels compiled for CUDA 12.4 (+cu124). This version mismatch can cause runtime incompatibilities and suboptimal GPU performance.

Affected Files

• runtimes/pytorch+llmcompressor/ubi9-python-3.12/Dockerfile.cuda (uses CUDA 12.6)
• runtimes/pytorch+llmcompressor/ubi9-python-3.12/pylock.toml (references +cu124 wheels)

Current State

Dockerfile CUDA version: 12.6

ENV CUDA_VERSION=12.6.3
ENV NV_CUDA_CUDART_VERSION=12.6.77-1

PyTorch wheel versions in pylock.toml: cu124

[[packages]]
name = "torch"
version = "2.6.0+cu124"

Solution

1. Update pyproject.toml PyTorch index

Update the PyTorch index URL in runtimes/pytorch+llmcompressor/ubi9-python-3.12/pyproject.toml:

[[tool.uv.index]]
name = "pytorch-cuda"
url = "https://download.pytorch.org/whl/cu126"  # Change from cu124
explicit = true

2. Regenerate lock file

Use uv to regenerate the lock file with CUDA 12.6 wheels:

cd runtimes/pytorch+llmcompressor/ubi9-python-3.12
uv lock --upgrade-package torch --upgrade-package torchvision

3. Verify alignment

Ensure pylock.toml contains +cu126 versions:

rg -n '(\+cu126)|download\.pytorch\.org/whl/cu126' pylock.toml

Acceptance Criteria

• pylock.toml references torch==2.6.0+cu126 and torchvision==0.21.0+cu126
• PyTorch index URL points to https://download.pytorch.org/whl/cu126
• Lock file contains correct wheel URLs and hashes for cu126 variants
• Container builds successfully with aligned CUDA versions
• PyTorch CUDA functionality works correctly in runtime environment

Context

• PR: RHAIENG-304: add uv pyproject.tomls for multiple notebooks and runtimes across Jupyter, RStudio, and VSCode #2145
• Review Comment: RHAIENG-304: add uv pyproject.tomls for multiple notebooks and runtimes across Jupyter, RStudio, and VSCode #2145 (comment)
• Requesting User: @jiridanek

This issue ensures optimal GPU performance and prevents potential runtime compatibility issues between CUDA runtime and PyTorch wheels.