@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
rocm_base 479 lines 20 KB
Raw
1 ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete
2 ARG TRITON_BRANCH="669b31a" # release/internal/3.8.x as of 09/17
3 ARG TRITON_REPO="https://github.com/ROCm/triton.git"
4 ARG PYTORCH_BRANCH="733fca1" # release/2.13 as of 09/17
5 ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git"
6 ARG PYTORCH_VISION_BRANCH="v0.28.0"
7 ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
8 ARG PYTORCH_AUDIO_BRANCH="v2.11.0"
9 ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
10 ARG FA_BRANCH="0e60e394"
11 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
12 ARG AITER_BRANCH="v0.1.23"
13 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
14 ARG MORI_BRANCH="v1.2.3.post1"
15 ARG MORI_REPO="https://github.com/ROCm/mori.git"
16 ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git"
17 ARG ROCPROFILER_SDK_COMMIT="2b22ab0195cc1461cd9abf3b969e9dd7c10af350"
18 ARG ROCM_RUNTIME_COMMIT="af70edc91f1dc5d228fd4c076df99e65e49c4521"
19
20 # Sccache configuration (only used in release pipeline)
21 ARG USE_SCCACHE
22 ARG SCCACHE_DOWNLOAD_URL
23 ARG SCCACHE_ENDPOINT
24 ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
25 ARG SCCACHE_REGION_NAME=us-west-2
26 ARG SCCACHE_S3_NO_CREDENTIALS=0
27
28
29 FROM ${BASE_IMAGE} AS base
30
31 ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
32 ENV ROCM_PATH=/opt/rocm
33 ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:
34
35 ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151
36 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
37 ENV AITER_ROCM_ARCH=gfx942;gfx950
38 # Note: Do not set MORI_GPU_ARCHS here, it is automatically inferred at runtime
39
40 # Required for RCCL in ROCm7.1
41 ENV HSA_NO_SCRATCH_RECLAIM=1
42
43 ARG PYTHON_VERSION=3.12
44 ENV PYTHON_VERSION=${PYTHON_VERSION}
45
46 RUN mkdir -p /app
47 WORKDIR /app
48 ENV DEBIAN_FRONTEND=noninteractive
49
50 # Install Python and other dependencies (plus gcc-13, selected later in build_pytorch).
51 RUN apt-get update -y \
52 && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \
53 && for i in 1 2 3; do \
54 add-apt-repository -y ppa:deadsnakes/ppa && add-apt-repository -y ppa:ubuntu-toolchain-r/test && break || \
55 { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
56 done \
57 && apt-get update -y \
58 && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \
59 python${PYTHON_VERSION}-lib2to3 python-is-python3 gcc-13 g++-13 \
60 && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
61 && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
62 && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
63 && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \
64 && python3 --version && python3 -m pip --version
65
66 RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' 'pybind11<3.1.0' Cython
67 RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/*
68
69 # Install sccache if USE_SCCACHE is enabled (for release builds)
70 ARG USE_SCCACHE
71 ARG SCCACHE_DOWNLOAD_URL
72 ARG SCCACHE_ENDPOINT
73 ARG SCCACHE_BUCKET_NAME
74 ARG SCCACHE_REGION_NAME
75 ARG SCCACHE_S3_NO_CREDENTIALS
76 RUN if [ "$USE_SCCACHE" = "1" ]; then \
77 echo "Installing sccache..." \
78 && SCCACHE_ARCH="x86_64" \
79 && SCCACHE_VERSION="v0.8.1" \
80 && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \
81 && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \
82 && tar -xzf /tmp/sccache.tar.gz -C /tmp \
83 && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
84 && chmod +x /usr/bin/sccache \
85 && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \
86 && sccache --version; \
87 fi
88
89 # Setup sccache for HIP compilation via HIP_CLANG_PATH
90 # This creates wrapper scripts in a separate directory and points HIP to use them
91 # This avoids modifying the original ROCm binaries which can break detection
92 # NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm)
93 # Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1
94 RUN if [ "$USE_SCCACHE" = "1" ]; then \
95 echo "Setting up sccache wrappers for HIP compilation..." \
96 && mkdir -p /opt/sccache-wrappers \
97 && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \
98 && chmod +x /opt/sccache-wrappers/clang++ \
99 && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \
100 && chmod +x /opt/sccache-wrappers/clang \
101 && echo "sccache wrappers created in /opt/sccache-wrappers"; \
102 fi
103
104 # Set sccache environment variables only when USE_SCCACHE=1
105 # This prevents S3 config from leaking into images when sccache is not used
106 ARG USE_SCCACHE
107 ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
108 ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
109 ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
110 ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
111
112
113 ###
114 ### Triton Build
115 ###
116 FROM base AS build_triton
117 ARG TRITON_BRANCH
118 ARG TRITON_REPO
119 RUN git clone ${TRITON_REPO}
120 RUN cd triton \
121 && git checkout ${TRITON_BRANCH} \
122 && git config --global user.email "you@example.com" && git config --global user.name "Your Name" \
123 && if [ ! -f setup.py ]; then cd python; fi \
124 && MAX_JOBS=16 python3 setup.py bdist_wheel --dist-dir=dist \
125 && mkdir -p /app/install && cp dist/*.whl /app/install
126
127
128 ###
129 ### AMD SMI Build
130 ###
131 FROM base AS build_amdsmi
132 RUN cd /opt/rocm/share/amd_smi \
133 && pip wheel . --wheel-dir=dist
134 RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install
135
136
137 ###
138 ### Pytorch build
139 ###
140 FROM base AS build_pytorch
141 ARG PYTORCH_BRANCH
142 ARG PYTORCH_REPO
143 ARG USE_SCCACHE
144
145 RUN apt-get update && apt-get install -y pkg-config liblzma-dev
146 # torch 2.12 emits float->__bf16 libcalls (__truncsfbf2) absent from gcc-11 libgcc; build with gcc-13.
147 RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-13 130 --slave /usr/bin/g++ g++ /usr/bin/g++-13
148 RUN git clone ${PYTORCH_REPO} pytorch
149 RUN cd pytorch && git checkout ${PYTORCH_BRANCH}
150 RUN cd pytorch \
151 && pip install -r requirements.txt && git submodule update --init --recursive
152 RUN cd pytorch && python3 tools/amd_build/build_amd.py \
153 && if [ "$USE_SCCACHE" = "1" ]; then \
154 export HIP_CLANG_PATH=/opt/sccache-wrappers \
155 && export CMAKE_C_COMPILER_LAUNCHER=sccache \
156 && export CMAKE_CXX_COMPILER_LAUNCHER=sccache \
157 && sccache --show-stats; \
158 fi \
159 && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \
160 && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
161 RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install
162
163
164 ###
165 ### Shared PyTorch runtime for dependent builds
166 ###
167 FROM base AS build_pytorch_runtime
168 RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-13 130 --slave /usr/bin/g++ g++ /usr/bin/g++-13
169 RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install,ro \
170 pip install /install/*.whl
171
172
173 ###
174 ### TorchVision build
175 ###
176 FROM build_pytorch_runtime AS build_torchvision
177 ARG PYTORCH_VISION_BRANCH
178 ARG PYTORCH_VISION_REPO
179 ARG USE_SCCACHE
180
181 RUN git clone ${PYTORCH_VISION_REPO} vision
182 RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \
183 && if [ "$USE_SCCACHE" = "1" ]; then \
184 export HIP_CLANG_PATH=/opt/sccache-wrappers \
185 && export CMAKE_C_COMPILER_LAUNCHER=sccache \
186 && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \
187 fi \
188 && python3 setup.py bdist_wheel --dist-dir=dist \
189 && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
190 RUN mkdir -p /app/install && cp /app/vision/dist/*.whl /app/install
191
192
193 ###
194 ### TorchAudio build
195 ###
196 FROM build_pytorch_runtime AS build_torchaudio
197 ARG PYTORCH_AUDIO_BRANCH
198 ARG PYTORCH_AUDIO_REPO
199 ARG USE_SCCACHE
200
201 RUN git clone ${PYTORCH_AUDIO_REPO} audio
202 RUN cd audio && git checkout ${PYTORCH_AUDIO_BRANCH} \
203 && git submodule update --init --recursive
204 RUN cd audio && pip install -r requirements.txt
205 RUN cd audio \
206 && if [ "$USE_SCCACHE" = "1" ]; then \
207 export HIP_CLANG_PATH=/opt/sccache-wrappers \
208 && export CMAKE_C_COMPILER_LAUNCHER=sccache \
209 && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \
210 fi \
211 && python3 setup.py bdist_wheel --dist-dir=dist \
212 && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
213 RUN mkdir -p /app/install && cp /app/audio/dist/*.whl /app/install
214
215
216 ###
217 ### MORI Build
218 ###
219 FROM base AS build_mori
220 ARG MORI_BRANCH
221 ARG MORI_REPO
222 RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
223 pip install /install/*.whl
224 RUN git clone ${MORI_REPO}
225 RUN cd mori \
226 && git checkout ${MORI_BRANCH} \
227 && git submodule update --init --recursive \
228 && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl
229 RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install
230
231
232 ###
233 ### FlashAttention Build
234 ###
235 FROM base AS build_fa
236 ARG FA_BRANCH
237 ARG FA_REPO
238 ARG USE_SCCACHE
239 RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
240 pip install /install/*.whl
241 RUN git clone ${FA_REPO}
242 RUN cd flash-attention \
243 && git checkout ${FA_BRANCH} \
244 && git submodule update --init \
245 && if [ "$USE_SCCACHE" = "1" ]; then \
246 export HIP_CLANG_PATH=/opt/sccache-wrappers \
247 && sccache --show-stats; \
248 fi \
249 && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \
250 && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
251 RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install
252
253
254 ###
255 ### AITER Build
256 ###
257 FROM base AS build_aiter
258 ARG AITER_BRANCH
259 ARG AITER_REPO
260 ARG USE_SCCACHE
261 RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
262 pip install /install/*.whl
263 RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO}
264 RUN cd aiter \
265 && git submodule update --init --recursive \
266 && pip install -r requirements.txt
267 RUN pip install pyyaml && cd aiter \
268 && if [ "$USE_SCCACHE" = "1" ]; then \
269 export HIP_CLANG_PATH=/opt/sccache-wrappers \
270 && sccache --show-stats; \
271 fi \
272 && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
273 && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
274 && ls /app/aiter/dist/*.whl
275 RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
276
277
278 ###
279 ### Final Build
280 ###
281
282 # Wheel release stage -
283 # only includes dependencies used by wheel release pipeline
284 FROM base AS debs_wheel_release
285 RUN mkdir /app/debs
286 RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
287 cp /install/*.whl /app/debs
288 RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
289 cp /install/*.whl /app/debs
290 RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
291 cp /install/*.whl /app/debs
292 RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
293 cp /install/*.whl /app/debs
294 RUN --mount=type=bind,from=build_torchvision,src=/app/install/,target=/install \
295 cp /install/*.whl /app/debs
296 RUN --mount=type=bind,from=build_torchaudio,src=/app/install/,target=/install \
297 cp /install/*.whl /app/debs
298 RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
299 cp /install/*.whl /app/debs
300
301 # Full debs stage - includes Mori (used by Docker releases)
302 FROM base AS debs
303 RUN mkdir /app/debs
304 RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
305 cp /install/*.whl /app/debs
306 RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
307 cp /install/*.whl /app/debs
308 RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
309 cp /install/*.whl /app/debs
310 RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
311 cp /install/*.whl /app/debs
312 RUN --mount=type=bind,from=build_torchvision,src=/app/install/,target=/install \
313 cp /install/*.whl /app/debs
314 RUN --mount=type=bind,from=build_torchaudio,src=/app/install/,target=/install \
315 cp /install/*.whl /app/debs
316 RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
317 cp /install/*.whl /app/debs
318 RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \
319 cp /install/*.whl /app/debs
320
321 ###
322 ### ROCr + CLR build
323 ###
324 # Pull in updates to ROCr + clr from ROCm 7.2.4 and beyond.
325 FROM base AS build_rocm_runtime
326 ARG ROCM_SYSTEMS_REPO
327 ARG ROCM_RUNTIME_COMMIT
328 RUN apt-get update \
329 && apt-get install -y --no-install-recommends g++ \
330 libelf-dev libdrm-dev libnuma-dev libdw-dev xxd \
331 && rm -rf /var/lib/apt/lists/*
332 RUN pip install --no-cache-dir CppHeaderParser
333 RUN git init -q /src && cd /src \
334 && git remote add origin ${ROCM_SYSTEMS_REPO} \
335 && git config core.sparseCheckout true \
336 && git config remote.origin.promisor true \
337 && git config remote.origin.partialclonefilter blob:none \
338 && git sparse-checkout init --cone \
339 && git sparse-checkout set projects/rocr-runtime projects/clr projects/hip shared cmake \
340 && git fetch --filter=blob:none --depth 1 origin ${ROCM_RUNTIME_COMMIT} \
341 && git checkout -q FETCH_HEAD
342 # ROCr's trap_handler and blit_shaders call find_package(Clang/LLVM REQUIRED) only to get
343 # IMPORTED executables. The ROCm dev image ships the binaries but no CMake package config, so
344 # supply one pointing at them.
345 RUN mkdir -p /opt/rocm-cmake-shim \
346 && printf 'if(NOT TARGET clang)\n add_executable(clang IMPORTED GLOBAL)\n set_target_properties(clang PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/clang")\nendif()\nset(Clang_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/ClangConfig.cmake \
347 && printf 'if(NOT TARGET llvm-objcopy)\n add_executable(llvm-objcopy IMPORTED GLOBAL)\n set_target_properties(llvm-objcopy PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/llvm-objcopy")\nendif()\nset(LLVM_FOUND TRUE)\nset(LLVM_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/LLVMConfig.cmake
348 # ROCr first, and installed into /opt/rocm, because the clr build below needs its headers
349 RUN cd /src/projects/rocr-runtime \
350 && cmake -B build -G Ninja \
351 -DCMAKE_BUILD_TYPE=Release \
352 -DCMAKE_PREFIX_PATH=/opt/rocm \
353 -DCMAKE_INSTALL_PREFIX=/opt/rocm \
354 -DClang_DIR=/opt/rocm-cmake-shim \
355 -DLLVM_DIR=/opt/rocm-cmake-shim \
356 && cmake --build build --parallel 16 \
357 && cmake --install build \
358 && cmake --install build --prefix /rocr-install --strip
359 RUN cd /src/projects/clr \
360 && cmake -B build -G Ninja \
361 -DCMAKE_BUILD_TYPE=Release \
362 -DCLR_BUILD_HIP=ON \
363 -DCLR_BUILD_OCL=OFF \
364 -DHIP_COMMON_DIR=/src/projects/hip \
365 -DROCM_PATH=/opt/rocm \
366 -DCMAKE_PREFIX_PATH=/opt/rocm \
367 -DCMAKE_INSTALL_PREFIX=/opt/rocm \
368 -DLLVM_DIR=/opt/rocm-cmake-shim \
369 -DHIP_LLVM_ROOT=/opt/rocm/llvm \
370 && cmake --build build --parallel 16 \
371 && cmake --install build --prefix /clr-install --strip
372 RUN mkdir -p /staging/lib \
373 && cp -P /rocr-install/lib/libhsa-runtime64.so* /staging/lib/ \
374 && cp -P /clr-install/lib/libamdhip64.so* /staging/lib/ \
375 && ls -l /staging/lib
376
377 # rocprofiler-sdk 1.3.2 fixes torch 2.12 profiler on ROCm 7.2.3
378 FROM base AS build_rocprofiler_sdk
379 ARG ROCM_SYSTEMS_REPO
380 ARG ROCPROFILER_SDK_COMMIT
381 ARG USE_SCCACHE
382 RUN apt-get update \
383 && apt-get install -y --no-install-recommends libdw-dev libelf-dev libsqlite3-dev libzstd-dev \
384 && rm -rf /var/lib/apt/lists/*
385 RUN git init -q /src && cd /src \
386 && git remote add origin ${ROCM_SYSTEMS_REPO} \
387 && git config core.sparseCheckout true \
388 && git config remote.origin.promisor true \
389 && git config remote.origin.partialclonefilter blob:none \
390 && git sparse-checkout init --cone \
391 && git sparse-checkout set projects/rocprofiler-sdk shared cmake \
392 && git fetch --filter=blob:none --depth 1 origin ${ROCPROFILER_SDK_COMMIT} \
393 && git checkout -q FETCH_HEAD
394 # Patch the ROCm 7.2.3 build of rocprofiler-sdk to include queue-interposition fixes that are
395 # absent from the pinned 7.14 build. This resolves hangs on torch.profiler on ROCm.
396 # TODO: Drop these once ${ROCPROFILER_SDK_COMMIT} is bumped to a build that already
397 # contains them, such as ROCm 10.0. See https://github.com/vllm-project/vllm/issues/54087.
398 ARG ROCPROFILER_SDK_PR_7796=fae7fd57ffb8a76e811d6410b8a0cee1e94d710a
399 ARG ROCPROFILER_SDK_PR_7924=f8db22adbe17686a9b05ae03367dd1ec73011653
400 RUN cd /src \
401 && HSA_PATHS='projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/*' \
402 && curl -fsSL -o /tmp/qi-7796.patch "${ROCM_SYSTEMS_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7796}.patch" \
403 && curl -fsSL -o /tmp/qi-7924.patch "${ROCM_SYSTEMS_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7924}.patch" \
404 && git apply -p1 --include="$HSA_PATHS" --verbose /tmp/qi-7796.patch \
405 && git apply -p1 --include="$HSA_PATHS" --verbose /tmp/qi-7924.patch
406 RUN cd /src/projects/rocprofiler-sdk \
407 && if [ "$USE_SCCACHE" = "1" ]; then \
408 export CMAKE_C_COMPILER_LAUNCHER=sccache \
409 && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \
410 fi \
411 && cmake -B build -G Ninja \
412 -DCMAKE_BUILD_TYPE=Release \
413 -DCMAKE_PREFIX_PATH=/opt/rocm \
414 -DCMAKE_INSTALL_PREFIX=/opt/rocm \
415 -DROCPROFILER_BUILD_TESTS=OFF \
416 -DROCPROFILER_BUILD_SAMPLES=OFF \
417 -DROCPROFILER_BUILD_DOCS=OFF \
418 && cmake --build build --parallel "$(nproc)" \
419 && cmake --install build --prefix /staging --strip
420
421 FROM base AS final
422 RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
423 pip install /install/*.whl
424
425 COPY --from=build_rocprofiler_sdk /staging/ /opt/rocm/
426 ENV ROCPROFILER_QUEUE_INTERPOSITION=0
427
428 COPY --from=build_rocm_runtime /staging/ /opt/rocm/
429 # Hard link the patched library onto every stock file name so that both the
430 # SONAME lookup and the hard coded paths resolve to it.
431 RUN set -eux; \
432 cd /opt/rocm/lib; \
433 ours_hsa="$(readlink -f libhsa-runtime64.so.1)"; \
434 ours_hip="$(readlink -f libamdhip64.so.7)"; \
435 for f in libhsa-runtime64.so.1.*; do \
436 [ "$f" = "${ours_hsa##*/}" ] || ln -f "$ours_hsa" "$f"; \
437 done; \
438 for f in libamdhip64.so.7.*; do \
439 [ "$f" = "${ours_hip##*/}" ] || ln -f "$ours_hip" "$f"; \
440 done
441
442 RUN ldconfig
443
444 ARG BASE_IMAGE
445 ARG TRITON_BRANCH
446 ARG TRITON_REPO
447 ARG PYTORCH_BRANCH
448 ARG PYTORCH_VISION_BRANCH
449 ARG PYTORCH_REPO
450 ARG PYTORCH_VISION_REPO
451 ARG PYTORCH_AUDIO_BRANCH
452 ARG PYTORCH_AUDIO_REPO
453 ARG FA_BRANCH
454 ARG FA_REPO
455 ARG AITER_BRANCH
456 ARG AITER_REPO
457 ARG MORI_BRANCH
458 ARG MORI_REPO
459 ARG ROCM_SYSTEMS_REPO
460 ARG ROCPROFILER_SDK_COMMIT
461 ARG ROCM_RUNTIME_COMMIT
462 RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
463 && echo "ROCM_SYSTEMS_REPO: ${ROCM_SYSTEMS_REPO}" >> /app/versions.txt \
464 && echo "ROCPROFILER_SDK_COMMIT: ${ROCPROFILER_SDK_COMMIT}" >> /app/versions.txt \
465 && echo "ROCM_RUNTIME_COMMIT: ${ROCM_RUNTIME_COMMIT}" >> /app/versions.txt \
466 && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
467 && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
468 && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \
469 && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \
470 && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \
471 && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \
472 && echo "PYTORCH_AUDIO_BRANCH: ${PYTORCH_AUDIO_BRANCH}" >> /app/versions.txt \
473 && echo "PYTORCH_AUDIO_REPO: ${PYTORCH_AUDIO_REPO}" >> /app/versions.txt \
474 && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \
475 && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \
476 && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \
477 && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \
478 && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \
479 && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt