Dockerfile.a5

#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# This file is a part of the vllm-ascend project.
#

FROM quay.io/ascend/cann:9.2.1-950-ubuntu22.04-py3.12

ENV DEBIAN_FRONTEND=noninteractive

WORKDIR /workspace

# Install clang-15 (for triton-ascend)
RUN apt-get update -y && \
    apt-get install -y git vim wget curl protobuf-compiler libprotobuf-dev net-tools gcc g++ cmake numactl libnuma-dev libibverbs-dev libhiredis-dev clang-15 && \
    update-alternatives --install /usr/bin/clang clang /usr/bin/clang-15 20 && \
    update-alternatives --install /usr/bin/clang++ clang++ /usr/bin/clang++-15 20 && \
    rm -rf /var/cache/apt/* && \
    rm -rf /var/lib/apt/lists/*

# Install modelscope (for fast download) and ray (for multinode)
RUN python3 -m pip install 'modelscope<1.38' 'ray>=2.47.1,<=2.48.0' 'protobuf>3.20.0' && \
    python3 -m pip cache purge

# Install vLLM
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
ARG VLLM_TAG=v0.33.0
ARG VLLM_COMMIT=""
ARG PYTORCH_INDEX_URL="https://download.pytorch.org/whl/cpu/"
# Primary pip index -> huaweicloud mirror. pypi.org via squid intermittently
# returns "No matching distribution" under high build parallelism; the mirror
# is fast and reliable from the cluster. Retries/timeout ride out transient
# index hiccups.
ENV PIP_INDEX_URL="https://repo.huaweicloud.com/repository/pypi/simple" \
    PIP_RETRIES=10 \
    PIP_TIMEOUT=60
ARG CARGO_BUILD_JOBS=16
RUN if [ -n "$VLLM_COMMIT" ]; then \
      git init /vllm-workspace/vllm && \
      git -C /vllm-workspace/vllm fetch --depth 1 $VLLM_REPO "$VLLM_COMMIT" && \
      git -C /vllm-workspace/vllm checkout FETCH_HEAD; \
    else \
      git clone --depth 1 -b $VLLM_TAG $VLLM_REPO /vllm-workspace/vllm; \
    fi && \
    VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index-url ${PYTORCH_INDEX_URL} && \
    # In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
    python3 -m pip uninstall -y triton && \
    # The same step also builds the optional Rust frontend (vllm-rs).
    cd /vllm-workspace/vllm && \
    export CARGO_BUILD_JOBS=$CARGO_BUILD_JOBS && \
    export PATH="/root/.cargo/bin:$PATH" && \
    export PROTOC_INCLUDE=/usr/include && \
    python3 -m pip install -r requirements/build/rust.txt && \
    bash tools/build_rust.sh && \
    python3 -m pip cache purge

# Install vllm-ascend
# Limit cargo parallelism to avoid OOM in constrained build environments.
# Overridable via --build-arg CARGO_BUILD_JOBS=N (e.g. 32-CPU/128Gi buildkitd
# runs up to 2 builds in parallel, so 16 jobs keep memory within ~64Gi per build).
# Limit opc/c++ build parallelism to avoid OOM. Mirrors
# .github/workflows/_build_csrc_cache.yaml.
# TODO: auto-detect buildkitd CPU cores and set MAX_JOBS accordingly

# Dependency versions used by the vllm-ascend build.
ARG SOC_VERSION="ascend950dt_9582"
ARG MAX_JOBS=16
ENV SOC_VERSION=$SOC_VERSION \
    TASK_QUEUE_ENABLE=1 \
    OMP_NUM_THREADS=1 \
    CARGO_BUILD_JOBS=$CARGO_BUILD_JOBS \
    MAX_JOBS=$MAX_JOBS
COPY . /vllm-workspace/vllm-ascend/

ARG ASCEND_INDEX_URL="https://mirrors.huaweicloud.com/ascend/repos/pypi \
    https://mirrors.huaweicloud.com/ascend/repos/pypi/nightly"
ARG COMPILE_CUSTOM_KERNELS=1
RUN export PIP_EXTRA_INDEX_URL="${ASCEND_INDEX_URL}" && \
    export VLLM_BATCH_INVARIANT=1 && \
    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
    source /usr/local/Ascend/nnal/atb/set_env.sh && \
    python3 -m pip install -e /vllm-workspace/vllm-ascend/ --extra-index-url ${PYTORCH_INDEX_URL} && \
    python3 -m pip cache purge

# Install cannbot wheels from the Ascend mirror (only aarch64)
ARG TARGETPLATFORM
RUN case "$TARGETPLATFORM" in \
  linux/arm64) \
    PIP_EXTRA_INDEX_URL="${ASCEND_INDEX_URL}" \
    python3 -m pip install "cannbotdsl===0.6.0+g91e7c8f" "cannbot-arena-net-ops===0.1.0" && \
    python3 -m pip cache purge ;; \
  linux/amd64) \
    echo "TARGET is x86_64, skip cannbot wheels install" ;; \
  *) echo "Unsupported platform: $TARGETPLATFORM" && exit 1 ;; \
esac

# ===== Install FLA =====
ARG FLA_VERSION="26.9.1"
ARG FLA_PACKAGE="flash-linear-attention-npu-a5"
RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
    # Install memfabric_hybrid kernels
    TORCH_DEVICE_BACKEND_AUTOLOAD=0 mfcli kernel install --soc-version A5 && \
    # Install flash-linear-attention-npu from the Ascend mirror
    PIP_EXTRA_INDEX_URL="${ASCEND_INDEX_URL}" \
    python3 -m pip install --force-reinstall --no-deps \
        "${FLA_PACKAGE}===${FLA_VERSION}" && \
    # Install mooncake-transfer-engine-npu from the Ascend nightly mirror
    PIP_EXTRA_INDEX_URL="${ASCEND_INDEX_URL}" \
    python3 -m pip install "mooncake-transfer-engine-npu===0.3.14rc3" && \
    python3 -m pip cache purge

# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
RUN echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib' >> ~/.bashrc

CMD ["/bin/bash"]