Report torch/torchvision versions and binary sizes for every validated build

.github/scripts/build_report.py

@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+"""Report the versions and sizes of the binaries a validation run installed.
+
+Emits a small markdown table that validate_binaries.sh prints to the job log and
+appends to the job summary. Values are read from the installed packages rather
+than from the build matrix, so a mismatch between what was requested and what
+pip resolved is visible in the report.
+
+Two sizes are reported because they answer different questions:
+
+* wheel     -- the compressed download size, i.e. what a user pulls from the
+               index. Only pip knows it, so the caller passes it in; it is
+               absent on the uv/wheel-variants path and when the wheel was
+               already satisfied.
+* installed -- unpacked bytes on disk, measured here from the installed
+               package, so it is available on every install path and every OS.
+
+Nothing in here may fail the validation job: every lookup that depends on how
+torch was built (CUDA, cuDNN, NCCL) degrades to "-" instead of raising.
+
+Note on imports: this is deliberately a script file rather than a heredoc piped
+to python. For a script, sys.path[0] is the script's own directory, so `import
+torch` cannot pick up the pytorch source checkout the validation runs from --
+which a `python -` heredoc would, since that puts the cwd on sys.path.
+
+Usage:
+  build_report.py --target-os linux --python-version 3.12 \
+      --gpu-arch-type cuda --gpu-arch-version 12.8 \
+      --torch-wheel-mb 812.4 --torchvision-wheel-mb 8.1
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from types import ModuleType
+
+
+def installed_size_mb(module: ModuleType) -> str:
+    """Total size of a package directory on disk, formatted as MB.
+
+    Files that vanish or cannot be stat'd mid-walk are skipped rather than
+    aborting the report.
+    """
+    path = getattr(module, "__file__", None)
+    if not path:
+        return "-"
+
+    total = 0
+    for dirpath, _, filenames in os.walk(os.path.dirname(path)):
+        for filename in filenames:
+            try:
+                total += os.path.getsize(os.path.join(dirpath, filename))
+            except OSError:
+                continue
+    return f"{total / 1024 / 1024:.1f} MB"
+
+
+def format_mb(value: str) -> str:
+    """Render a caller-supplied megabyte figure, or "-" when it is unknown."""
+    return f"{value} MB" if value else "-"
+
+
+def cudnn_version(torch: ModuleType) -> str:
+    """Decode torch's packed cuDNN version integer, e.g. 91002 -> 9.10.2."""
+    try:
+        packed = torch.backends.cudnn.version()
+    except Exception:
+        return "-"
+    if not packed:
+        return "-"
+    return f"{packed // 10000}.{packed % 10000 // 100}.{packed % 100}"
+
+
+def nccl_version(torch: ModuleType) -> str:
+    """Format torch's NCCL version tuple, e.g. (2, 30, 7) -> 2.30.7."""
+    try:
+        return ".".join(str(part) for part in torch.cuda.nccl.version())
+    except Exception:
+        return "-"
+
+
+def collect_rows(args: argparse.Namespace) -> list[tuple[str, str]]:
+    arch = args.gpu_arch_type
+    if args.gpu_arch_version:
+        arch = f"{arch} {args.gpu_arch_version}"
+    rows = [("build", f"{args.target_os} / py{args.python_version} / {arch}")]
+
+    try:
+        import torch
+    except Exception as exc:
+        rows.append(("torch", f"import failed: {exc}"))
+    else:
+        rows.extend(
+            [
+                ("torch", torch.__version__),
+                ("torch wheel", format_mb(args.torch_wheel_mb)),
+                ("torch installed", installed_size_mb(torch)),
+                ("CUDA", torch.version.cuda or "-"),
+                ("cuDNN", cudnn_version(torch)),
+                ("NCCL", nccl_version(torch)),
+            ]
+        )
+
+    try:
+        import torchvision
+    except Exception:
+        rows.append(("torchvision", "-"))
+    else:
+        rows.extend(
+            [
+                ("torchvision", torchvision.__version__),
+                ("torchvision wheel", format_mb(args.torchvision_wheel_mb)),
+                ("torchvision installed", installed_size_mb(torchvision)),
+            ]
+        )
+
+    return rows
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--target-os", default="?")
+    parser.add_argument("--python-version", default="?")
+    parser.add_argument("--gpu-arch-type", default="cpu")
+    parser.add_argument("--gpu-arch-version", default="")
+    parser.add_argument(
+        "--torch-wheel-mb",
+        default="",
+        help="Compressed torch wheel size in MB, as read from pip's output",
+    )
+    parser.add_argument(
+        "--torchvision-wheel-mb",
+        default="",
+        help="Compressed torchvision wheel size in MB, as read from pip's output",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    rows = collect_rows(parse_args())
+    print("| field | value |")
+    print("| --- | --- |")
+    for field, value in rows:
+        print(f"| {field} | {value} |")
+
+
+if __name__ == "__main__":
+    main()

.github/scripts/validate_binaries.sh

@@ -247,6 +247,29 @@ cleanup_conda_env() {
     fi
 }
 
+# Read a wheel's compressed download size, in MB, out of a captured pip log.
+#
+# $1 = log file, $2 = distribution name as it appears in the wheel filename.
+# Prints the size, or nothing when that wheel is absent from the log (already
+# satisfied, or installed from a local file). "<dist>-[0-9]" keeps a request for
+# "torch" from matching the torchvision-/torchaudio- wheels.
+parse_wheel_size_mb() {
+    local log_file="$1" dist="$2" frag size unit
+    frag=$(grep -oiE "${dist}-[0-9][^ /]*\.whl \([0-9.]+ ?[kKmMgG]i?B\)" "${log_file}" | tail -1 || true)
+    if [[ -z ${frag} ]]; then
+        return 0
+    fi
+    size=$(echo "${frag}" | sed -E 's/.*\(([0-9.]+) ?([A-Za-z]+)\)$/\1/')
+    unit=$(echo "${frag}" | sed -E 's/.*\(([0-9.]+) ?([A-Za-z]+)\)$/\2/')
+    case ${unit} in
+        B)             awk "BEGIN{printf \"%.1f\", ${size}/1024/1024}" ;;
+        kB|KB|kiB|KiB) awk "BEGIN{printf \"%.1f\", ${size}/1024}" ;;
+        MB|MiB)        awk "BEGIN{printf \"%.1f\", ${size}}" ;;
+        GB|GiB)        awk "BEGIN{printf \"%.1f\", ${size}*1024}" ;;
+        *) echo "::warning::wheel-size: unrecognized size unit '${unit}' for ${dist}" >&2 ;;
+    esac
+}
+
 # Fail the build if the installed torch wheel exceeds a hard size ceiling.
 #
 # Scope: Linux x86_64 + aarch64 wheels only, excluding ROCm (whose wheels are
