Skip to content

fix(docker): build native binaries for every arch and un-pin the base image #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
maple52046 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(docker): build native binaries for every arch and un-pin the base image #1

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
110 changes: 83 additions & 27 deletions Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,20 +12,50 @@
# - base: framework + benchtestkit prerequisites (UCX, MPICH/mpirun,
# rccl-tests).
# - downstream: benchtestkit's own prerequisites and benchtestkit itself.
#
# Build context, next to this Dockerfile:
# pkgs/ - the two hipBLASLt client .debs (see the ARGs in the final stage). The
# public radeon repo does not publish hipblaslt-benchmarks, so the GEMM
# benchmarks' hipblaslt-bench has to come from locally built packages.
#
# Last verified against BASE_IMAGE
# rocm/ali-private:ubuntu22.04_rocm7.2.4_cp313_torch2.12.0_20260726
# (ROCm 7.2.4, Python 3.13, PyTorch 2.12.0, PYTORCH_ROCM_ARCH="gfx942;gfx950"):
# `benchtestkit check-env` reports every requirement present, and the three native
# binaries are fat binaries covering both gfx942 and gfx950.

# Stage `framework`: build a framework image inline from a torch base.
# Alternative: if a prebuilt framework image is available, replace this whole
# stage (the aiter install below) with `FROM <framework-image> AS framework`.
# BenchtestKit is fine with an older, pinned framework, so there is no need to
# keep rebuilding aiter from the latest commit.
FROM rocm/ali-private:ubuntu22.04_rocm7.2.0.43_cp312_torch2.9.1_7e1940d_20260129 AS framework

# Remove the internal apt repo baked into this base image. The base was built
# against an AMD-internal site that is no longer reachable, and a shipped image
# should not carry an internal repo at all (every later `apt` step would
# otherwise fail trying to refresh it). This is base-image-specific: revisit /
# update or drop this line whenever the base image is changed.
RUN apt remove -y amdgpu-install-internal
#
# The torch base is an ARG so the same file can target another ROCm / Python /
# PyTorch baseline without being edited. The default is the combination this file
# was last verified against (see the header).
ARG BASE_IMAGE="rocm/ali-private:ubuntu22.04_rocm7.2.4_cp313_torch2.12.0_20260726"
FROM ${BASE_IMAGE} AS framework

# Remove the internal apt repo if the base image carries one. Older bases were built
# against an AMD-internal site that is no longer reachable, and a shipped image should
# not carry an internal repo at all (every later `apt` step would otherwise fail trying
# to refresh it). Guarded because newer bases do not ship the package, and `apt remove`
# on a package that is not installed exits non-zero and would fail the build.
RUN if dpkg-query -W -f='${Status}' amdgpu-install-internal 2>/dev/null \
| grep -q "install ok installed"; then \
apt remove -y amdgpu-install-internal; \
else \
echo "amdgpu-install-internal is not installed; nothing to remove"; \
fi

# Repair cmake if the base image's is broken. Some torch bases install the pip `cmake`
# package under one interpreter while /usr/local/bin/cmake runs another, so
# `from cmake import cmake` resolves to the /opt/rocm/lib/cmake namespace directory and
# raises "cannot import name 'cmake' from 'cmake'". TransferBench and rccl-tests both
# need cmake, so check first and only reinstall when it is actually broken.
RUN cmake --version >/dev/null 2>&1 \
|| pip install --no-cache-dir "cmake<4"
RUN cmake --version

## Install aiter (develop mode; JIT kernels are left to compile lazily at runtime)
# - Pull the latest commit of the branch; ADD does not fetch submodules.
Expand All @@ -37,7 +67,15 @@ RUN git submodule sync && git submodule update --init --recursive
RUN pip install -r requirements.txt && pip cache purge
ENV AITER_USE_SYSTEM_TRITON=1
RUN GPU_ARCHS="${PYTORCH_ROCM_ARCH:-gfx942}" python3 setup.py develop
ENV PYTHONPATH="/opt/aiter:${PYTHONPATH}"
# Plain "${PYTHONPATH}" leaves a trailing colon when the base image does not set the
# variable, and an empty PYTHONPATH entry puts the current working directory on
# sys.path -- enough for a local ./cmake directory to shadow the installed cmake module.
# Only append the previous value when there is one: this yields "/opt/aiter" on a base
# without PYTHONPATH and "/opt/aiter:<inherited>" on one with it.
# BuildKit prints an "UndefinedVar" lint warning here on bases that do not set the
# variable; that is precisely the case ${var:+...} exists to handle. Do not "fix" it by
# declaring `ARG PYTHONPATH` -- an ARG shadows the inherited ENV and the append is lost.
ENV PYTHONPATH="/opt/aiter${PYTHONPATH:+:${PYTHONPATH}}"

