# Tested base images (the apt ROCm install below must match the base's ROCm

# version: the base's rocm_sdk wheels warn on every import of a mori built

# against another one, and a different major stacks two ROCm installs):

# rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.14.0 (default)

# rocm/pytorch:rocm10.1.0_ubuntu22.04_py3.10_pytorch_release_2.14.0

# rocm/jax:rocm10.1.0_ubuntu24.04_py3.12_jax_release_0.10.2 (JAX CI)

# rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0 (ROCM_APT_* below)

ARG BASE_IMAGE=rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.14.0

FROM ${BASE_IMAGE}

# Space-separated gfx targets to pull ROCm device packages for. The CCO CI

# image uses "gfx950 gfx1250".

ARG GPU_ARCHS="gfx942 gfx950"

# apt ROCm to install when the base image ships no dev stack (see below).

# {ubuntu} in ROCM_APT_LIST expands to the base's release (2404, 2204). 7.14

# lives in a different repo than 10.x, so the URLs are args too:

# ROCM_APT_VERSION=7.14

# ROCM_APT_GPG=https://repo.amd.com/rocm/packages-multi-arch/gpg/rocm.gpg

# ROCM_APT_LIST="https://repo.amd.com/rocm/packages-multi-arch/ubuntu2404 stable main"

ARG ROCM_APT_VERSION=10.1

ARG ROCM_APT_GPG=https://stable.repo.amd.com/rocm/gpg/packages.gpg

ARG ROCM_APT_LIST="https://stable.repo.amd.com/rocm/core/packages/ubuntu{ubuntu} stable main"

RUN apt-get update && \

apt-get install -y --no-install-recommends \

git ibverbs-utils libibverbs-dev \

openmpi-bin libopenmpi-dev \

libdw1 locales \

libgrpc-dev libgrpc++-dev libprotobuf-dev protobuf-compiler-grpc \

libnuma-dev libdrm-dev \

cmake && \

rm -rf /var/lib/apt/lists/*

# The 10.1 torch and jax images run the system Python, which has python3 but no

# python.

RUN command -v python >/dev/null || ln -s "$(command -v python3)" /usr/local/bin/python

# ── apt ROCm dev toolchain (only when the base image lacks one) ──

#

# The rocm pytorch and jax images have no /opt/rocm: ROCm comes from the pip

# `rocm-sdk-core` wheel, which ships runtime .so files and headers but no CMake

# package configs, so `project(... HIP)` cannot find hip-lang-config.cmake.

# The apt ROCm below installs under /opt/rocm/core-${ROCM_APT_VERSION} and

# registers /opt/rocm/{bin,lib,include,...} via update-alternatives, the layout

# mori expects.

#

# Older base images (the 7.2.x ones) already ship a full ROCm dev stack, so

# probe for the same file mori's CMakeLists checks and skip if it is there.

RUN set -ex; \

if [ -e /opt/rocm/lib/cmake/hip-lang/hip-lang-config.cmake ]; then \

echo "ROCm dev stack already present at /opt/rocm — skipping apt ROCm ${ROCM_APT_VERSION}"; \

else \

apt-get update; \

apt-get install -y --no-install-recommends wget gnupg ca-certificates; \

mkdir -p --mode=0755 /etc/apt/keyrings; \

wget -qO - "${ROCM_APT_GPG}" \

| gpg --dearmor > /etc/apt/keyrings/amdrocm.gpg; \

ubuntu="$(. /etc/os-release && echo "${VERSION_ID}" | tr -d .)"; \

echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdrocm.gpg] $(echo "${ROCM_APT_LIST}" | sed "s/{ubuntu}/${ubuntu}/g")" \

> /etc/apt/sources.list.d/amdrocm.list; \

apt-get update; \

for arch in ${GPU_ARCHS}; do pkgs="${pkgs} amdrocm-core-dev${ROCM_APT_VERSION}-${arch}"; done; \

apt-get install -y --no-install-recommends ${pkgs}; \

rm -rf /var/lib/apt/lists/*; \

fi

# Declared here, next to their only use: every RUN after an ARG has it in its

# cache key, so up top a new FlyDSL version would rebuild the ROCm layers too.

# 1 installs flydsl, which only the cco Python examples need.

ARG WITH_FLYDSL=0

# The version aiter pins, so CI runs mori against the FlyDSL it ships with.

# mori itself also runs on 0.2.x and 0.3.x; nightly.yml's flydsl-compat job

# builds this image with an older one to keep those paths covered.

ARG FLYDSL_VERSION=0.4.0

RUN if [ "${WITH_FLYDSL}" = "1" ]; then pip install --no-cache-dir "flydsl==${FLYDSL_VERSION}"; fi

# ci_cco.yml's fused GEMM job runs its single-GPU and 2-/4-rank cases under xdist.

RUN pip install --no-cache-dir "pytest-xdist==3.8.0"

# Last, so adding them does not rebuild the layers above. setup.py compiles

# through ccache when it is on PATH (ci_run.sh keeps the cache in a volume), and

# the internode jobs pin ranks with numactl.

RUN apt-get update && \

apt-get install -y --no-install-recommends ccache numactl && \

rm -rf /var/lib/apt/lists/*