Skip to content

Add Dockerfile for building ROCm-based Strata Docker images #1307

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

Closed
jintakhan wants to merge 12 commits into from
Closed

Add Dockerfile for building ROCm-based Strata Docker images #1307

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
f69eb3b
Re-add Dockerfile for ROCm-based Docker builds
jintakhan Oct 7, 2026
e607ece
Merge branch 'Niko1221:main' into docker-rocm
jintakhan Oct 7, 2026
d729736
Fix Dockerfile.rocm header; pin ROCm version to 7.2.4; add backend ar…
jintakhan Oct 7, 2026
539436a
Merge branch 'main' into docker-rocm
jintakhan Oct 7, 2026
d7fc89a
Make BASE_IMAGE a configurable build arg
jintakhan Oct 8, 2026
809a8d6
Include amd-smi description in docker-entrypoint.sh
jintakhan Oct 8, 2026
3692721
Merge branch 'Niko1221:main' into docker-rocm
jintakhan Oct 9, 2026
b01ab8b
Merge branch 'Niko1221:main' into docker-rocm
jintakhan Oct 10, 2026
aa4af9a
Merge branch 'Niko1221:main' into docker-rocm
jintakhan Oct 12, 2026
f57cb37
Merge branch 'Niko1221:main' into docker-rocm
jintakhan Oct 14, 2026
041cd99
Update docker-entrypoint.sh
jintakhan Oct 15, 2026
9316433
Merge branch 'Niko1221:main' into docker-rocm
jintakhan Oct 18, 2026
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
148 changes: 148 additions & 0 deletions Dockerfile.rocm
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,148 @@
# syntax=docker/dockerfile:1
#
# Strata: Qwen3.8-Flash-Next on AMD GPUs via ROCm (RX 7900 XT / XTX (gfx1100),
# RX 9070 series / Radeon AI PRO R9700 (gfx1201) - maintainers-validated; see
# docs/AMD_HIP.md). Vision is CPU-only on the HIP backend (build_vision_cpu).
#
# The engine is compiled during docker build, so the first container start only
# downloads the model (~70 GB) and starts the server. docker build has no GPU,
# so the gfx targets are fixed here: the engine is a fat binary with a code
# object per listed arch. Keep both archs unless you want a faster build for a
# single card (narrow HIP_ARCHITECTURES to just yours).
#
# setup.py's rocm_root() accepts a system ROCm >= 7.0 when hipcc, libhipblas and
# the hip-lang dev files are present (the -full base has all three). This image
# defaults to 7.14.1, but newer ROCm versions may also be used.
#
# Build:
# docker build -t strata -f Dockerfile.rocm .
# docker build -t strata -f Dockerfile.rocm --build-arg HIP_ARCHITECTURES=gfx1100 . # 7900 XT only
# docker build -t strata -f Dockerfile.rocm --build-arg BASE_IMAGE=rocm/dev-ubuntu-26.04:10.1.0-full . # ROCm 10.1.0 build
#
# Run (host needs the amdgpu driver; the container needs KFD + DRI passed through):
# docker run --rm \
# --device=/dev/kfd --device=/dev/dri --group-add video \
# -p 8080:8080 \
# --ulimit memlock=-1 \
# -v strata-data:/data \
# -e MODEL=IQ2_XS \
# strata
#
# hipBLASLt tuning (docs/AMD_HIP.md): tune_hipblaslt is built into the image but
# cannot run at build time (no GPU in the build). Run it once inside a running
# container on the target card, save the table to the /data volume, and mount or
# point STRATA_HIPBLASLT_TUNING at it on later runs. The table is pinned to the
# exact hipBLASLt version that generated it - re-tune after any ROCm bump.
#
# Setup choices are env vars, read by docker-entrypoint.sh: FAMILY, MODEL, CONTEXT,
# VISION (no | yes | cpu), KV (int8 | q4_0 | k8v4), GPU (one card) or GPUS ("0,2"
# or "all", with LAYER_SPLIT), LOW_RAM (auto | on | off), HOST, PORT, API_KEY.
#
# Only the model files, the prepared pack, the MTP layer and the install config
# live in the /data volume; the engine is part of the image. Strata loads 32-62 GB
# into RAM, so a capped container needs -e LOW_RAM=on: setup.py reads the RAM from
# /proc/meminfo, which here is the host's total, not the container's limit. Add an
# API key before exposing the port to a network: -e API_KEY=<secret>. Pass
# -e REINSTALL=1 to change the model settings later.
#
# --gpus all on a host with two usable cards: setup takes both (the layer split is
# its recommended default). Pin one card with -e GPU=0, or name them with
# -e GPUS=0,2. A volume set up for one card switches to the pair on its first start
# on a two-card host unless GPU or GPUS pins it. LOW_RAM=on runs on one card.

