update npu dockerfile (#1755)

docker/Dockerfile.ascend

@@ -9,14 +9,31 @@ ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \
 SHELL ["/bin/bash", "-c"]
 
 # ---------- System dependencies ----------
-RUN rm -f /etc/apt/apt.conf.d/docker-clean && \
+RUN set -eux; \
-    find /etc/apt/apt.conf.d -maxdepth 1 -type f | xargs -r grep -l "APT::Update::Post-Invoke\|docker-clean" | xargs -r rm -f && \
+    . /etc/os-release; \
-    apt-get update -y && \
+    case "${ID,,}" in \
-    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+        ubuntu) \
-        gcc g++ cmake ninja-build libnuma-dev libgl1 libglib2.0-0 libsm6 libxext6 libxrender1 \
+            rm -f /etc/apt/apt.conf.d/docker-clean; \
-        wget git curl jq vim build-essential ca-certificates && \
+            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 clean && \
+            apt-get update -y; \
-    rm -rf /var/lib/apt/lists/*
+            DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+                gcc g++ cmake ninja-build libnuma-dev libgl1 libglib2.0-0 libsm6 libxext6 libxrender1 \
+                wget git curl jq vim build-essential ca-certificates; \
+            apt-get clean; \
+            rm -rf /var/lib/apt/lists/* \
+            ;; \
+        openeuler) \
+            yum install -y \
+                gcc gcc-c++ cmake ninja-build numactl-devel mesa-libGL glib2 libSM libXext libXrender \
+                wget git curl jq vim make ca-certificates; \
+            yum clean all; \
+            rm -rf /var/cache/yum \
+            ;; \
+        *) \
+            echo "Unsupported base image OS: ${ID}" >&2; \
+            exit 1 \
+            ;; \
+    esac
 
 RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \
     pip config set global.extra-index-url "https://pypi.org/simple" && \
@@ -30,14 +47,14 @@ RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \
 # ---------- 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.18.0 https://github.com/vllm-project/vllm && \
+    git clone --depth 1 --branch {vllm_git_ref} https://github.com/vllm-project/vllm && \
-    git clone --depth 1 --branch v0.18.0 https://github.com/vllm-project/vllm-ascend.git
+    git clone --depth 1 --branch {vllm_ascend_git_ref} https://github.com/vllm-project/vllm-ascend.git
 
 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.post2 torchvision==0.24.0 && \
+    pip install torch=={torch_version} torch_npu=={torch_npu_version} torchvision=={torchvision_version} && \
     # Install vllm
     cd vllm && VLLM_TARGET_DEVICE=empty pip install -v -e . && cd .. && \
     # Install vllm-ascend
@@ -46,14 +63,13 @@ RUN ARCH=$(uname -m) && \
 # ---------- Clone training-side repositories ----------
 RUN git clone --depth 1 --branch {megatron_branch} https://github.com/NVIDIA/Megatron-LM.git /Megatron-LM && \
     git clone --depth 1 --branch {mindspeed_branch} 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_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 install --no-cache-dir mcore-bridge -i https://pypi.org/simple/ -U && \
     cd /ms-swift && pip install --no-cache-dir -e .
 
 # ---------- Pin torch to the correct version + torch_npu ----------
@@ -65,13 +81,13 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
     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; \
+            torch=={torch_version} torchvision=={torchvision_version} torchaudio=={torchaudio_version}; \
     else \
         pip install --no-cache-dir --force-reinstall --no-deps \
-            torch==2.9.0 torchvision==0.24.0 torchaudio==2.9.0; \
+            torch=={torch_version} torchvision=={torchvision_version} torchaudio=={torchaudio_version}; \
     fi && \
     pip install --no-cache-dir --force-reinstall --no-deps \
-        torch_npu==2.9.0.post2 && \
+        torch_npu=={torch_npu_version} && \
     rm -rf /root/.cache/pip
 
 # ---------- Remove CUDA-only dependencies pulled in by vllm (they cause missing libtorch_cuda.so errors on NPU) ----------
@@ -126,36 +142,19 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
     pip install --no-cache-dir omegaconf==2.3.0 && \
     pip cache purge
 
-# ---------- Reinstall triton-ascend for the selected CANN version ----------
+# ---------- Install training and evaluation dependencies ----------
+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 && \
+    TORCH_DEVICE_BACKEND_AUTOLOAD=0 pip install --no-cache-dir "deepspeed<0.19" ray liger_kernel pre-commit -U && \
+    pip cache purge
+
+# ---------- Install triton-ascend ----------
 RUN set -eux; \
     pip uninstall -y triton || true; \
     pip uninstall -y triton-ascend || true; \
-    case "${CANN_VERSION}" in \
+    pip install --no-cache-dir --force-reinstall \
-        8.5.*) \
+        triton-ascend=={triton_ascend_version} \
-            pip install --no-cache-dir --force-reinstall triton-ascend==3.2.0; \
+        --extra-index-url=https://triton-ascend.osinfra.cn/pypi/simple
-            ;; \
-        9.0.0) \
-            PY_ABI="cp$(python -c 'import sys; print(f"{sys.version_info.major}{sys.version_info.minor}")')"; \
-            case "${PY_ABI}" in \
-                cp310|cp311|cp312|cp313) ;; \
-                *) echo "Unsupported Python ABI for triton-ascend 3.2.1: ${PY_ABI}" >&2; exit 1 ;; \
-            esac; \
-            ARCH="$(uname -m)"; \
-            case "${ARCH}" in \
-                aarch64|x86_64) ;; \
-                *) echo "Unsupported architecture for triton-ascend 3.2.1: ${ARCH}" >&2; exit 1 ;; \
-            esac; \
-            WHEEL_NAME="triton_ascend-3.2.1-${PY_ABI}-${PY_ABI}-manylinux_2_27_${ARCH}.manylinux_2_28_${ARCH}.whl"; \
-            WHEEL_PATH="/tmp/${WHEEL_NAME}"; \
-            curl -fL "https://gitcode.com/Ascend/triton-ascend/releases/download/v3.2.1/${WHEEL_NAME}" -o "${WHEEL_PATH}"; \
-            pip install --no-cache-dir --force-reinstall "${WHEEL_PATH}"; \
-            rm -f "${WHEEL_PATH}"; \
-            ;; \
-        *) \
-            echo "Unsupported CANN_VERSION for triton-ascend install: ${CANN_VERSION}" >&2; \
-            exit 1; \
-            ;; \
-    esac
 
 RUN echo 'source /usr/local/Ascend/ascend-toolkit/set_env.sh' >> /root/.bashrc && \
     echo '[ -f /usr/local/Ascend/nnal/atb/set_env.sh ] && source /usr/local/Ascend/nnal/atb/set_env.sh' >> /root/.bashrc && \

docker/OVERVIEW.ascend.md

@@ -10,7 +10,8 @@ ms-swift Ascend images provide a ready-to-use ms-swift environment for Huawei As
 - Build template: `docker/Dockerfile.ascend`
 - Build entrypoint: `docker/build_image.py --image_type ascend`
 - Default base image: `quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11`
-- Default output tag: `${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-<arch>`
+- Supported base OSes: Ubuntu and openEuler, selected from the CANN base-image tag
+- Default output tag: `${DOCKER_REGISTRY}:main-cann8.5.1-torch_npu2.9.0.post2-a3-ubuntu22.04-py3.11-<arch>`
 - Ascend runtime environment is sourced from `/usr/local/Ascend/ascend-toolkit/set_env.sh`
 - If available, NNAL/ATB runtime is sourced from `/usr/local/Ascend/nnal/atb/set_env.sh`
 
@@ -22,45 +23,46 @@ The Ascend Dockerfile installs and configures:
 | --- | --- |
 | CANN | inherited from the selected `quay.io/ascend/cann` base image |
 | Python | inherited from the base image tag, for example `py3.11` |
-| PyTorch | `torch==2.9.0` |
+| PyTorch | `torch==2.9.0` by default; configurable with `--torch_version` |
-| torch-npu | `torch_npu==2.9.0.post2` |
+| torch-npu | `torch_npu==2.9.0.post2` by default; configurable with `--torch_npu_version` |
-| torchvision / torchaudio | `torchvision==0.24.0`, `torchaudio==2.9.0` |
+| torchvision / torchaudio | `torchvision==0.24.0`, `torchaudio==2.9.0` by default; pass both explicitly when overriding `--torch_version` |
-| vLLM | source install from `vllm-project/vllm`, default branch `v0.18.0` |
+| vLLM | source install from `vllm-project/vllm`, default `0.18.0`; configurable with `--vllm_version` |
-| vLLM Ascend | source install from `vllm-project/vllm-ascend`, default branch `v0.18.0` |
+| vLLM Ascend | source install from `vllm-project/vllm-ascend`, default `0.18.0`; configurable with `--vllm_ascend_version` |
 | Megatron-LM | source checkout, default branch `v0.15.3` |
 | MindSpeed | source checkout, default branch `core_r0.15.3` |
-| mcore-bridge | source checkout from `modelscope/mcore-bridge` |
+| mcore-bridge | latest release from PyPI |
 | ms-swift | source checkout from `modelscope/ms-swift`, default branch `main` |
 | ModelScope | source checkout from `modelscope/modelscope`, default branch `master` |
-| triton-ascend | `3.2.0` for CANN `8.5.*`; local wheel install of `3.2.1` for CANN `9.0.0` |
+| triton-ascend | CANN `8.5.*` defaults to `3.2.0`; CANN `9.0.*` defaults to `3.2.1`; configurable with `--triton_ascend_version` and installed from the Triton Ascend PyPI index |
 
 ## Supported Tag Format
 
 Images built by `docker/build_image.py --image_type ascend` use this tag format:
 
 ```text
-${DOCKER_REGISTRY}:<swift-branch>-<atlas-hardware>-<python-tag>-<cann-version-tag>-<os-tag>-<arch>
+${DOCKER_REGISTRY}:<swift-branch>-<cann-version-tag>-torch_npu<torch-npu-version>-<atlas-hardware>-<os-tag>-<python-tag>-<arch>
 ```
 
 | Field | Example | Description |
 | --- | --- | --- |
 | `swift-branch` | `main` | ms-swift branch used during image build |
-| `atlas-hardware` | `A2`, `A3`, `300I`, `A5` | Derived from `--soc_version` |
+| `cann-version-tag` | `cann8.5.1`, `cann9.0.0` | Parsed from the CANN base image tag |
-| `python-tag` | `py311` | Derived from `--python_version` |
+| `torch-npu-version` | `2.9.0.post2` | From `--torch_npu_version`; defaults to `2.9.0.post2` |
-| `cann-version-tag` | `CANN8.5.1`, `CANN9.0.0` | Parsed from the CANN base image tag |
+| `atlas-hardware` | `a2`, `a3`, `300i`, `a5` | Derived from `--soc_version` |
-| `os-tag` | `ubuntu22.04` | Parsed from the CANN base image tag |
+| `os-tag` | `ubuntu22.04`, `openeuler24.03` | Parsed from the CANN base-image tag; prevents tags for different OSes from colliding |
-| `arch` | `arm`, `x86` | Derived from host architecture or `--arch` |
+| `python-tag` | `py3.11` | Parsed from the CANN base image tag |
+| `arch` | `aarch64`, `x86_64` | Derived from host architecture or `--arch` |
 
 Default example on an ARM64 host:
 
 ```text
-${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-arm
+${DOCKER_REGISTRY}:main-cann8.5.1-torch_npu2.9.0.post2-a3-ubuntu22.04-py3.11-aarch64
 ```
 
 A2 / CANN 9.0.0 example:
 
 ```text
-${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm
+${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64
 ```
 
 ## Build Locally
@@ -85,6 +87,39 @@ python docker/build_image.py \
   --soc_version ascend910b1
 ```
 
+Build an openEuler image. The system-dependency layer automatically uses `yum`; Ubuntu images continue to use `apt-get`.
+
+```bash
+python docker/build_image.py \
+  --image_type ascend \
+  --base_image quay.io/ascend/cann:8.5.1-a3-openeuler24.03-py3.11 \
+  --soc_version ascend910_9391
+```
+
+Override the PyTorch stack. `--torch_version` must match the base version of
+`--torch_npu_version`; when overriding PyTorch, pass its matching torchvision
+and torchaudio versions explicitly.
+
+```bash
+python docker/build_image.py \
+  --image_type ascend \
+  --torch_version 2.9.0 \
+  --torch_npu_version 2.9.0.post2 \
+  --torchvision_version 0.24.0 \
+  --torchaudio_version 2.9.0
+```
+
+Override the vLLM stack or triton-ascend. The vLLM version arguments select
+the matching Git tag, for example `0.18.0` selects `v0.18.0`.
+
+```bash
+python docker/build_image.py \
+  --image_type ascend \
+  --vllm_version 0.18.0 \
+  --vllm_ascend_version 0.18.0 \
+  --triton_ascend_version 3.2.1
+```
+
 Override Megatron or MindSpeed source branches when needed:
 
 ```bash
@@ -94,11 +129,11 @@ python docker/build_image.py \
   --mindspeed_branch core_r0.15.3
 ```
 
-For slow networks, Linux hosts can use Docker host networking after the root `Dockerfile` is generated:
+To run the rendered Dockerfile manually, use:
 
 ```bash
-docker build --network host \
+docker build \
-  -t ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \
+  -t ${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 \
   -f Dockerfile .
 ```
 
@@ -119,7 +154,7 @@ docker run --rm -it \
   -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
   -v /etc/ascend_install.info:/etc/ascend_install.info \
   -v /mnt/workspace:/mnt/workspace \
-  ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \
+  ${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 \
   bash
 ```
 
@@ -147,7 +182,8 @@ pip show ms-swift modelscope torch-npu triton-ascend
 ## Notes
 
 - CANN, firmware, and driver versions must be compatible with each other.
-- CANN `8.5.*` and CANN `9.0.0` use different `triton-ascend` install paths in this Dockerfile.
+- Ubuntu base images install system dependencies through `apt-get`; openEuler base images install the corresponding RPM packages through `yum`.
+- `triton-ascend` is installed from `https://triton-ascend.osinfra.cn/pypi/simple`; select a version compatible with the chosen CANN, Python, and architecture.
 - The image is intended for Ascend NPU ms-swift workflows. CUDA-only packages pulled in by dependencies are removed when they conflict with NPU runtime libraries.
 - Use a fixed image tag for production jobs instead of relying on a moving branch name.
 

docker/OVERVIEW.ascend.zh.md

@@ -10,7 +10,8 @@ ms-swift Ascend 镜像面向华为昇腾 Atlas NPU，提供可直接使用的 ms
 - 构建模板：`docker/Dockerfile.ascend`
 - 构建入口：`docker/build_image.py --image_type ascend`
 - 默认基础镜像：`quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11`
-- 默认输出 tag：`${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-<arch>`
+- 支持的基础 OS：Ubuntu 和 openEuler，由 CANN 基础镜像 tag 选择
+- 默认输出 tag：`${DOCKER_REGISTRY}:main-cann8.5.1-torch_npu2.9.0.post2-a3-ubuntu22.04-py3.11-<arch>`
 - Ascend runtime 环境来自 `/usr/local/Ascend/ascend-toolkit/set_env.sh`
 - 如果镜像内存在 NNAL/ATB，则会加载 `/usr/local/Ascend/nnal/atb/set_env.sh`
 
@@ -22,45 +23,46 @@ Ascend Dockerfile 会安装和配置：
 | --- | --- |
 | CANN | 继承自选定的 `quay.io/ascend/cann` 基础镜像 |
 | Python | 继承自基础镜像 tag，例如 `py3.11` |
-| PyTorch | `torch==2.9.0` |
+| PyTorch | 默认 `torch==2.9.0`；可通过 `--torch_version` 配置 |
-| torch-npu | `torch_npu==2.9.0.post2` |
+| torch-npu | 默认 `torch_npu==2.9.0.post2`；可通过 `--torch_npu_version` 配置 |
-| torchvision / torchaudio | `torchvision==0.24.0`，`torchaudio==2.9.0` |
+| torchvision / torchaudio | 默认 `torchvision==0.24.0`、`torchaudio==2.9.0`；覆盖 `--torch_version` 时必须同时显式传入两者 |
-| vLLM | 从 `vllm-project/vllm` 源码安装，默认分支 `v0.18.0` |
+| vLLM | 从 `vllm-project/vllm` 源码安装，默认 `0.18.0`；可通过 `--vllm_version` 配置 |
-| vLLM Ascend | 从 `vllm-project/vllm-ascend` 源码安装，默认分支 `v0.18.0` |
+| vLLM Ascend | 从 `vllm-project/vllm-ascend` 源码安装，默认 `0.18.0`；可通过 `--vllm_ascend_version` 配置 |
 | Megatron-LM | 源码 checkout，默认分支 `v0.15.3` |
 | MindSpeed | 源码 checkout，默认分支 `core_r0.15.3` |
-| mcore-bridge | 来自 `modelscope/mcore-bridge` 的源码 checkout |
+| mcore-bridge | PyPI 上的最新发布版 |
 | ms-swift | 来自 `modelscope/ms-swift` 的源码 checkout，默认分支 `main` |
 | ModelScope | 来自 `modelscope/modelscope` 的源码 checkout，默认分支 `master` |
-| triton-ascend | CANN `8.5.*` 安装 `3.2.0`；CANN `9.0.0` 下载并本地安装 `3.2.1` wheel |
+| triton-ascend | CANN `8.5.*` 默认 `3.2.0`；CANN `9.0.*` 默认 `3.2.1`；可通过 `--triton_ascend_version` 配置，并从 Triton Ascend PyPI 源安装 |
 
 ## 支持的 Tag 格式
 
 通过 `docker/build_image.py --image_type ascend` 构建的镜像使用以下 tag 格式：
 
 ```text
-${DOCKER_REGISTRY}:<swift-branch>-<atlas-hardware>-<python-tag>-<cann-version-tag>-<os-tag>-<arch>
+${DOCKER_REGISTRY}:<swift-branch>-<cann-version-tag>-torch_npu<torch-npu-version>-<atlas-hardware>-<os-tag>-<python-tag>-<arch>
 ```
 
 | 字段 | 示例 | 说明 |
 | --- | --- | --- |
 | `swift-branch` | `main` | 构建镜像时使用的 ms-swift 分支 |
-| `atlas-hardware` | `A2`、`A3`、`300I`、`A5` | 从 `--soc_version` 推导 |
+| `cann-version-tag` | `cann8.5.1`、`cann9.0.0` | 从 CANN 基础镜像 tag 解析 |
-| `python-tag` | `py311` | 从 `--python_version` 推导 |
+| `torch-npu-version` | `2.9.0.post2` | 来自 `--torch_npu_version`，默认 `2.9.0.post2` |
-| `cann-version-tag` | `CANN8.5.1`、`CANN9.0.0` | 从 CANN 基础镜像 tag 解析 |
+| `atlas-hardware` | `a2`、`a3`、`300i`、`a5` | 从 `--soc_version` 推导 |
-| `os-tag` | `ubuntu22.04` | 从 CANN 基础镜像 tag 解析 |
+| `os-tag` | `ubuntu22.04`、`openeuler24.03` | 从 CANN 基础镜像 tag 解析；避免不同 OS 的镜像 tag 冲突 |
-| `arch` | `arm`、`x86` | 从宿主机架构或 `--arch` 推导 |
+| `python-tag` | `py3.11` | 从 CANN 基础镜像 tag 解析 |
+| `arch` | `aarch64`、`x86_64` | 从宿主机架构或 `--arch` 推导 |
 
 ARM64 宿主机上的默认示例：
 
 ```text
-${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-arm
+${DOCKER_REGISTRY}:main-cann8.5.1-torch_npu2.9.0.post2-a3-ubuntu22.04-py3.11-aarch64
 ```
 
 A2 / CANN 9.0.0 示例：
 
 ```text
-${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm
+${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64
 ```
 
 ## 本地构建
@@ -85,6 +87,36 @@ python docker/build_image.py \
   --soc_version ascend910b1
 ```
 
+构建 openEuler 镜像。系统依赖层会自动使用 `yum`，Ubuntu 镜像继续使用 `apt-get`：
+
+```bash
+python docker/build_image.py \
+  --image_type ascend \
+  --base_image quay.io/ascend/cann:8.5.1-a3-openeuler24.03-py3.11 \
+  --soc_version ascend910_9391
+```
+
+覆盖 PyTorch 版本组。`--torch_version` 必须与 `--torch_npu_version` 的基础版本一致；覆盖 PyTorch 时，必须显式传入匹配的 torchvision 和 torchaudio 版本：
+
+```bash
+python docker/build_image.py \
+  --image_type ascend \
+  --torch_version 2.9.0 \
+  --torch_npu_version 2.9.0.post2 \
+  --torchvision_version 0.24.0 \
+  --torchaudio_version 2.9.0
+```
+
+覆盖 vLLM 版本组或 triton-ascend。vLLM 版本参数会选择对应 Git tag，例如 `0.18.0` 会选择 `v0.18.0`。
+
+```bash
+python docker/build_image.py \
+  --image_type ascend \
+  --vllm_version 0.18.0 \
+  --vllm_ascend_version 0.18.0 \
+  --triton_ascend_version 3.2.1
+```
+
 需要时可以覆盖 Megatron 或 MindSpeed 源码分支：
 
 ```bash
@@ -94,11 +126,11 @@ python docker/build_image.py \
   --mindspeed_branch core_r0.15.3
 ```
 
-如果构建时网络较慢，Linux 宿主机可以在根目录 `Dockerfile` 生成后使用 host network 构建：
+如需手工构建生成后的根目录 `Dockerfile`，可使用：
 
 ```bash
-docker build --network host \
+docker build \
-  -t ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \
+  -t ${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 \
   -f Dockerfile .
 ```
 
@@ -119,7 +151,7 @@ docker run --rm -it \
   -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
   -v /etc/ascend_install.info:/etc/ascend_install.info \
   -v /mnt/workspace:/mnt/workspace \
-  ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \
+  ${DOCKER_REGISTRY}:main-cann9.0.0-torch_npu2.9.0.post2-a2-ubuntu22.04-py3.11-aarch64 \
   bash
 ```
 
@@ -147,7 +179,8 @@ pip show ms-swift modelscope torch-npu triton-ascend
 ## 注意事项
 
 - CANN、firmware 和 driver 版本必须互相兼容。
-- 这个 Dockerfile 对 CANN `8.5.*` 和 CANN `9.0.0` 使用不同的 `triton-ascend` 安装路径。
+- Ubuntu 基础镜像通过 `apt-get` 安装系统依赖；openEuler 基础镜像通过 `yum` 安装对应 RPM 包。
+- `triton-ascend` 从 `https://triton-ascend.osinfra.cn/pypi/simple` 安装；请选择与 CANN、Python 和架构兼容的版本。
 - 该镜像面向 Ascend NPU 上的 ms-swift 工作流。依赖安装过程中引入且与 NPU runtime 冲突的 CUDA-only 包会被移除。
 - 生产任务建议使用固定镜像 tag，不要依赖浮动分支名。
 

docker/build_image.py

@@ -469,20 +469,33 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
 
 class AscendImageBuilder(StableGPUImageBuilder):
 
+    _DEFAULT_TORCH_VERSION = '2.9.0'
+    _DEFAULT_TORCHVISION_VERSION = '0.24.0'
+    _DEFAULT_TORCHAUDIO_VERSION = '2.9.0'
+    _DEFAULT_TORCH_NPU_VERSION = '2.9.0.post2'
+    _DEFAULT_VLLM_VERSION = '0.18.0'
+    _DEFAULT_VLLM_ASCEND_VERSION = '0.18.0'
+    _DEFAULT_TRITON_ASCEND_VERSIONS = {
+        '8.5': '3.2.0',
+        '9.0': '3.2.1',
+    }
     _CANN_VERSION_PATTERN = re.compile(r'^\d+(?:\.[0-9A-Za-z]+)+$')
     _OS_TAG_PATTERN = re.compile(r'^[A-Za-z]+[0-9][0-9A-Za-z.]*$')
+    _PYTHON_TAG_PATTERN = re.compile(r'^py\d+\.\d+$', re.IGNORECASE)
+    _TORCH_NPU_VERSION_PATTERN = re.compile(
+        r'^(?P<torch_version>\d+\.\d+\.\d+)(?:\.post\d+)?$')
 
     @staticmethod
     def _normalize_arch(arch: str = None) -> str:
         arch = arch or platform.machine()
         arch = arch.lower()
         arch_mapping = {
-            'x86': 'x86',
+            'x86': 'x86_64',
-            'x86_64': 'x86',
+            'x86_64': 'x86_64',
-            'amd64': 'x86',
+            'amd64': 'x86_64',
-            'arm': 'arm',
+            'arm': 'aarch64',
-            'aarch64': 'arm',
+            'aarch64': 'aarch64',
-            'arm64': 'arm',
+            'arm64': 'aarch64',
         }
         if arch not in arch_mapping:
             raise ValueError(f'Unsupported architecture: {arch}. '
@@ -522,25 +535,107 @@ class AscendImageBuilder(StableGPUImageBuilder):
 
         cann_version = parts[0]
         os_tag = parts[2]
+        python_tag = parts[3]
         if not cls._CANN_VERSION_PATTERN.fullmatch(cann_version):
             raise ValueError(f'Invalid CANN version in Ascend base image tag: '
                              f'{cann_version}')
         if not cls._OS_TAG_PATTERN.fullmatch(os_tag):
             raise ValueError(
                 f'Invalid OS tag in Ascend base image tag: {os_tag}')
+        if not cls._PYTHON_TAG_PATTERN.fullmatch(python_tag):
+            raise ValueError(
+                f'Invalid Python tag in Ascend base image tag: {python_tag}')
 
-        return cann_version, f'CANN{cann_version}', os_tag
+        return cann_version, f'CANN{cann_version}', os_tag, python_tag
+
+    @staticmethod
+    def _get_os_family(os_tag: str) -> str:
+        os_tag = os_tag.lower()
+        if os_tag.startswith('ubuntu'):
+            return 'ubuntu'
+        if os_tag.startswith('openeuler'):
+            return 'openeuler'
+        raise ValueError(f'Unsupported Ascend base image OS tag: {os_tag}. '
+                         'Supported OS families are Ubuntu and openEuler.')
+
+    @classmethod
+    def _init_torch_versions(cls, args) -> None:
+        torch_version_specified = args.torch_version is not None
+        torchvision_version_specified = args.torchvision_version is not None
+        torchaudio_version_specified = args.torchaudio_version is not None
+
+        if torch_version_specified:
+            if (not torchvision_version_specified
+                    or not torchaudio_version_specified):
+                raise ValueError(
+                    'When overriding --torch_version for an Ascend image, also '
+                    'pass matching --torchvision_version and '
+                    '--torchaudio_version.')
+        elif torchvision_version_specified or torchaudio_version_specified:
+            raise ValueError(
+                '--torchvision_version and --torchaudio_version require an '
+                'explicit --torch_version for an Ascend image.')
+
+        args.torch_version = args.torch_version or cls._DEFAULT_TORCH_VERSION
+        args.torchvision_version = (
+            args.torchvision_version or cls._DEFAULT_TORCHVISION_VERSION)
+        args.torchaudio_version = (
+            args.torchaudio_version or cls._DEFAULT_TORCHAUDIO_VERSION)
+        args.torch_npu_version = (
+            args.torch_npu_version or cls._DEFAULT_TORCH_NPU_VERSION)
+
+        match = cls._TORCH_NPU_VERSION_PATTERN.fullmatch(
+            args.torch_npu_version)
+        if not match:
+            raise ValueError('Invalid --torch_npu_version. Expected '
+                             '<major>.<minor>.<patch> or '
+                             '<major>.<minor>.<patch>.post<revision>.')
+        if args.torch_version != match.group('torch_version'):
+            raise ValueError(
+                '--torch_version must exactly match the base version of '
+                f'--torch_npu_version, got torch={args.torch_version} and '
+                f'torch_npu={args.torch_npu_version}.')
+
+    @classmethod
+    def _init_component_versions(cls, args) -> None:
+        args.vllm_version = args.vllm_version or cls._DEFAULT_VLLM_VERSION
+        args.vllm_ascend_version = (
+            args.vllm_ascend_version or cls._DEFAULT_VLLM_ASCEND_VERSION)
+        args.vllm_git_ref = cls._get_vllm_git_ref(args.vllm_version)
+        args.vllm_ascend_git_ref = cls._get_vllm_git_ref(
+            args.vllm_ascend_version)
+
+        if not args.triton_ascend_version:
+            cann_series = '.'.join(args.cann_version.split('.')[:2])
+            try:
+                args.triton_ascend_version = (
+                    cls._DEFAULT_TRITON_ASCEND_VERSIONS[cann_series])
+            except KeyError as e:
+                raise ValueError('No default triton-ascend version for CANN '
+                                 f'{args.cann_version}. Please pass '
+                                 '--triton_ascend_version explicitly.') from e
+
+    @staticmethod
+    def _get_vllm_git_ref(version: str) -> str:
+        return version if version.startswith('v') else f'v{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/cann:8.5.1-a3-ubuntu22.04-py3.11'
+        self._init_torch_versions(args)
         args.arch = self._normalize_arch(args.arch)
         args.atlas_hardware = self._get_atlas_hardware(args.soc_version)
-        args.cann_version, args.cann_version_tag, args.os_tag = (
+        (args.cann_version, args.cann_version_tag, args.os_tag,
-            self._get_cann_os_tags(args.base_image))
+         args.ascend_python_tag) = (
+             self._get_cann_os_tags(args.base_image))
+        self._get_os_family(args.os_tag)
+        self._init_component_versions(args)
         return super().init_args(args)
 
+    def _generate_python_tag(self, _python_version: str) -> str:
+        return self.args.ascend_python_tag
+
     def generate_dockerfile(self) -> str:
         extra_content = """
 RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
@@ -550,6 +645,19 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
             content = content.replace('{base_image}', self.args.base_image)
             content = content.replace('{soc_version}', self.args.soc_version)
             content = content.replace('{cann_version}', self.args.cann_version)
+            content = content.replace('{torch_version}',
+                                      self.args.torch_version)
+            content = content.replace('{torchvision_version}',
+                                      self.args.torchvision_version)
+            content = content.replace('{torchaudio_version}',
+                                      self.args.torchaudio_version)
+            content = content.replace('{torch_npu_version}',
+                                      self.args.torch_npu_version)
+            content = content.replace('{vllm_git_ref}', self.args.vllm_git_ref)
+            content = content.replace('{vllm_ascend_git_ref}',
+                                      self.args.vllm_ascend_git_ref)
+            content = content.replace('{triton_ascend_version}',
+                                      self.args.triton_ascend_version)
             content = content.replace('{extra_content}', extra_content)
             content = content.replace('{cur_time}', formatted_time)
             content = content.replace('{install_ms_deps}', 'False')
@@ -563,11 +671,12 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
         return content
 
     def image(self) -> str:
-        return (
+        tag = (f'{self.args.swift_branch}-{self.args.cann_version_tag}-'
-            f'{docker_registry}:{self.args.swift_branch}-'
+               f'torch_npu{self.args.torch_npu_version}-'
-            f'{self.args.atlas_hardware}-{self.args.python_tag}-'
+               f'{self.args.atlas_hardware}-{self.args.os_tag}-'
-            f'{self.args.cann_version_tag}-{self.args.os_tag}-{self.args.arch}'
+               f'{self.args.python_tag}-'
-        )
+               f'{self.args.arch}')
+        return f'{docker_registry}:{tag.lower()}'
 
     def push(self):
         return 0
@@ -579,12 +688,15 @@ parser.add_argument('--image_type', type=str)
 parser.add_argument('--python_version', type=str, default='3.12.13')
 parser.add_argument('--ubuntu_version', type=str, default='22.04')
 parser.add_argument('--torch_version', type=str, default=None)
+parser.add_argument('--torch_npu_version', type=str, default=None)
 parser.add_argument('--torchvision_version', type=str, default=None)
 parser.add_argument('--cuda_version', type=str, default=None)
 parser.add_argument('--ci_image', type=int, default=0)
 parser.add_argument('--torchaudio_version', type=str, default=None)
 parser.add_argument('--tf_version', type=str, default=None)
 parser.add_argument('--vllm_version', type=str, default=None)
+parser.add_argument('--vllm_ascend_version', type=str, default=None)
+parser.add_argument('--triton_ascend_version', type=str, default=None)
 parser.add_argument('--lmdeploy_version', type=str, default=None)
 parser.add_argument('--flashattn_version', type=str, default=None)
 parser.add_argument('--autogptq_version', type=str, default=None)

tests/msdatasets/test_stream_load.py

@@ -44,15 +44,12 @@ class TestStreamLoad(unittest.TestCase):
                     pass
 
             self.assertIs(HfFileSystem._open, hf_datasets_util._hf_fs_open)
-            self.assertIsNot(
+            self.assertIsNot(hf_datasets_util._hf_fs_open_original,
-                hf_datasets_util._hf_fs_open_original,
+                             hf_datasets_util._hf_fs_open)
-                hf_datasets_util._hf_fs_open)
+            self.assertIs(HfFileSystem.__init__,
-            self.assertIs(
+                          hf_datasets_util._hf_fs_init_with_cookie)
-                HfFileSystem.__init__,
+            self.assertIsNot(hf_datasets_util._hf_fs_init_original,
-                hf_datasets_util._hf_fs_init_with_cookie)
+                             hf_datasets_util._hf_fs_init_with_cookie)
-            self.assertIsNot(
-                hf_datasets_util._hf_fs_init_original,
-                hf_datasets_util._hf_fs_init_with_cookie)
         finally:
             HfFileSystem._open = hf_fs_open_before
             HfFileSystem.__init__ = hf_fs_init_before
@@ -75,7 +72,8 @@ class TestStreamLoad(unittest.TestCase):
                     'load_dataset',
                     side_effect=RuntimeError('load failed')):
                 with self.assertRaises(RuntimeError):
-                    with hf_datasets_util.load_dataset_with_ctx(streaming=True):
+                    with hf_datasets_util.load_dataset_with_ctx(
+                            streaming=True):
                         pass
 
             self.assertIs(HfFileSystem._open, hf_fs_open_clean)