[WIP] torch: add nightly version

arg.json

@@ -162,20 +162,36 @@
   ],
   "pytorch": [
     {
-      "pytorch_version": "2.1.0",
+      "torch_version": "2.1.0",
+      "torch_vision_version": "0.16.0",
+      "torch_audio_version": "2.1.0",
+      "torch_npu_version": "2.1.0.post3",
       "cann_tag": "8.0.rc1-910b-ubuntu22.04-py3.8",
       "tags": [
         "2.1.0",
         "ubuntu-python3.8-cann8.0.rc1-pytorch2.1.0"
       ]
     },
     {
-      "pytorch_version": "2.2.0",
+      "torch_version": "2.2.0",
+      "torch_vision_version": "0.17.0",
+      "torch_audio_version": "2.2.0",
+      "torch_npu_version": "2.2.0",
       "cann_tag": "8.0.rc1-910b-ubuntu22.04-py3.8",
       "tags": [
         "2.2.0",
         "ubuntu-python3.8-cann8.0.rc1-pytorch2.2.0"
       ]
+    },
+    {
+      "torch_version": "nightly",
+      "torch_vision_version": "nightly",
+      "torch_audio_version": "nightly",
+      "torch_npu_version": "nightly",
+      "cann_tag": "8.0.rc3.alpha002-910b-ubuntu22.04-py3.9",
+      "tags": [
+        "nightly"
+      ]
     }
   ],
   "mindspore": [

docker-bake.hcl

@@ -68,7 +68,10 @@ target "pytorch" {
   }
   args = {
     BASE_VERSION = "${item.cann_tag}"
-    PYTORCH_VERSION = "${item.pytorch_version}"
+    TORCH_VERSION = "${item.torch_version}"
+    TORCH_VISION_VERSION = "${item.torch_vision_version}"
+    TORCH_AUDIO_VERSION = "${item.torch_audio_version}"
+    TORCH_NPU_VERSION = "${item.torch_npu_version}"
   }
   tags = generate_tags("pytorch", "${item.tags}")
 }

pytorch/Dockerfile

@@ -2,30 +2,22 @@ ARG BASE_VERSION=latest
 FROM ascendai/cann:${BASE_VERSION} as official
 
 # Arguments
-ARG PYTORCH_VERSION=2.2.0
+ARG TORCH_VERSION=2.1.0
+ARG TORCH_VISION_VERSION=0.16.0
+ARG TORCH_AUDIO_VERSION=2.1.0
+ARG TORCH_NPU_VERSION=2.1.0.post3
 
 # Change the default shell
 SHELL [ "/bin/bash", "-c" ]
 
 # Install pytorch, torch-npu and related packages
-RUN if [ "${PYTORCH_VERSION}" == "2.1.0" ]; then \
-        TORCH_VISION_VERSION=0.16.0; \
-        TORCH_AUDIO_VERSION=2.1.0; \
-        TORCH_NPU_VERSION=2.1.0; \
-    elif [ "${PYTORCH_VERSION}" == "2.2.0" ]; then \
-        TORCH_VISION_VERSION=0.17.0; \
-        TORCH_AUDIO_VERSION=2.2.0; \
-        TORCH_NPU_VERSION=2.2.0; \
-    else \
-        echo "Not supported version: ${PYTORCH_VERSION}."; \
-        exit 1; \
-    fi && \
+RUN \
     # Uninstall the latest numpy and sympy first, as the right versions will be installed again \
     # after installing following packages \
     pip uninstall -y numpy sympy && \
     pip install --no-cache-dir --index-url https://download.pytorch.org/whl/cpu \
-        torch==${PYTORCH_VERSION} \
+        torch==${TORCH_VERSION} \
         torchvision==${TORCH_VISION_VERSION} \
         torchaudio==${TORCH_AUDIO_VERSION} && \
-    pip install --no-cache-dir \
+    pip install --no-cache-dir --index-url https://pypi.org/simple \
         torch-npu==${TORCH_NPU_VERSION}

pytorch/nightly.Dockerfile

@@ -0,0 +1,25 @@
+# Multi-arch builders, see https://github.com/docker/buildx/issues/805
+FROM pytorch/manylinuxaarch64-builder:cpu-aarch64-main AS builder-arm64
+FROM pytorch/manylinux-builder:cpu-main AS builder-amd64
+FROM builder-${TARGETARCH} AS builder
+
+# Get source code and install dependencies
+ARG PY_VERSION=3.9
+RUN git clone https://gitee.com/ascend/pytorch.git --depth 1
+RUN python${PY_VERSION} -m pip install --no-cache-dir -r pytorch/requirements.txt
+
+# The distribution package can be found at /pytorch/dist/torch_npu-xxx.whl
+RUN bash pytorch/ci/build.sh --python=${PY_VERSION}
+
+ARG BASE_VERSION
+FROM ascendai/cann:${BASE_VERSION} as official
+
+# Change the default shell
+SHELL [ "/bin/bash", "-c" ]
+
+# Install pytorch, torch-npu and related packages
+RUN \
+    # Uninstall the latest numpy and sympy first, as the right versions will be installed again \
+    # after installing following packages \
+    pip uninstall -y numpy sympy && \
+    # Nightly build \