ARG BASE_IMAGE=rocm/dev-ubuntu-24.04:7.2.4-complete
FROM ${BASE_IMAGE}

# STRATA_EXECV=1: setup.py replaces itself with the server, so the server is PID 1
# and docker stop's SIGTERM reaches it (see setup.start). Normal Linux starts, which
# don't set it, keep spawning the server as a child.
ENV DEBIAN_FRONTEND=noninteractive PYTHONUNBUFFERED=1 LANG=C.UTF-8 STRATA_EXECV=1 BACKEND=hip

RUN apt-get update && apt-get install -y --no-install-recommends \
build-essential ca-certificates curl git libatomic1 libgomp1 \
python3 python3-pip python3-venv unzip \
&& rm -rf /var/lib/apt/lists/*

WORKDIR /opt/strata
COPY . .

# Target all maintainer- and community-validated archs for a fat build by default.
# Narrow to one for a faster build or add your own (e.g. gfx1030).
ARG HIP_ARCHITECTURES=gfx1100;gfx1101;gfx1200;gfx1201
ARG BUILD_VISION=1

# fail the build early if the base image ever lacks what setup.py's rocm_root()
# needs to accept a system ROCm (hipcc, libhipblas, hip-lang dev files)
RUN test -x /opt/rocm/bin/hipcc \
&& ls /opt/rocm/lib/libhipblas.so* >/dev/null \
&& ls /opt/rocm/lib/cmake/hip-lang/hip-lang-config.cmake /opt/rocm/lib64/cmake/hip-lang/hip-lang-config.cmake 2>/dev/null | grep -q . \
|| { echo "base image lacks the ROCm development files"; exit 1; }

RUN python3 -m venv .venv \
&& .venv/bin/pip install --no-cache-dir --upgrade pip \
&& .venv/bin/pip install --no-cache-dir -r requirements.txt \
&& chmod +x setup.sh docker-entrypoint.sh

# llama.cpp at the pinned commit, then the engine built exactly the way
# setup.py's build_engine_hip() builds it (same defs, same BUILD.json stamp), so
# the first container start reuses the image engine instead of recompiling.
# Vision encoder remains CPU-only for now on HIP; tools/vision is built
# without CUDA here.
RUN .venv/bin/python - <<'PYEOF'
import json, os, pathlib, shutil
import setup


# Target all maintainer- and community-validated archs for a fat build by default.
# Narrow to one for a faster build or add your own (e.g. gfx1030).
llama = setup.get_llama_cpp()
archs = [a for a in os.environ.get("HIP_ARCHITECTURES", "gfx1100;gfx1101;gfx1200;gfx1201").strip().strip('"').replace(",", ";").split(";") if a]
root, libs = setup.rocm_root(archs)
assert root == pathlib.Path("/opt/rocm"), f"rocm_root picked {root}, not the base image's ROCm"
bitcode = next((p for p in (root / "lib" / "llvm" / "amdgcn" / "bitcode", root / "amdgcn" / "bitcode") if p.is_dir()),
root / "amdgcn" / "bitcode")
os.environ.update({"HIP_PLATFORM": "amd", "HIP_COMPILER": "clang", "HIP_RUNTIME": "rocclr",
"ROCM_PATH": str(root), "HIP_PATH": str(root)})
os.environ["LD_LIBRARY_PATH"] = os.pathsep.join(libs + [os.environ.get("LD_LIBRARY_PATH", "")]).rstrip(os.pathsep)
os.environ["PATH"] = os.pathsep.join([str(root / "bin"), str(root / "llvm" / "bin"), os.environ.get("PATH", "")])

setup.cmake_build(setup.ROOT, setup.ROOT / "build-hip", "strata",
["-DSTRATA_ENABLE_HIP=ON", "-DSTRATA_ENABLE_CUDA=OFF", "-DSTRATA_BUILD_TESTS=OFF",
"-DSTRATA_PREFILL_MMQ=ON", "-DCMAKE_HIP_ARCHITECTURES=" + ";".join(archs),
f"-DCMAKE_HIP_COMPILER={root / 'llvm' / 'bin' / 'clang++'}",
f"-DCMAKE_HIP_COMPILER_ROCM_ROOT={root}",
"-DCMAKE_PREFIX_PATH=" + ";".join([str(root), *libs]),
f"-DCMAKE_HIP_FLAGS=--rocm-path={root} --rocm-device-lib-path={bitcode}",
f"-DSTRATA_GGML_DIR={llama}"], None, "")

setup.cmake_build(setup.ROOT, setup.ROOT / "build-hip", "tune_hipblaslt", [], None, "")

eng = setup.ROOT / "engine"
eng.mkdir(exist_ok=True)
shutil.copy2(setup.ROOT / "build-hip" / setup.EXE, eng / setup.EXE)
shutil.copy2(setup.ROOT / "build-hip" / "tune_hipblaslt", eng / "tune_hipblaslt")
meta = {"source": "local-hip", "backend": "hip", "version": setup.source_version(), "archs": archs,
"vision": "none", "lib_dirs": libs, "src": setup.source_hash(setup.ENGINE_SOURCES),
"vision_src": setup.source_hash(setup.VISION_SOURCES)}
if os.environ.get("BUILD_VISION", "1") == "1":
setup.cmake_build(setup.ROOT / "tools" / "vision", setup.ROOT / "build-vision", "strata-vision",
[f"-DLLAMA_DIR={llama}", "-DSTRATA_VISION_CUDA=OFF"], None, "")
shutil.copy2(setup.ROOT / "build-vision" / "bin" / setup.VEXE, eng / setup.VEXE)
meta["vision"] = "cpu"
(eng / "BUILD.json").write_text(json.dumps(meta, indent=1))
PYEOF

# the cmake trees are build-time only; the engine itself is what the container needs
RUN rm -rf build build-hip build-vision

VOLUME ["/data"]
EXPOSE 8080

# /health is answered before the API key gate, so it works with or without one.
# The port only opens after the model loads (1-3 minutes, longer on a first run),
# so the start period is generous: a too short one marks a still-loading container
# unhealthy and a restart policy would kill it mid-download.
HEALTHCHECK --interval=30s --timeout=5s --start-period=600s --retries=3 \
CMD curl -fs "http://127.0.0.1:${PORT:-8080}/health" || exit 1

ENTRYPOINT ["./docker-entrypoint.sh"]
5 changes: 4 additions & 1 deletion docker-entrypoint.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,9 +26,10 @@ PORT="${PORT:-8080}"
API_KEY="${API_KEY:-}"
KV="${KV:-}" # int8 | q4_0 | k8v4; empty: setup.py's own default (int8)
GPUS="${GPUS:-}" # "0,2" or "all": one model across several cards (docs/MULTI_GPU.md)
GPU="${GPU:-}" # one card, numbered as nvidia-smi numbers them
GPU="${GPU:-}" # one card, numbered as nvidia-smi/amd-smi numbers them
LAYER_SPLIT="${LAYER_SPLIT:-}" # with GPUS: where each later card's layers start (default: auto)
LOW_RAM="${LOW_RAM:-auto}" # on: the experts come from the pack's experts.bin, not from RAM
BACKEND="${BACKEND:-}" # cuda | hip | sycl; set to appropriate backend for your GPU
GGUF_DIR="${GGUF_DIR:-}" # a mounted folder with GGUF files you already have: no download
RESIDENT_BUDGET_GIB="${RESIDENT_BUDGET_GIB:-}" # UD-Q4_K_XL: GiB of experts kept in RAM (default: setup's pick)
KV_STREAMING="${KV_STREAMING:-}" # auto | on | off; empty: setup.py's own default (auto)
Expand Down Expand Up @@ -83,6 +84,7 @@ elif [ "${REINSTALL:-0}" = "1" ] || [ ! -f "$cfg" ]; then
if [ -n "$GGUF_DIR" ]; then set -- "$@" --gguf-dir "$GGUF_DIR"; fi
if [ -n "$RESIDENT_BUDGET_GIB" ]; then set -- "$@" --resident-budget-gib "$RESIDENT_BUDGET_GIB"; fi
if [ -n "$KV_STREAMING" ]; then set -- "$@" --kv-streaming "$KV_STREAMING"; fi
if [ -n "$BACKEND" ]; then set -- "$@" --backend "$BACKEND"; fi
.venv/bin/python setup.py --setup --yes "$@"
[ -e "/opt/strata/strata-$tag.json" ] && { cmp -s "/opt/strata/strata-$tag.json" "$cfg" || cp -f "/opt/strata/strata-$tag.json" "$cfg"; }
echo "Config: $cfg (from MODEL $MODEL, just set up)"
Expand All @@ -103,4 +105,5 @@ set -- --port "$PORT"
if [ -n "$GPUS" ]; then set -- "$@" --gpus "$GPUS"; fi
if [ -n "$GPU" ]; then set -- "$@" --gpu "$GPU"; fi
if [ -n "$LAYER_SPLIT" ]; then set -- "$@" --layer-split "$LAYER_SPLIT"; fi
if [ -n "$BACKEND" ]; then set -- "$@" --backend "$BACKEND"; fi
exec .venv/bin/python setup.py "$@"