Torch-NPU主干和上游社区集成测试

.ci/docker/requirements-ci.txt

@@ -1,7 +1,51 @@
-# Python dependencies required for unit tests
+# Python dependencies required for NPU tests
+# Based on upstream PyTorch .ci/docker/requirements-ci.txt
 
-mypy==1.9.0
-# Pin MyPy version because new errors are likely to appear with each release
-#Description: linter
-#Pinned versions: 1.9.0
-#test that import: test_typing.py, test_type_hints.py
+# pytest and plugins
+pytest==7.3.2
+pytest-xdist==3.3.1
+pytest-flakefinder==1.1.0
+pytest-rerunfailures>=10.3
+pytest-subtests==0.13.1
+pytest-timeout>=2.3.1
+xdoctest==1.3.0
+
+# test utilities
+hypothesis==6.56.4
+expecttest==0.3.0
+parameterized==0.8.1
+
+# numpy (version per Python version)
+numpy==1.26.2; python_version >= "3.11" and python_version < "3.14"
+
+# scientific packages
+scipy==1.14.1; python_version > "3.11" and python_version < "3.14"
+scikit-image==0.22.0
+pillow==12.1.1
+pywavelets==1.7.0; python_version >= "3.12"
+
+# core utilities
+networkx==2.8.8
+optree==0.13.0; python_version < "3.14"
+opt-einsum==3.3
+filelock==3.20.3
+sympy==1.13.3
+
+# build/serialization
+pyyaml==6.0.3
+packaging==24.0
+typing-extensions==4.12.2; python_version < "3.14"
+pyzstd
+setuptools>=70.1.0,<82
+zstandard
+
+# ONNX support
+onnx==1.20.0
+onnxscript==0.6.2
+protobuf==6.33.5
+
+# misc
+psutil
+jinja2==3.1.6
+tqdm>=4.66.0
+click

.github/actions/setup-npu-test-env/action.yml

@@ -0,0 +1,146 @@
+name: 'Setup NPU Test Environment'
+description: 'Common environment setup for NPU upstream tests - checkout, cache, install PyTorch/torch_npu/triton-ascend, test dependencies'
+
+inputs:
+  python_version:
+    required: true
+    type: string
+    description: Python version to use
+  torch_wheel_artifact:
+    required: true
+    type: string
+    description: Name of the torch wheel artifact
+  torch_npu_wheel_artifact:
+    required: true
+    type: string
+    description: Name of the torch_npu wheel artifact
+  pytorch_src_artifact:
+    required: true
+    type: string
+    description: Name of the PyTorch source artifact
+
+env:
+  # PyPI 缓存 URL（用于加速 pip 下载）
+  PYPI_CACHE_URL: 'http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple'
+
+runs:
+  using: 'composite'
+  steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          repository: ${{ github.repository }}
+          ref: ${{ github.ref }}
+          fetch-depth: 1
+          path: ascend_pytorch
+
+      - name: Setup cache directories
+        run: |
+          mkdir -p /github/home/.cache/pip
+          chmod -R 777 /github/home/.cache
+
+      - name: Cache pip
+        uses: actions/cache@v4
+        with:
+          path: /github/home/.cache/pip
+          key: pip-arm-collect-py${{ inputs.python_version }}
+          restore-keys: |
+            pip-arm-collect-py${{ inputs.python_version }}-
+            pip-arm-collect-
+
+      - name: Download built torch wheel
+        uses: actions/download-artifact@v4
+        with:
+          name: ${{ inputs.torch_wheel_artifact }}
+          path: torch-wheel-artifact
+
+      - name: Download built torch_npu wheel
+        uses: actions/download-artifact@v4
+        with:
+          name: ${{ inputs.torch_npu_wheel_artifact }}
+          path: torch-npu-wheel-artifact
+
+      - name: Download PyTorch source and test code
+        uses: actions/download-artifact@v4
+        with:
+          name: ${{ inputs.pytorch_src_artifact }}
+          path: pytorch-src-artifact
+
+      - name: Extract PyTorch source
+        run: |
+          tar -xzf pytorch-src-artifact/pytorch-src.tar.gz
+
+      - name: Install built PyTorch and torch_npu
+        run: |
+          source /usr/local/Ascend/cann/set_env.sh 2>/dev/null || true
+          source /usr/local/Ascend/nnal/atb/set_env.sh 2>/dev/null || true
+
+          PIP=pip${{ inputs.python_version }}
+          PYTHON=python${{ inputs.python_version }}
+          export PIP_CACHE_DIR=/github/home/.cache/pip
+
+          # Configure pip to use PyPI cache for faster downloads
+          if [ -n "${{ env.PYPI_CACHE_URL }}" ]; then
+            $PIP config set global.index-url ${{ env.PYPI_CACHE_URL }}
+            $PIP config set global.trusted-host "cache-service.nginx-pypi-cache.svc.cluster.local"
+            echo "pip index-url configured: ${{ env.PYPI_CACHE_URL }}"
+          fi
+
+          $PIP install --upgrade pip
+
+          # Install built torch wheel
+          TORCH_WHL=$(ls torch-wheel-artifact/*.whl | head -1)
+          $PIP install "${TORCH_WHL}"
+
+          # Install built torch_npu wheel
+          TORCH_NPU_WHL=$(ls torch-npu-wheel-artifact/*.whl | head -1)
+          $PIP install "${TORCH_NPU_WHL}"
+
+          echo "Installed PyTorch and torch_npu from built wheels"
+          echo "torch: ${TORCH_WHL}"
+          echo "torch_npu: ${TORCH_NPU_WHL}"
+
+      - name: Install test dependencies
+        run: |
+          PIP=pip${{ inputs.python_version }}
+          export PIP_CACHE_DIR=/github/home/.cache/pip
+          cd pytorch-src
+
+          # Core test dependencies
+          $PIP install pytest pytest-timeout pytest-xdist hypothesis zstandard pyyaml
+          $PIP install pytest-rerunfailures pytest-flakefinder
+          $PIP install 'pytest-subtests==0.13.1' 'xdoctest==1.1.0' 'pulp>=2.9'
+
+          # Optional dependencies for ONNX tests
+          # These are not in PyTorch requirements.txt but needed by specific tests
+          $PIP install onnxruntime onnxscript onnx-ir ml-dtypes || true
+
+          # torchvision for ONNX model tests (install without deps to bypass torch version check)
+          # PyPI torchvision requires exact torch version (torch==2.11.0), but we have dev build
+          # Use --no-deps to skip torch dependency, we already have our compiled torch installed
+          $PIP install numpy pillow || true
+          $PIP install torchvision --no-deps || true
+
+          # Other optional dependencies
+          $PIP install parameterized pandas || true
+          $PIP install opencv-python || true
+
+          # PyTorch requirements (if exists)
+          if [ -f requirements.txt ]; then
+            $PIP install -r requirements.txt || true
+          fi
+
+      - name: Verify NPU availability
+        run: |
+          source /usr/local/Ascend/cann/set_env.sh 2>/dev/null || true
+          source /usr/local/Ascend/nnal/atb/set_env.sh 2>/dev/null || true
+
+          PYTHON=python${{ inputs.python_version }}
+          $PYTHON -c "
+          import torch
+          print(f'torch: {torch.__version__}')
+          import torch_npu
+          print(f'torch_npu: {torch_npu.__version__}')
+          print(f'NPU available: {torch.npu.is_available()}')
+          print(f'NPU count: {torch.npu.device_count()}')
+          "

.github/docker/pytorch-npu-builder.Dockerfile

@@ -0,0 +1,152 @@
+# 基于 PyPA manylinux 2_28 aarch64 镜像 (与 PyTorch 主干一致)
+FROM quay.io/pypa/manylinux_2_28_aarch64
+
+ARG GCCTOOLSET_VERSION=13
+
+# CANN 包下载 URL（通过 build-arg 传入）
+ARG CANN_TOOLKIT_URL
+ARG CANN_A3OPS_URL
+ARG CANN_NNAL_URL
+ARG CANN_VERSION
+
+# Language variables
+ENV LC_ALL=en_US.UTF-8
+ENV LANG=en_US.UTF-8
+ENV LANGUAGE=en_US.UTF-8
+
+# 安装必要的 OS 包 (与 PyTorch 官方 Dockerfile 一致)
+RUN yum -y install epel-release && \
+    yum -y update && \
+    yum install -y \
+        autoconf \
+        automake \
+        bison \
+        bzip2 \
+        curl \
+        diffutils \
+        file \
+        git \
+        less \
+        libffi-devel \
+        libgomp \
+        make \
+        openssl-devel \
+        patch \
+        perl \
+        unzip \
+        util-linux \
+        wget \
+        which \
+        xz \
+        yasm \
+        zstd \
+        sudo \
+        gcc-toolset-${GCCTOOLSET_VERSION}-gcc \
+        gcc-toolset-${GCCTOOLSET_VERSION}-gcc-c++ \
+        gcc-toolset-${GCCTOOLSET_VERSION}-gcc-gfortran \
+        gcc-toolset-${GCCTOOLSET_VERSION}-gdb && \
+    yum install -y --enablerepo=powertools ninja-build && \
+    rm -rf /var/cache/yum
+
+# 确保使用正确的 devtoolset
+ENV PATH=/opt/rh/gcc-toolset-${GCCTOOLSET_VERSION}/root/usr/bin:$PATH
+ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${GCCTOOLSET_VERSION}/root/usr/lib64:/opt/rh/gcc-toolset-${GCCTOOLSET_VERSION}/root/usr/lib:$LD_LIBRARY_PATH
+
+# git 2.36+ 需要配置 safe.directory
+RUN git config --global --add safe.directory "*"
+
+# ============================================================
+# 预装所有 Python 版本（镜像支持多 Python 版本）
+# ============================================================
+# manylinux 镜像已包含 cp310-cp310, cp311-cp311, cp312-cp312, cp313-cp313
+# 默认使用 Python 3.11（可通过环境变量切换）
+
+ENV DEFAULT_PYTHON_VERSION=3.11
+ENV PATH=/opt/python/cp311-cp311/bin:$PATH
+
+# 创建 Python 版本切换脚本
+RUN printf '#!/bin/bash\n\
+# Python 版本切换辅助脚本\n\
+# 使用方法: source /usr/local/bin/switch_python.sh 3.11\n\
+\n\
+PYTHON_VERSION="${1:-3.11}"\n\
+\n\
+case "$PYTHON_VERSION" in\n\
+    3.10) PYTHON_DIR="cp310-cp310" ;;\n\
+    3.11) PYTHON_DIR="cp311-cp311" ;;\n\
+    3.12) PYTHON_DIR="cp312-cp312" ;;\n\
+    3.13) PYTHON_DIR="cp313-cp313" ;;\n\
+    *) echo "Unsupported Python version: $PYTHON_VERSION"; return 1 ;;\n\
+esac\n\
+\n\
+export PATH=/opt/python/$PYTHON_DIR/bin:$PATH\n\
+echo "Switched to Python $PYTHON_VERSION ($(python --version))"\n\
+' > /usr/local/bin/switch_python.sh && \
+    chmod +x /usr/local/bin/switch_python.sh
+
+# 为每个 Python 版本安装常用包
+RUN for py_dir in cp310-cp310 cp311-cp311 cp312-cp312 cp313-cp313; do \
+        /opt/python/$py_dir/bin/pip install --upgrade pip setuptools wheel; \
+    done
+
+# ============================================================
+# 安装 CANN（使用传入的 URL）
+# ============================================================
+
+WORKDIR /root
+
+RUN mkdir -p cann && cd cann && \
+    curl -O "${CANN_TOOLKIT_URL}" && \
+    curl -O "${CANN_A3OPS_URL}" && \
+    curl -O "${CANN_NNAL_URL}" && \
+    chmod +x Ascend-cann*.run && \
+    ./Ascend-cann-toolkit*.run --full --quiet --install-path=/usr/local/Ascend && \
+    ./Ascend-cann-A3*.run --install --quiet --install-path=/usr/local/Ascend && \
+    source /usr/local/Ascend/cann/set_env.sh && \
+    ./Ascend-cann-nnal*.run --install --quiet --install-path=/usr/local/Ascend && \
+    rm -rf cann
+
+# 设置环境变量
+ENV CANN_PATH=/usr/local/Ascend/cann
+ENV NNAL_PATH=/usr/local/Ascend/nnal
+ENV ASCEND_HOME=/usr/local/Ascend
+ENV CANN_VERSION=${CANN_VERSION}
+
+# 添加 CANN 环境初始化脚本
+RUN printf '#!/bin/bash\n\
+source /usr/local/Ascend/cann/set_env.sh 2>/dev/null || true\n\
+source /usr/local/Ascend/nnal/atb/set_env.sh 2>/dev/null || true\n\
+' > /etc/profile.d/cann_env.sh && \
+    chmod +x /etc/profile.d/cann_env.sh
+
+# ============================================================
+# 预安装 pytest 等测试依赖（为所有 Python 版本）
+# ============================================================
+
+RUN for py_dir in cp310-cp310 cp311-cp311 cp312-cp312 cp313-cp313; do \
+        /opt/python/$py_dir/bin/pip install pytest pytest-timeout pytest-xdist hypothesis pyyaml zstandard cmake ninja; \
+    done
+
+# ============================================================
+# 设置工作目录和默认命令
+# ============================================================
+
+WORKDIR /workspace
+
+# 创建 welcome 消息
+RUN printf '\n\
+========================================\n\
+PyTorch NPU Builder Image\n\
+========================================\n\
+CANN Version: %s\n\
+Python Versions: 3.10, 3.11, 3.12, 3.13 (default: 3.11)\n\
+\n\
+To switch Python version:\n\
+  source /usr/local/bin/switch_python.sh 3.12\n\
+\n\
+To setup CANN environment:\n\
+  source /etc/profile.d/cann_env.sh\n\
+========================================\n\
+\n' "${CANN_VERSION}" > /etc/motd
+
+CMD ["bash"]