s3_management: source PyPI-hosted NVIDIA packages from pypi.org (#8473)

get_package_source_url() routed any package whose name starts with
"nvidia-" or "cuda-" to pypi.nvidia.com. A handful of NVIDIA-named
packages are not published there -- cuda-bindings, cuda-pathfinder,
cuda-python and nvidia-ml-py live on PyPI proper -- so fetching their
simple index 404s. The failure is quiet:
upload_package_using_simple_index() catches the fetch error and returns,
leaving the mirrored index frozen at whatever was last seeded.

The prefix predicate carries two meanings that had to be separated
before the exception could be expressed. is_nvidia_package() also gates
which packages are initialized for CUDA targets only
(get_packages_for_target), so excluding these four inside it would have
leaked them into cpu and rocm targets. It keeps the plain prefix
heuristic; the new uses_nvidia_index() layers the exception on top and
is what the URL builder and the "NVIDIA"/"PyPI" source label now
consult, so the printed label always matches the index actually fetched.

Test Plan:

The four listed packages resolve to pypi.org/simple/, nvidia-cublas-cu12
and cuda-runtime-cu13 still resolve to pypi.nvidia.com, numpy to
pypi.org and pytorch-triton-rocm to repo.amd.com. is_nvidia_package()
returns True for all six NVIDIA names, confirming the CUDA-target gate
is unchanged.

cc @atalman

Author: tinglvv

s3_management/update_dependencies.py

@@ -882,11 +882,27 @@
 PREVIEW_NUMPY_VERSION = "2.6.0"
 
 
+# NVIDIA packages that are published to PyPI rather than pypi.nvidia.com. The
+# nvidia-/cuda- prefix heuristic would send these to a 404 and silently freeze
+# their mirrored index at whatever was last seeded.
+PYPI_HOSTED_NVIDIA_PACKAGES = {
+    "cuda-bindings",
+    "cuda-pathfinder",
+    "cuda-python",
+    "nvidia-ml-py",
+}
+
+
 def is_nvidia_package(pkg_name: str) -> bool:
-    """Check if a package is from NVIDIA and should use pypi.nvidia.com"""
+    """Check if a package is from NVIDIA and is therefore CUDA-target only"""
     return pkg_name.startswith("nvidia-") or pkg_name.startswith("cuda-")
 
 
+def uses_nvidia_index(pkg_name: str) -> bool:
+    """Check if a package is mirrored from pypi.nvidia.com rather than PyPI"""
+    return is_nvidia_package(pkg_name) and pkg_name not in PYPI_HOSTED_NVIDIA_PACKAGES
+
+
 def is_amd_package(pkg_name: str) -> bool:
     """Check if a package is from AMD and should use repo.amd.com"""
     name = pkg_name.lower()
@@ -895,7 +911,7 @@ def is_amd_package(pkg_name: str) -> bool:
 
 def get_package_source_url(pkg_name: str) -> str:
     """Get the source URL for a package based on its type"""
-    if is_nvidia_package(pkg_name):
+    if uses_nvidia_index(pkg_name):
         return f"https://pypi.nvidia.com/{pkg_name}/"
     if is_amd_package(pkg_name):
         return f"https://repo.amd.com/rocm/whl-multi-arch/{pkg_name}/"
@@ -1106,7 +1122,7 @@ def upload_package_using_simple_index(
     Works for both NVIDIA and non-NVIDIA packages.
     """
     source_url = get_package_source_url(pkg_name)
-    if is_nvidia_package(pkg_name):
+    if uses_nvidia_index(pkg_name):
         source_label = "NVIDIA"
     elif is_amd_package(pkg_name):
         source_label = "AMD"