Re-add nvidia-cudnn-cu12 9.1.0.70 to the mirrored cudnn index

Add a RETAINED_WHEELS mechanism to re-inject wheels we still host that the
upstream simple index has stopped publishing, and use it for
nvidia-cudnn-cu12 9.1.0.70 on the root, cu121 and cu124 indexes across
nightly, test and prod.

Author: atalman

s3_management/update_dependencies.py

@@ -1,7 +1,7 @@
 import os
 import re
 import time
-from typing import Dict, List
+from typing import Dict, List, NamedTuple
 from urllib.parse import urljoin
 
 import boto3  # type: ignore[import-untyped]
@@ -227,6 +227,14 @@
         {
             "project": "torch",
         },
+        {
+            "project": "torch",
+            "target": "cu121",
+        },
+        {
+            "project": "torch",
+            "target": "cu124",
+        },
         {
             "project": "torch",
             "target": "cu126",
@@ -919,6 +927,32 @@
 }
 
 
+class RetainedWheel(NamedTuple):
+    filename: str
+    url: str
+    sha256: str
+    targets: tuple[str, ...]
+
+
+# Wheels that download.pytorch.org still hosts but the upstream simple index has
+# stopped publishing, so mirroring alone drops them and pinned installs stop
+# resolving. Keyed by normalised package name; injected into the index for the
+# listed targets ("" is the channel root) on nightly, test and prod.
+RETAINED_WHEELS: Dict[str, List[RetainedWheel]] = {
+    "nvidia_cudnn_cu12": [
+        RetainedWheel(
+            filename="nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl",
+            url=(
+                "https://download.pytorch.org/whl/cu124/"
+                "nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl"
+            ),
+            sha256="165764f44ef8c61fcdfdfdbe769d687e06374059fbb388b6c89ecb0e28793a6f",
+            targets=("", "cu121", "cu124"),
+        ),
+    ],
+}
+
+
 def is_nvidia_package(pkg_name: str) -> bool:
     """Check if a package is from NVIDIA and is therefore CUDA-target only"""
     return pkg_name.startswith("nvidia-") or pkg_name.startswith("cuda-")
@@ -1136,6 +1170,39 @@ def append_preview_numpy_wheels(html: str, pkg_name: str, prefix: str) -> str:
     return f"{html}\n{block}\n"
 
 
+def append_retained_wheels(html: str, pkg_name: str, prefix: str) -> str:
+    """Merge :data:`RETAINED_WHEELS` links for *pkg_name* into *html*.
+
+    The target is the last segment of *prefix* ("whl/nightly/cu124" -> "cu124",
+    "whl/nightly" -> ""). Wheels the upstream index still lists are skipped, so
+    this becomes a no-op if a version is ever restored upstream.
+    """
+    entries = RETAINED_WHEELS.get(normalize_pkg_name(pkg_name))
+    if not entries:
+        return html
+
+    last_segment = prefix.rsplit("/", 1)[-1]
+    target = last_segment if is_valid_target(last_segment) else ""
+
+    additions = [
+        f'    <a href="{entry.url}#sha256={entry.sha256}">{entry.filename}</a><br/>'
+        for entry in entries
+        if target in entry.targets and entry.filename not in html
+    ]
+
+    if not additions:
+        return html
+
+    print(
+        f"INFO: Merging {len(additions)} retained wheel link(s) "
+        f"for {pkg_name} under {prefix}"
+    )
+    block = "\n".join(additions)
+    if "</body>" in html:
+        return html.replace("</body>", f"{block}\n  </body>", 1)
+    return f"{html}\n{block}\n"
+
+
 def upload_package_using_simple_index(
     pkg_name: str,
     prefix: str,
@@ -1168,6 +1235,9 @@ def upload_package_using_simple_index(
     # published on the scientific-python nightly wheelhouse so pip can resolve them.
     raw_html = append_preview_numpy_wheels(raw_html, pkg_name, prefix)
 
+    # Re-add wheels we still host that upstream has stopped publishing.
+    raw_html = append_retained_wheels(raw_html, pkg_name, prefix)
+
     # Upload modified index.html with absolute links
     upload_index_html(pkg_name, prefix, raw_html, source_url, dry_run=dry_run)