@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
zen 98 lines 3.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # This Dockerfile builds AMD Zen CPU images for vLLM with zentorch.
5 # It layers on top of the CPU image produced by docker/Dockerfile.cpu.
6 #
7 # Supported platforms:
8 # - linux/amd64 (Zen2 and newer; Zen5 is the canonical CI target)
9 #
10 # Build targets:
11 # vllm-openai-zen (default): vLLM + zentorch for serving
12 # vllm-zen-test : vllm-openai-zen + test deps for CI
13 #
14 # Build workflow (see .buildkite/image_build/image_build_zen_cpu.sh):
15 # 1. Obtain a CPU base image that already has vLLM installed. In CI this is the
16 # shared `<repo>:<commit>-cpu` image (docker/Dockerfile.cpu --target
17 # vllm-test); either the vllm-test or vllm-openai target works as the base
18 # since both install vLLM into the same venv this stage extends.
19 # 2. docker build -f docker/Dockerfile.zen --build-arg BASE_IMAGE=<base>:<tag> \
20 # --target vllm-zen-test -t <out>:<tag> .
21
22 ARG BASE_IMAGE
23 ARG PYTHON_VERSION=3.12
24 ARG ZENTORCH_VERSION
25
26 ######################### ZEN OPENAI IMAGE #########################
27 FROM ${BASE_IMAGE} AS vllm-openai-zen
28
29 ARG TARGETARCH
30 ARG ZENTORCH_VERSION
31
32 RUN if [ "${TARGETARCH}" != "amd64" ]; then \
33 echo "ERROR: Dockerfile.zen requires --platform=linux/amd64"; \
34 exit 1; \
35 fi
36
37 WORKDIR /vllm-workspace
38
39 RUN --mount=type=cache,target=/root/.cache/uv \
40 if [ -n "${ZENTORCH_VERSION}" ]; then \
41 uv pip install "zentorch==${ZENTORCH_VERSION}"; \
42 else \
43 uv pip install "vllm[zen]"; \
44 fi
45
46 LABEL org.opencontainers.image.title="vLLM Zen CPU"
47 LABEL org.opencontainers.image.description="vLLM with zentorch for AMD Zen CPUs"
48 LABEL org.opencontainers.image.vendor="vLLM Project"
49 LABEL org.opencontainers.image.source="https://github.com/vllm-project/vllm"
50 LABEL ai.vllm.build.target-arch="${TARGETARCH}"
51 LABEL ai.vllm.build.python-version="${PYTHON_VERSION}"
52 LABEL ai.vllm.build.zentorch-version="${ZENTORCH_VERSION:-unpinned}"
53
54 ENTRYPOINT ["vllm", "serve"]
55
56 ######################### ZEN TEST IMAGE #########################
57 FROM vllm-openai-zen AS vllm-zen-test
58
59 WORKDIR /vllm-workspace
60
61 COPY requirements/common.txt requirements/common.txt
62 COPY requirements/test/cuda.in requirements/test/zen.in
63
64 RUN sed -i '/mamba_ssm/d' requirements/test/zen.in && \
65 sed -i 's/^torch==.*/torch==2.13.0/g' requirements/test/zen.in && \
66 sed -i 's/torchaudio.*/torchaudio/g' requirements/test/zen.in && \
67 sed -i 's/torchvision.*/torchvision/g' requirements/test/zen.in && \
68 # zentorch parity: keep sentence-transformers pinned where vllm-test-deps did.
69 sed -i 's/^sentence-transformers.*/sentence-transformers==5.3.0/g' requirements/test/zen.in && \
70 uv pip compile requirements/test/zen.in \
71 -o requirements/test/zen.txt \
72 --index-strategy unsafe-best-match --torch-backend cpu
73
74 RUN --mount=type=cache,target=/root/.cache/uv \
75 uv pip install -r requirements/test/zen.txt --torch-backend cpu
76
77 # Reassert zentorch's torch in case test-deps moved it.
78 RUN --mount=type=cache,target=/root/.cache/uv \
79 if [ -n "${ZENTORCH_VERSION}" ]; then \
80 uv pip install --force-reinstall --no-deps "zentorch==${ZENTORCH_VERSION}"; \
81 else \
82 uv pip install --force-reinstall --no-deps "vllm[zen]"; \
83 fi
84
85 ADD ./tests/ ./tests/
86 ADD ./examples/ ./examples/
87 ADD ./benchmarks/ ./benchmarks/
88 ADD ./vllm/collect_env.py .
89 ADD ./docker/ ./docker/
90 ADD ./.buildkite/ ./.buildkite/
91
92 RUN --mount=type=cache,target=/root/.cache/uv \
93 uv pip install -e tests/vllm_test_utils
94
95 ENV HF_XET_HIGH_PERFORMANCE=1
96 ENV HF_HUB_DOWNLOAD_TIMEOUT=60
97
98 ENTRYPOINT []