## vLLM is intentionally NOT installed. The paged_attention benchmark defaults
## to the aiter backend (it calls aiter's paged-attention kernel directly, which
Expand Down Expand Up @@ -114,12 +152,12 @@ FROM base AS final

## Basic arguments and environment parameters
ARG BUILD_DATE
ARG ROCM_VERSION="7.2.0"
ARG ROCM_VERSION="7.2.4"
# Set the full OCI annotation set: labels are inherited from the base image, so
# override every standard key to avoid stale/leaked base values. Keys we cannot
# fill are set empty on purpose to clear any inherited value.
LABEL org.opencontainers.image.title="Benchtestkit on Ubuntu 22.04 with ROCm ${ROCM_VERSION}" \
org.opencontainers.image.description="A Docker image for benchtestkit based on Ubuntu 22.04 with ROCm ${ROCM_VERSION}, PyTorch 2.9.1 and Python 3.10, built by DCGPU System Eng TWN, AMD Inc." \
org.opencontainers.image.description="A Docker image for benchtestkit based on Ubuntu 22.04 with ROCm ${ROCM_VERSION}, PyTorch 2.12.0 and Python 3.13, built by DCGPU System Eng TWN, AMD Inc. vLLM is intentionally absent: paged_attention uses the aiter backend; install ROCm vLLM manually in a running container to use its vllm backend." \
org.opencontainers.image.version="${BUILD_DATE:-20260617}" \
org.opencontainers.image.created="${BUILD_DATE:-20260617}" \
org.opencontainers.image.authors="Chen-Hao Ku <Bill.Ku@amd.com>" \
Expand All @@ -142,36 +180,54 @@ LABEL org.opencontainers.image.title="Benchtestkit on Ubuntu 22.04 with ROCm ${R
# Install hipblaslt-bench (the GEMM benchmarks' workload tool). This base image
# does not ship it - hipBLASLt is present but its clients/benchmarks are a
# separate package - so the .deb is installed from the build context (bind-mount
# avoids baking it into a layer). The .debs must sit next to this Dockerfile at
# build time. This is base-image-specific: update the .deb versions, or drop
# this step entirely if a future base image already includes hipblaslt-bench.
ARG HIPBLASLT_BENCHMARKS_DEB="hipblaslt-benchmarks_1.2.1.70200-43~22.04_amd64.deb"
ARG HIPBLASLT_CLIENTS_COMMON_DEB="hipblaslt-clients-common_1.2.1.70200-43~22.04_amd64.deb"
RUN --mount=type=bind,source=${HIPBLASLT_BENCHMARKS_DEB},target=/tmp/${HIPBLASLT_BENCHMARKS_DEB} \
--mount=type=bind,source=${HIPBLASLT_CLIENTS_COMMON_DEB},target=/tmp/${HIPBLASLT_CLIENTS_COMMON_DEB} \
# avoids baking it into a layer). This is base-image-specific: update the .deb
# versions, or drop this step entirely if a future base image already includes
# hipblaslt-bench.
# The .debs live in a pkgs/ subdirectory of the build context rather than beside this
# Dockerfile, so the `COPY . /opt/benchtestkit` below does not drag them into the image.
ARG HIPBLASLT_BENCHMARKS_DEB="hipblaslt-benchmarks_1.2.2-dabb6df2_amd64.deb"
ARG HIPBLASLT_CLIENTS_COMMON_DEB="hipblaslt-clients-common_1.2.2-dabb6df2_amd64.deb"
RUN --mount=type=bind,source=pkgs,target=/tmp/pkgs \
apt-get update && \
apt install -y /tmp/${HIPBLASLT_BENCHMARKS_DEB} /tmp/${HIPBLASLT_CLIENTS_COMMON_DEB} libopenblas0 && \
apt install -y /tmp/pkgs/${HIPBLASLT_BENCHMARKS_DEB} \
/tmp/pkgs/${HIPBLASLT_CLIENTS_COMMON_DEB} libopenblas0 && \
apt-get clean && rm -rf /var/lib/apt/lists/*

# Copy the benchtestkit sources (including vendored native code and the pinned
# TransferBench submodule). See .dockerignore for what is excluded.
WORKDIR /opt/benchtestkit
COPY . /opt/benchtestkit
# pkgs/ only carries the .debs installed above; it is not part of the sources.
RUN rm -rf /opt/benchtestkit/pkgs

# Build the vendored native benchmark binaries and expose them on PATH so
# `benchtestkit check-env` finds them.
# `benchtestkit check-env` finds them. Each one is built for *every* arch in
# PYTORCH_ROCM_ARCH, so a single image runs on any of them and nothing has to be
# compiled at run time.
# - hip-stream (HBM bandwidth, BabelStream HIP backend)
RUN make -C vendor/babelstream-hip && \
# Its HIP.make passes no --offload-arch at all, which leaves the target up to hipcc's
# detection -- and there is no GPU to detect during `docker build`. It does append
# EXTRA_FLAGS, so the archs go in as one repeated flag each.
RUN ARCH_FLAGS="$(echo "${PYTORCH_ROCM_ARCH:-gfx942}" | tr ';' '\n' | sed 's/^/--offload-arch=/' | tr '\n' ' ')" && \
echo "hip-stream arch flags: ${ARCH_FLAGS}" && \
make -C vendor/babelstream-hip EXTRA_FLAGS="${ARCH_FLAGS}" && \
ln -sf /opt/benchtestkit/vendor/babelstream-hip/hip-stream \
/usr/local/bin/hip-stream
# - p2pBandwidthLatencyTest (GPU peer-to-peer bandwidth)
# Its Makefile passes GPU_TARGETS straight into a single --offload-arch, so a
# ';'-separated multi-arch value (e.g. "gfx942;gfx950") would break; build for
# the first arch only (the cmake-based TransferBench below handles the list).
RUN GPU_TARGETS="$(echo "${PYTORCH_ROCM_ARCH:-gfx942}" | cut -d';' -f1)" \
make -C vendor/p2pBandwidthLatencyTest && \
# Its Makefile interpolates GPU_TARGETS into a single --offload-arch, and an unquoted
# ';'-separated value is additionally split by the shell into a second command (which
# surfaces as a confusing "ld.lld: error: undefined symbol: main"). The Makefile does
# append CPPFLAGS verbatim, so pass the first arch through GPU_TARGETS and every
# remaining arch as extra --offload-arch flags. No Makefile or source change needed.
RUN FIRST_ARCH="$(echo "${PYTORCH_ROCM_ARCH:-gfx942}" | cut -d';' -f1)" && \
EXTRA_ARCHS="$(echo "${PYTORCH_ROCM_ARCH:-gfx942}" | cut -d';' -s -f2- | tr ';' '\n' | sed 's/^/--offload-arch=/' | tr '\n' ' ')" && \
echo "p2p GPU_TARGETS=${FIRST_ARCH} CPPFLAGS=${EXTRA_ARCHS}" && \
make -C vendor/p2pBandwidthLatencyTest GPU_TARGETS="${FIRST_ARCH}" CPPFLAGS="${EXTRA_ARCHS}" && \
ln -sf /opt/benchtestkit/vendor/p2pBandwidthLatencyTest/p2pBandwidthLatencyTest \
/usr/local/bin/p2pBandwidthLatencyTest
/usr/local/bin/p2pBandwidthLatencyTest && \
find vendor/p2pBandwidthLatencyTest -maxdepth 1 -type f \
\( -name '*.o' -o -name '*.s' -o -name '*.bc' -o -name '*.ii' -o -name '*.cui' \
-o -name '*.hipi' -o -name '*.tmp' -o -name '*-host-*' -o -name '*-hip-*' \) -delete
# - TransferBench (PCIe bandwidth; cmake submodule)
RUN cmake -S vendor/TransferBench -B vendor/TransferBench/build \
-DCMAKE_BUILD_TYPE=Release \
Expand Down