feat: build and publish Nemotron OCR wheels for Python 3.11, 3.12, and 3.13

.github/workflows/huggingface-nightly.yml

@@ -168,8 +168,11 @@ jobs:
     needs: nightly_coordinate
     if: ${{ github.event_name != 'workflow_dispatch' || inputs.package == 'all' || inputs.package == 'nemotron-ocr-v2' }}
     # Nemotron OCR packages need nvcc/CUDA headers to build their extension.
-    # Build with Python 3.12 to match upstream package constraints and
-    # avoid producing an extension for the wrong Python ABI.
+    # Build a per-interpreter extension for each supported Python (3.11, 3.12,
+    # 3.13): torch cpp_extension compiles against the active interpreter's ABI,
+    # so setup-python selects both the build interpreter and the wheel ABI tag.
+    # We relax the upstream requires-python (see --set-requires-python) so the
+    # resulting wheels stay pip-installable on every targeted interpreter.
     # Matrix across x86_64 and aarch64 so the wheel works on ARM hosts
     # (e.g. DGX Spark) as well as conventional x86 machines.
     runs-on: ${{ matrix.platform.runner }}
@@ -183,6 +186,13 @@ jobs:
           - runner: ubuntu-24.04-arm
             cuda_image: nvidia/cuda:13.0.0-devel-ubuntu24.04
             arch: aarch64
+        python:
+          - version: "3.11"
+            tag: cp311
+          - version: "3.12"
+            tag: cp312
+          - version: "3.13"
+            tag: cp313
         ocr:
           - id: nemotron-ocr-v2
             url: https://huggingface.co/nvidia/nemotron-ocr-v2
@@ -194,8 +204,34 @@ jobs:
         shell: bash
         run: |
           set -euo pipefail
-          apt-get update
-          apt-get install -y --no-install-recommends \
+
+          # GitHub arm runners often cannot reach Ubuntu HTTP mirrors (port 80).
+          # Rewrite apt sources to HTTPS before the first update.
+          for apt_list in /etc/apt/sources.list /etc/apt/sources.list.d/*.list /etc/apt/sources.list.d/*.sources; do
+            [ -f "${apt_list}" ] || continue
+            sed -i \
+              -e 's|http://ports.ubuntu.com|https://ports.ubuntu.com|g' \
+              -e 's|http://archive.ubuntu.com|https://archive.ubuntu.com|g' \
+              -e 's|http://security.ubuntu.com|https://security.ubuntu.com|g' \
+              "${apt_list}"
+          done
+
+          apt_retry() {
+            local attempt=1
+            local max_attempts=5
+            until apt-get "$@"; do
+              if (( attempt >= max_attempts )); then
+                return 1
+              fi
+              attempt=$((attempt + 1))
+              echo "apt-get $* failed (attempt ${attempt}/${max_attempts}); retrying..."
+              sleep $((attempt * 5))
+              apt-get update || true
+            done
+          }
+
+          apt_retry update
+          apt_retry install -y --no-install-recommends \
             ca-certificates \
             git \
             git-lfs \
@@ -213,7 +249,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: "3.12"
+          python-version: "${{ matrix.python.version }}"
 
       - name: Decide upload target
         id: target
@@ -230,7 +266,7 @@ jobs:
           echo "upload_to=${upload_to}" >> "$GITHUB_OUTPUT"
           echo "skip_existing=${skip_existing}" >> "$GITHUB_OUTPUT"
 
-      - name: Build ${{ matrix.ocr.id }} (and maybe upload)
+      - name: Build ${{ matrix.ocr.id }} (and maybe upload) for Python ${{ matrix.python.version }}
         env:
           NIGHTLY_DATE_SUFFIX: ${{ needs.nightly_coordinate.outputs.nightly_date_suffix }}
           PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
@@ -267,6 +303,14 @@ jobs:
             arch_env_arg="--build-env ARCH=arm64"
           fi
 
+          # The sdist is interpreter-agnostic and identical across every leg of
+          # the platform/Python matrix, so build (and upload) it exactly once
+          # (x86_64 + 3.12); every other leg produces only its ABI-tagged wheel.
+          sdist_arg="--skip-sdist"
+          if [[ "${{ matrix.platform.arch }}" == "x86_64" && "${{ matrix.python.version }}" == "3.12" ]]; then
+            sdist_arg=""
+          fi
+
           release_type="nightly"
           expected_version="${{ matrix.ocr.nightly_base_version }}.dev"
           version_args=(--nightly-base-version "${{ matrix.ocr.nightly_base_version }}")
@@ -285,6 +329,8 @@ jobs:
             --dist-dir "dist-out" \
             --project-subdir "nemotron-ocr" \
             "${version_args[@]}" \
+            --set-requires-python ">=3.11,<3.14" \
+            ${sdist_arg} \
             --hatch-force-platform-wheel \
             --auditwheel-repair \
             --auditwheel-exclude libtorch_cpu.so \
@@ -338,6 +384,8 @@ jobs:
               "torch": f"~={os.environ['OCR_TORCH_VERSION']}",
               "torchvision": f"~={os.environ['OCR_TORCHVISION_VERSION']}",
           }
+          expected_license = "Apache-2.0"
+          expected_license_classifier = "License :: OSI Approved :: Apache Software License"
 
           def expanded_compatible_specifier(specifier):
               if not specifier.startswith("~="):
@@ -373,6 +421,18 @@ jobs:
               if f"Name: {expected_project}" not in metadata:
                   raise SystemExit(f"Built wheel metadata does not declare project {expected_project!r}")
               metadata_version = metadata_message.get("Version")
+              metadata_license = metadata_message.get("License")
+              if metadata_license != expected_license:
+                  raise SystemExit(
+                      "Built wheel metadata does not declare expected license "
+                      f"{expected_license!r}; got {metadata_license!r}"
+                  )
+              metadata_classifiers = metadata_message.get_all("Classifier", [])
+              if expected_license_classifier not in metadata_classifiers:
+                  raise SystemExit(
+                      "Built wheel metadata does not declare expected license classifier "
+                      f"{expected_license_classifier!r}; got {metadata_classifiers!r}"
+                  )
               if release_type == "stable" and metadata_version != expected_version:
                   raise SystemExit(
                       "Built wheel metadata does not declare expected version "
@@ -427,5 +487,5 @@ jobs:
         if: always()
         uses: actions/upload-artifact@v4
         with:
-          name: dist-${{ matrix.ocr.id }}-${{ matrix.platform.arch }}
+          name: dist-${{ matrix.ocr.id }}-${{ matrix.platform.arch }}-${{ matrix.python.tag }}
           path: dist-out/${{ matrix.ocr.id }}/*

.gitignore

@@ -275,3 +275,7 @@ spool/
 
 # Local git worktrees
 .worktrees/
+
+# Local huggingface cache and build directories
+.work*/
+dist*

