Dockerfile 4.62 KB
Newer Older
Simon Mo's avatar
Simon Mo committed
1
2
3
4
# The vLLM Dockerfile is used to construct vLLM image that can be directly used
# to run the OpenAI compatible server.

#################### BASE BUILD IMAGE ####################
5
FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 AS dev
Stephen Krider's avatar
Stephen Krider committed
6
7

RUN apt-get update -y \
8
    && apt-get install -y python3-pip git
Stephen Krider's avatar
Stephen Krider committed
9

10
11
12
13
14
15
# Workaround for https://github.com/openai/triton/issues/2507 and
# https://github.com/pytorch/pytorch/issues/107960 -- hopefully
# this won't be needed for future versions of this docker image
# or future versions of triton.
RUN ldconfig /usr/local/cuda-12.1/compat/

Stephen Krider's avatar
Stephen Krider committed
16
17
18
19
20
21
WORKDIR /workspace

# install build and runtime dependencies
COPY requirements.txt requirements.txt
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -r requirements.txt
22

Stephen Krider's avatar
Stephen Krider committed
23
24
25
26
# install development dependencies
COPY requirements-dev.txt requirements-dev.txt
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -r requirements-dev.txt
Simon Mo's avatar
Simon Mo committed
27
28
#################### BASE BUILD IMAGE ####################

Stephen Krider's avatar
Stephen Krider committed
29

Simon Mo's avatar
Simon Mo committed
30
#################### EXTENSION BUILD IMAGE ####################
Stephen Krider's avatar
Stephen Krider committed
31
32
FROM dev AS build

33
34
35
36
37
# install build dependencies
COPY requirements-build.txt requirements-build.txt
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -r requirements-build.txt

Stephen Krider's avatar
Stephen Krider committed
38
39
40
# copy input files
COPY csrc csrc
COPY setup.py setup.py
bnellnm's avatar
bnellnm committed
41
42
COPY cmake cmake
COPY CMakeLists.txt CMakeLists.txt
Stephen Krider's avatar
Stephen Krider committed
43
44
45
46
COPY requirements.txt requirements.txt
COPY pyproject.toml pyproject.toml
COPY vllm/__init__.py vllm/__init__.py

Simon Mo's avatar
Simon Mo committed
47
# cuda arch list used by torch
48
49
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.6 8.9 9.0+PTX'
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
Stephen Krider's avatar
Stephen Krider committed
50
# max jobs used by Ninja to build extensions
51
52
ARG max_jobs=2
ENV MAX_JOBS=${max_jobs}
53
54
55
# number of threads used by nvcc
ARG nvcc_threads=8
ENV NVCC_THREADS=$nvcc_threads
56
57
# make sure punica kernels are built (for LoRA)
ENV VLLM_INSTALL_PUNICA_KERNELS=1
58

Stephen Krider's avatar
Stephen Krider committed
59
RUN python3 setup.py build_ext --inplace
Simon Mo's avatar
Simon Mo committed
60
#################### EXTENSION Build IMAGE ####################
Stephen Krider's avatar
Stephen Krider committed
61

62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
#################### FLASH_ATTENTION Build IMAGE ####################
FROM dev as flash-attn-builder
# max jobs used for build
ARG max_jobs=2
ENV MAX_JOBS=${max_jobs}
# flash attention version
ARG flash_attn_version=v2.5.6
ENV FLASH_ATTN_VERSION=${flash_attn_version}

WORKDIR /usr/src/flash-attention-v2

# Download the wheel or build it if a pre-compiled release doesn't exist
RUN pip --verbose wheel flash-attn==${FLASH_ATTN_VERSION} \
    --no-build-isolation --no-deps --no-cache-dir

#################### FLASH_ATTENTION Build IMAGE ####################
Simon Mo's avatar
Simon Mo committed
78
79

#################### TEST IMAGE ####################
Stephen Krider's avatar
Stephen Krider committed
80
81
82
83
84
# image to run unit testing suite
FROM dev AS test

# copy pytorch extensions separately to avoid having to rebuild
# when python code changes
Simon Mo's avatar
Simon Mo committed
85
86
87
88
WORKDIR /vllm-workspace
# ADD is used to preserve directory structure
ADD . /vllm-workspace/
COPY --from=build /workspace/vllm/*.so /vllm-workspace/vllm/
89
90
91
# Install flash attention (from pre-built wheel)
RUN --mount=type=bind,from=flash-attn-builder,src=/usr/src/flash-attention-v2,target=/usr/src/flash-attention-v2 \
    pip install /usr/src/flash-attention-v2/*.whl --no-cache-dir
Simon Mo's avatar
Simon Mo committed
92
93
94
95
# ignore build dependencies installation because we are using pre-complied extensions
RUN rm pyproject.toml
RUN --mount=type=cache,target=/root/.cache/pip VLLM_USE_PRECOMPILED=1 pip install . --verbose
#################### TEST IMAGE ####################
Stephen Krider's avatar
Stephen Krider committed
96
97


Simon Mo's avatar
Simon Mo committed
98
#################### RUNTIME BASE IMAGE ####################
99
100
101
102
# We used base cuda image because pytorch installs its own cuda libraries.
# However cupy depends on cuda libraries so we had to switch to the runtime image
# In the future it would be nice to get a container with pytorch and cuda without duplicating cuda
FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04 AS vllm-base
Stephen Krider's avatar
Stephen Krider committed
103
104
105
106
107
108
109
110
111

# libnccl required for ray
RUN apt-get update -y \
    && apt-get install -y python3-pip

WORKDIR /workspace
COPY requirements.txt requirements.txt
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -r requirements.txt
112
113
114
115
116

# Install flash attention (from pre-built wheel)
RUN --mount=type=bind,from=flash-attn-builder,src=/usr/src/flash-attention-v2,target=/usr/src/flash-attention-v2 \
    pip install /usr/src/flash-attention-v2/*.whl --no-cache-dir

Simon Mo's avatar
Simon Mo committed
117
#################### RUNTIME BASE IMAGE ####################
Stephen Krider's avatar
Stephen Krider committed
118
119


Simon Mo's avatar
Simon Mo committed
120
#################### OPENAI API SERVER ####################
Stephen Krider's avatar
Stephen Krider committed
121
122
# openai api server alternative
FROM vllm-base AS vllm-openai
123
# install additional dependencies for openai api server
Stephen Krider's avatar
Stephen Krider committed
124
RUN --mount=type=cache,target=/root/.cache/pip \
125
    pip install accelerate hf_transfer modelscope
Stephen Krider's avatar
Stephen Krider committed
126

127
COPY --from=build /workspace/vllm/*.so /workspace/vllm/
128
COPY vllm vllm
Stephen Krider's avatar
Stephen Krider committed
129
130

ENTRYPOINT ["python3", "-m", "vllm.entrypoints.openai.api_server"]
Simon Mo's avatar
Simon Mo committed
131
#################### OPENAI API SERVER ####################