-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathDockerfile
141 lines (102 loc) · 4.34 KB
/
Dockerfile
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
ARG BASE_NAME=cpu
###############################################################################
# NVIDIA BASE IMAGE
FROM nvcr.io/nvidia/pytorch:24.11-py3 AS nvidia
RUN apt-get update -y && apt-get install -y python3-venv
ENV VIRTUAL_ENV=/app/.venv
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
RUN python -m venv $VIRTUAL_ENV
RUN . $VIRTUAL_ENV/bin/activate
ARG MAX_JOBS=8
# Put HPC-X MPI in the PATH, i.e. mpirun
ENV PATH=/opt/hpcx/ompi/bin:$PATH
ENV LD_LIBRARY_PATH=/opt/hpcx/ompi/lib:$LD_LIBRARY_PATH
ARG TORCH_VERSION="2.4.0"
RUN pip install uv
RUN uv pip install torch==${TORCH_VERSION}
RUN uv pip install xformers==0.0.27.post2
ENV BASE_NAME=nvidia
###############################################################################
# CPU BASE IMAGE
FROM ubuntu:24.04 AS cpu
RUN --mount=type=cache,target=/var/cache/apt \
apt-get update -y \
&& apt-get install -y python3 python3-pip python3-venv \
openmpi-bin libopenmpi-dev libpmix-dev
ENV VIRTUAL_ENV=/app/.venv
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
RUN python3 -m venv $VIRTUAL_ENV
RUN . $VIRTUAL_ENV/bin/activate
ARG MAX_JOBS=4
#ENV DNNL_DEFAULT_FPMATH_MODE=F32
ARG TORCH_VERSION="2.4.0"
RUN pip install uv
RUN uv pip install torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cpu
ENV BASE_NAME=cpu
###############################################################################
# AMD BASE IMAGE
FROM gdiamos/rocm-base AS amd
ARG MAX_JOBS=8
ENV BASE_NAME=amd
RUN pip install amdsmi
###############################################################################
# VLLM BUILD STAGE
FROM ${BASE_NAME} AS vllm
RUN --mount=type=cache,target=/var/cache/apt \
apt-get update -y \
&& apt-get install -y curl ccache git vim numactl gcc-12 g++-12 libomp-dev libnuma-dev \
&& apt-get install -y ffmpeg libsm6 libxext6 libgl1 libdnnl-dev \
&& update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12
ARG INSTALL_ROOT=/app/cray
COPY ./requirements.txt ${INSTALL_ROOT}/requirements.txt
COPY ./test/requirements-pytest.txt ${INSTALL_ROOT}/requirements-pytest.txt
COPY ./infra/requirements-vllm-build.txt ${INSTALL_ROOT}/requirements-vllm-build.txt
RUN uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements.txt
RUN uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements-vllm-build.txt
RUN uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements-pytest.txt
WORKDIR ${INSTALL_ROOT}
COPY ./infra/cray_infra/vllm ${INSTALL_ROOT}/infra/cray_infra/vllm
COPY ./infra/setup.py ${INSTALL_ROOT}/infra/cray_infra/setup.py
COPY ./infra/CMakeLists.txt ${INSTALL_ROOT}/infra/cray_infra/CMakeLists.txt
COPY ./infra/cmake ${INSTALL_ROOT}/infra/cray_infra/cmake
COPY ./infra/csrc ${INSTALL_ROOT}/infra/cray_infra/csrc
COPY ./infra/requirements-vllm.txt ${INSTALL_ROOT}/infra/cray_infra/requirements.txt
WORKDIR ${INSTALL_ROOT}/infra/cray_infra
ARG VLLM_TARGET_DEVICE=cpu
ARG TORCH_CUDA_ARCH_LIST="7.0 8.6"
# Build vllm python package
RUN \
--mount=type=cache,target=/root/.cache/pip \
--mount=type=cache,target=/root/.cache/ccache \
MAX_JOBS=${MAX_JOBS} \
TORCH_CUDA_ARCH_LIST=${PYTORCH_ROCM_ARCH:-${TORCH_CUDA_ARCH_LIST}} \
VLLM_TARGET_DEVICE=${VLLM_TARGET_DEVICE} \
python ${INSTALL_ROOT}/infra/cray_infra/setup.py bdist_wheel && \
pip install ${INSTALL_ROOT}/infra/cray_infra/dist/*.whl && \
rm -rf ${INSTALL_ROOT}/infra/cray_infra/dist
WORKDIR ${INSTALL_ROOT}
###############################################################################
# MAIN IMAGE
FROM vllm AS infra
RUN apt-get update -y \
&& apt-get install -y slurm-wlm libslurm-dev \
build-essential \
less curl wget net-tools vim iputils-ping \
&& rm -rf /var/lib/apt/lists/*
# Setup python path
ENV PYTHONPATH="${PYTHONPATH}:${INSTALL_ROOT}/infra"
ENV PYTHONPATH="${PYTHONPATH}:${INSTALL_ROOT}/sdk"
ENV PYTHONPATH="${PYTHONPATH}:${INSTALL_ROOT}/ml"
ENV PYTHONPATH="${PYTHONPATH}:${INSTALL_ROOT}/test"
RUN mkdir -p ${INSTALL_ROOT}/jobs
# Copy the rest of the platform code
COPY ./infra ${INSTALL_ROOT}/infra
COPY ./sdk ${INSTALL_ROOT}/sdk
COPY ./test ${INSTALL_ROOT}/test
COPY ./cray ${INSTALL_ROOT}/cray
COPY ./ml ${INSTALL_ROOT}/ml
COPY ./scripts ${INSTALL_ROOT}/scripts
# Build SLURM plugin
RUN /app/cray/infra/slurm_src/compile.sh
ENV SLURM_CONF=${INSTALL_ROOT}/infra/slurm_configs/slurm.conf
ENV HIP_FORCE_DEV_KERNARG=1