update ascend dockerfile (#1687)

docker/Dockerfile.ascend

@@ -1,14 +1,13 @@
 FROM {base_image}
 
-# Use bash so that `source` and other bash builtins work in all following RUN steps.
+ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \
+    PIP_DEFAULT_TIMEOUT=300 \
+    PIP_RETRIES=10 \
+    SOC_VERSION={soc_version}
+
 SHELL ["/bin/bash", "-c"]
 
-ENV PIP_DISABLE_PIP_VERSION_CHECK=1
-ENV PIP_DEFAULT_TIMEOUT=300
-ENV PIP_RETRIES=10
-ENV TRANSFORMERS_VERBOSITY=error
-ENV TRANSFORMERS_NO_ADVISORY_WARNINGS=1
-
+# ---------- System dependencies ----------
 RUN rm -f /etc/apt/apt.conf.d/docker-clean && \
     find /etc/apt/apt.conf.d -maxdepth 1 -type f | xargs -r grep -l "APT::Update::Post-Invoke\|docker-clean" | xargs -r rm -f && \
     apt-get update -y && \
@@ -19,23 +18,62 @@ RUN rm -f /etc/apt/apt.conf.d/docker-clean && \
     rm -rf /var/lib/apt/lists/*
 
 RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \
-    pip config set install.trusted-host mirrors.aliyun.com
+    pip config set install.trusted-host mirrors.aliyun.com && \
+    ARCH=$(uname -m) && \
+    if [ "$ARCH" = "x86_64" ]; then \
+        pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \
+    fi
 
 {extra_content}
+# ---------- Install vllm + vllm-ascend ----------
+RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    if [ -f /usr/local/Ascend/nnal/atb/set_env.sh ]; then source /usr/local/Ascend/nnal/atb/set_env.sh; fi && \
+    git clone --depth 1 --branch v0.14.0 https://github.com/vllm-project/vllm && \
+    git clone --depth 1 --branch v0.14.0rc1 https://github.com/vllm-project/vllm-ascend.git
 
-# Reuse the vllm-ascend base image and only add the extra repos we need.
-# --depth 1 keeps the image smaller; branch/tag names work with shallow clone.
+RUN ARCH=$(uname -m) && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    # Install torch & torch_npu & torchvision
+    pip install torch==2.9.0 torch_npu==2.9.0 torchvision==0.24.0 && \
+    # Install vllm
+    cd vllm && VLLM_TARGET_DEVICE=empty pip install -v -e . && cd .. && \
+    # Install vllm-ascend
+    cd vllm-ascend && pip install -v -e . && cd ..
+
+# ---------- Clone training-side repositories ----------
 RUN git clone --depth 1 --branch v0.15.3 https://github.com/NVIDIA/Megatron-LM.git /Megatron-LM && \
     git clone --depth 1 --branch core_r0.15.3 https://gitcode.com/Ascend/MindSpeed.git /MindSpeed && \
     GIT_LFS_SKIP_SMUDGE=1 git clone --depth 1 -b {swift_branch} --single-branch https://github.com/modelscope/ms-swift.git /ms-swift && \
     git clone --depth 1 https://github.com/modelscope/mcore-bridge.git /mcore-bridge
 
+# ---------- Install training-side repositories ----------
 RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
     if [ -f /usr/local/Ascend/nnal/atb/set_env.sh ]; then source /usr/local/Ascend/nnal/atb/set_env.sh; fi && \
     cd /MindSpeed && pip install --no-cache-dir -e . && \
     cd /mcore-bridge && pip install --no-cache-dir -e . && \
-    pip cache purge
+    cd /ms-swift && pip install --no-cache-dir -e .
 
+# ---------- Pin torch to the correct version + torch_npu ----------
+# x86: must force-install the CPU build from pytorch.org/whl/cpu
+# aarch64: PyPI only provides the CPU build, so install it directly from the Aliyun mirror
+RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    if [ -f /usr/local/Ascend/nnal/atb/set_env.sh ]; then source /usr/local/Ascend/nnal/atb/set_env.sh; fi && \
+    ARCH=$(uname -m) && \
+    if [ "$ARCH" = "x86_64" ]; then \
+        pip install --no-cache-dir --force-reinstall --no-deps \
+            --index-url https://download.pytorch.org/whl/cpu \
+            torch==2.9.0 torchvision==0.24.0 torchaudio==2.9.0; \
+    else \
+        pip install --no-cache-dir --force-reinstall --no-deps \
+            torch==2.9.0 torchvision==0.24.0 torchaudio==2.9.0; \
+    fi && \
+    pip install --no-cache-dir --force-reinstall --no-deps \
+        torch_npu==2.9.0 && \
+    rm -rf /root/.cache/pip
+
+# ---------- Remove CUDA-only dependencies pulled in by vllm (they cause missing libtorch_cuda.so errors on NPU) ----------
+RUN pip uninstall -y flashinfer tvm-ffi torch-c-dlpack-ext 2>/dev/null || true
 ARG INSTALL_MS_DEPS={install_ms_deps}
 
 ENV MEGATRON_LM_PATH=/Megatron-LM
@@ -68,6 +106,8 @@ fi
 ARG CUR_TIME={cur_time}
 RUN echo $CUR_TIME
 
+RUN pip install --no-cache-dir --no-build-isolation OpenCC
+
 RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
     if [ -f /usr/local/Ascend/nnal/atb/set_env.sh ]; then source /usr/local/Ascend/nnal/atb/set_env.sh; fi && \
     pip install --no-cache-dir -U funasr scikit-learn && \

docker/build_image.py

@@ -1,5 +1,6 @@
 import argparse
 import os
+import platform
 import subprocess
 from datetime import datetime
 from typing import Any
@@ -461,10 +462,46 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
 
 class AscendImageBuilder(StableGPUImageBuilder):
 
+    @staticmethod
+    def _normalize_arch(arch: str = None) -> str:
+        arch = arch or platform.machine()
+        arch = arch.lower()
+        arch_mapping = {
+            'x86': 'x86',
+            'x86_64': 'x86',
+            'amd64': 'x86',
+            'arm': 'arm',
+            'aarch64': 'arm',
+            'arm64': 'arm',
+        }
+        if arch not in arch_mapping:
+            raise ValueError(f'Unsupported architecture: {arch}. '
+                             'Please pass --arch x86 or --arch arm.')
+        return arch_mapping[arch]
+
+    @staticmethod
+    def _get_atlas_hardware(soc_version: str) -> str:
+        soc_version = soc_version.lower()
+        atlas_mapping = {
+            'ascend910b1': 'A2',
+            'ascend910_9391': 'A3',
+            'ascend310p1': '300I',
+        }
+        if soc_version.startswith('ascend950'):
+            return 'A5'
+        if soc_version not in atlas_mapping:
+            raise ValueError(
+                f'Unsupported soc_version: {soc_version}. '
+                'Supported values are ascend910b1, ascend910_9391, '
+                'ascend310p1, and values starting with ascend950.')
+        return atlas_mapping[soc_version]
+
     def init_args(self, args) -> Any:
         if not args.base_image:
             # Reuse the prebuilt vllm-ascend image to avoid rebuilding its stack.
-            args.base_image = 'quay.io/ascend/vllm-ascend:v0.14.0rc1-a3'
+            args.base_image = 'quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11'
+        args.arch = self._normalize_arch(args.arch)
+        args.atlas_hardware = self._get_atlas_hardware(args.soc_version)
         return super().init_args(args)
 
     def generate_dockerfile(self) -> str:
@@ -474,6 +511,7 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
         with open('docker/Dockerfile.ascend', 'r') as f:
             content = f.read()
             content = content.replace('{base_image}', self.args.base_image)
+            content = content.replace('{soc_version}', self.args.soc_version)
             content = content.replace('{extra_content}', extra_content)
             content = content.replace('{cur_time}', formatted_time)
             content = content.replace('{install_ms_deps}', 'False')
@@ -484,8 +522,9 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
 
     def image(self) -> str:
         return (
-            f'{docker_registry}:{self.args.base_image.split(":")[-1]}-torch2.7.1'
-            f'-{self.args.modelscope_version}-ascend-test')
+            f'{docker_registry}:{self.args.swift_branch}-'
+            f'{self.args.atlas_hardware}-{self.args.python_tag}-{self.args.arch}'
+        )
 
     def push(self):
         return 0
@@ -510,6 +549,8 @@ parser.add_argument('--optimum_version', type=str, default=None)
 parser.add_argument('--modelscope_branch', type=str, default='master')
 parser.add_argument('--modelscope_version', type=str, default='9.99.0')
 parser.add_argument('--swift_branch', type=str, default='main')
+parser.add_argument('--soc_version', type=str, default='ascend910_9391')
+parser.add_argument('--arch', type=str, choices=['x86', 'arm'], default=None)
 parser.add_argument('--dry_run', type=int, default=0)
 args = parser.parse_args()