@@ -269,27 +292,13 @@ check_wheel_size() {
         return 0
     fi
 
-    # Pull the torch wheel's size off pip's Downloading/Using-cached line, e.g.
-    #   Downloading torch-2.10.0.dev...-linux_x86_64.whl (812.4 MB)
-    # torch-[0-9] isolates the torch wheel from torchvision-/torchaudio-.
-    local frag
-    frag=$(grep -oiE "torch-[0-9][^ /]*\.whl \([0-9.]+ ?[kKmMgG]i?B\)" "${log_file}" | tail -1 || true)
-    if [[ -z ${frag} ]]; then
+    local size_mb
+    size_mb=$(parse_wheel_size_mb "${log_file}" torch)
+    if [[ -z ${size_mb} ]]; then
         echo "::warning::wheel-size check: could not find the torch wheel size in the pip output; skipping"
         return 0
     fi
 
-    local size unit size_mb
-    size=$(echo "${frag}" | sed -E 's/.*\(([0-9.]+) ?([A-Za-z]+)\)$/\1/')
-    unit=$(echo "${frag}" | sed -E 's/.*\(([0-9.]+) ?([A-Za-z]+)\)$/\2/')
-    case ${unit} in
-        B)             size_mb=$(awk "BEGIN{printf \"%.1f\", ${size}/1024/1024}") ;;
-        kB|KB|kiB|KiB) size_mb=$(awk "BEGIN{printf \"%.1f\", ${size}/1024}") ;;
-        MB|MiB)        size_mb=$(awk "BEGIN{printf \"%.1f\", ${size}}") ;;
-        GB|GiB)        size_mb=$(awk "BEGIN{printf \"%.1f\", ${size}*1024}") ;;
-        *) echo "::warning::wheel-size check: unrecognized size unit '${unit}'; skipping"; return 0 ;;
-    esac
-
     # Always surface the measured size (as an annotation) whether or not the
     # check passes, so it is visible on the run summary of a successful job too.
     echo "::notice::torch wheel size: ${size_mb} MB (arch=${MATRIX_GPU_ARCH_TYPE:-cpu} os=${TARGET_OS} py=${MATRIX_PYTHON_VERSION:-?}); ceiling ${threshold_mb} MB"
