[NPU] add support for npu x86_64 image release (#14127)
Signed-off-by: 22dimensions <waitingwind@foxmail.com>
This commit is contained in:
+17
-5
@@ -6,13 +6,15 @@ ARG PYTHON_VERSION=py3.11
|
||||
FROM quay.io/ascend/cann:$CANN_VERSION-$DEVICE_TYPE-$OS-$PYTHON_VERSION
|
||||
|
||||
# Update pip & apt sources
|
||||
ARG TARGETARCH
|
||||
ARG CANN_VERSION
|
||||
ARG DEVICE_TYPE
|
||||
ARG PIP_INDEX_URL="https://pypi.org/simple/"
|
||||
ARG APTMIRROR=""
|
||||
ARG PYTORCH_VERSION="2.8.0"
|
||||
ARG TORCHVISION_VERSION="0.23.0"
|
||||
ARG PTA_URL="https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.8.0/torch_npu-2.8.0.post2-cp311-cp311-manylinux_2_28_aarch64.whl"
|
||||
ARG PTA_URL_ARM64="https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.8.0/torch_npu-2.8.0.post2-cp311-cp311-manylinux_2_28_aarch64.whl"
|
||||
ARG PTA_URL_AMD64="https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.8.0/torch_npu-2.8.0.post2-cp311-cp311-manylinux_2_28_x86_64.whl"
|
||||
ARG SGLANG_TAG=main
|
||||
ARG ASCEND_CANN_PATH=/usr/local/Ascend/ascend-toolkit
|
||||
ARG SGLANG_KERNEL_NPU_TAG=main
|
||||
@@ -20,6 +22,16 @@ ARG SGLANG_KERNEL_NPU_TAG=main
|
||||
ARG PIP_INSTALL="python3 -m pip install --no-cache-dir"
|
||||
ARG DEVICE_TYPE
|
||||
|
||||
RUN if [ "$TARGETARCH" = "amd64" ]; then \
|
||||
echo "Using x86_64 dependencies"; \
|
||||
echo "PTA_URL=$PTA_URL_AMD64" >> /etc/environment_new; \
|
||||
elif [ "$TARGETARCH" = "arm64" ]; then \
|
||||
echo "Using aarch64 dependencies"; \
|
||||
echo "PTA_URL=$PTA_URL_ARM64" >> /etc/environment_new; \
|
||||
else \
|
||||
echo "Unsupported TARGETARCH: $TARGETARCH"; exit 1; \
|
||||
fi
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
# Define environments
|
||||
@@ -63,13 +75,13 @@ RUN ${PIP_INSTALL} sglang-router
|
||||
|
||||
|
||||
### Install PyTorch and PTA
|
||||
RUN (${PIP_INSTALL} torch==${PYTORCH_VERSION} torchvision==${TORCHVISION_VERSION} --index-url https://download.pytorch.org/whl/cpu) \
|
||||
RUN . /etc/environment_new && \
|
||||
(${PIP_INSTALL} torch==${PYTORCH_VERSION} torchvision==${TORCHVISION_VERSION} --index-url https://download.pytorch.org/whl/cpu) \
|
||||
&& (${PIP_INSTALL} ${PTA_URL})
|
||||
|
||||
|
||||
# TODO: install from pypi released triton-ascend
|
||||
RUN (${PIP_INSTALL} pybind11) \
|
||||
&& (${PIP_INSTALL} triton-ascend)
|
||||
## Install triton-ascend
|
||||
RUN (${PIP_INSTALL} pybind11 triton-ascend)
|
||||
|
||||
# Install SGLang
|
||||
RUN git clone https://github.com/sgl-project/sglang --branch $SGLANG_TAG && \
|
||||
|
||||
Reference in New Issue
Block a user