ci/scripts/build_ocr_wheels_local.sh

@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+set -euo pipefail
+export DEBIAN_FRONTEND=noninteractive
+apt-get update
+apt-get install -y --no-install-recommends \
+  ca-certificates git git-lfs build-essential ninja-build patchelf \
+  software-properties-common python3-pip
+add-apt-repository -y ppa:deadsnakes/ppa
+apt-get update
+apt-get install -y --no-install-recommends \
+  python3.11 python3.11-dev python3.11-venv \
+  python3.13 python3.13-dev python3.13-venv
+git lfs install
+arch_env_arg=""
+if [[ "$(uname -m)" == "aarch64" ]]; then
+  arch_env_arg="--build-env ARCH=arm64"
+fi
+mkdir -p dist-out/nemotron-ocr-v2
+for PY_VER in 3.11 3.13; do
+  PYTHON="python${PY_VER}"
+  TAG="cp${PY_VER//./}"
+  STAGING="dist-staging-${TAG}"
+  VENV_DIR="/tmp/.venv-build-${TAG}"
+  echo "=== Building nemotron-ocr for Python ${PY_VER} ==="
+  "${PYTHON}" --version
+  "${PYTHON}" -m pip install --break-system-packages "packaging>=24"
+  "${PYTHON}" ci/scripts/nightly_build_publish.py \
+    --repo-id nemotron-ocr-v2 \
+    --repo-url https://huggingface.co/nvidia/nemotron-ocr-v2 \
+    --work-dir ".work-${TAG}" \
+    --dist-dir "${STAGING}" \
+    --venv-dir "${VENV_DIR}" \
+    --project-subdir nemotron-ocr \
+    --nightly-base-version 2.0.1 \
+    --set-requires-python ">=3.11,<3.14" \
+    --skip-sdist \
+    --hatch-force-platform-wheel \
+    --auditwheel-repair \
+    --auditwheel-exclude libtorch_cpu.so \
+    --auditwheel-exclude libtorch_cuda.so \
+    --auditwheel-exclude libtorch.so \
+    --auditwheel-exclude libc10.so \
+    --auditwheel-exclude libc10_cuda.so \
+    --auditwheel-exclude libtorch_python.so \
+    --build-no-isolation \
+    --venv-pip-install hatchling \
+    --venv-pip-install "setuptools>=68" \
+    --venv-pip-install ninja \
+    --venv-pip-install "torch==${OCR_TORCH_VERSION}" \
+    --venv-pip-install "torchvision==${OCR_TORCHVISION_VERSION}" \
+    --pin-runtime-dependency torch \
+    --pin-runtime-dependency torchvision \
+    --build-env BUILD_CPP_EXTENSION=1 \
+    --build-env BUILD_CPP_FORCE=1 \
+    ${arch_env_arg}
+  cp "${STAGING}/nemotron-ocr-v2/"*.whl dist-out/nemotron-ocr-v2/
+done
+ls -lh dist-out/nemotron-ocr-v2/