@@ -299,6 +308,44 @@ check_wheel_size() {
     fi
 }
 
+# Report what this build actually installed, and how big it is.
+#
+# Complements check_wheel_size, which only measures linux/linux-aarch64 pip
+# wheels and exists to enforce a ceiling: this runs for every build on every OS
+# and never fails, so windows/macos/ROCm sizes are visible too.
+#
+# The report itself is built by build_report.py; see that script for what the
+# two size figures mean. Runs while the env is still active, because
+# cleanup_conda_env removes it on non-linux.
+write_build_report() {
+    local torch_wheel_mb="${1:-}" vision_wheel_mb="${2:-}"
+    local report
+
+    report=$("${PYTHON_RUN}" "${SCRIPT_DIR}/build_report.py" \
+        --target-os "${TARGET_OS}" \
+        --python-version "${MATRIX_PYTHON_VERSION:-?}" \
+        --gpu-arch-type "${MATRIX_GPU_ARCH_TYPE:-cpu}" \
+        --gpu-arch-version "${MATRIX_GPU_ARCH_VERSION:-}" \
+        --torch-wheel-mb "${torch_wheel_mb}" \
+        --torchvision-wheel-mb "${vision_wheel_mb}") || report="| field | value |
+| --- | --- |
+| report | failed to collect |"
+
+    echo "--- Build report"
+    echo "${report}"
+
+    # The job summary file is not reachable from inside the validation
+    # container, so only append when the runner actually exposes a writable one.
+    # stderr is redirected before the append so a non-writable path fails quietly
+    if [[ -n ${GITHUB_STEP_SUMMARY:-} ]] && : 2>/dev/null >>"${GITHUB_STEP_SUMMARY}"; then
+        {
+            echo "### ${MATRIX_PACKAGE_TYPE:-wheel}: ${TARGET_OS} / py${MATRIX_PYTHON_VERSION:-?} / ${MATRIX_GPU_ARCH_TYPE:-cpu} ${MATRIX_GPU_ARCH_VERSION:-}"
+            echo "${report}"
+            echo
+        } >> "${GITHUB_STEP_SUMMARY}"
+    fi
+}
+
 #######################################
 # Main Script
 #######################################
@@ -393,6 +440,8 @@ if [[ ${MATRIX_PACKAGE_TYPE} == 'wheel' ]]; then
 fi
 
 # Install packages
+TORCH_WHEEL_MB=""
+TORCHVISION_WHEEL_MB=""
 if [[ ${USE_WHEEL_VARIANTS:-} == 'true' ]]; then
     install_wheel_variants
 else
@@ -404,6 +453,8 @@ else
     WHEEL_INSTALL_LOG="$(mktemp)"
     eval "${INSTALLATION}" 2>&1 | tee "${WHEEL_INSTALL_LOG}"
     check_wheel_size "${WHEEL_INSTALL_LOG}"
+    TORCH_WHEEL_MB="$(parse_wheel_size_mb "${WHEEL_INSTALL_LOG}" torch)"
+    TORCHVISION_WHEEL_MB="$(parse_wheel_size_mb "${WHEEL_INSTALL_LOG}" torchvision)"
     rm -f "${WHEEL_INSTALL_LOG}"
 fi
 
@@ -413,6 +464,9 @@ install_numpy_1x
 # Run tests
 run_smoke_tests "${TEST_SUFFIX}"
 
+# Report versions and sizes for this build
+write_build_report "${TORCH_WHEEL_MB}" "${TORCHVISION_WHEEL_MB}"
+
 # Restore PATH for macos-arm64
 if [[ ${TARGET_OS} == 'macos-arm64' ]]; then
     export PATH=${OLD_PATH}

mypy.ini

@@ -70,4 +70,12 @@ disable_error_code = var-annotated, import-untyped
 disable_error_code = index, no-redef, no-untyped-def, import-not-found, no-untyped-call
 
 [mypy-tools.rockset_migration.*]
-disable_error_code = arg-type, attr-defined, return-value, misc, func-returns-value, no-untyped-def, assignment, var-annotated, unused-ignore, truthy-function
+disable_error_code = arg-type, attr-defined, return-value, misc, func-returns-value, no-untyped-def, assignment, var-annotated, unused-ignore, truthy-function
+
+# torch and torchvision are not installed in the lint environment. Scripts that
+# import them do so defensively at runtime and degrade when they are absent.
+[mypy-torch.*]
+ignore_missing_imports = True
+
+[mypy-torchvision.*]
+ignore_missing_imports = True