Dockerfile.rocm

# Copyright (C) 2025 Advanced Micro Devices, Inc. All rights reserved.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

ARG NODE_IMAGE=docker.io/library/node:22-bookworm-slim
ARG BASE_IMAGE=ubuntu:24.04

FROM ${NODE_IMAGE} AS runtime-status-builder

ARG JUPYTERLAB_VERSION="4.5.6"
ARG NPM_REGISTRY=
ARG PNPM_VERSION="10.27.0"

WORKDIR /build/runtime

RUN apt-get update && apt-get install -y --no-install-recommends \
    python3 \
    python3-pip \
    python3-venv \
    && rm -rf /var/lib/apt/lists/*

RUN python3 -m venv /opt/jupyter-build && \
    /opt/jupyter-build/bin/python -m pip install --no-cache-dir --upgrade pip && \
    /opt/jupyter-build/bin/python -m pip install --no-cache-dir jupyterlab=="${JUPYTERLAB_VERSION}"

ENV PATH="/opt/jupyter-build/bin:${PATH}"

RUN corepack enable && corepack prepare pnpm@"${PNPM_VERSION}" --activate

COPY runtime/package.json runtime/pnpm-workspace.yaml runtime/pnpm-lock.yaml runtime/.npmrc ./
COPY runtime/hub/frontend/tsconfig.base.json ./hub/frontend/tsconfig.base.json
COPY runtime/shared/runtime-status /build/runtime/shared/runtime-status
COPY runtime/notebook/jupyterlab-runtime-status /build/runtime/notebook/jupyterlab-runtime-status

RUN rm -rf \
    /build/runtime/shared/runtime-status/node_modules \
    /build/runtime/shared/runtime-status/dist \
    /build/runtime/notebook/jupyterlab-runtime-status/node_modules \
    /build/runtime/notebook/jupyterlab-runtime-status/lib \
    /build/runtime/notebook/jupyterlab-runtime-status/auplc_jupyterlab_runtime_status/labextension

RUN if [ -n "${NPM_REGISTRY}" ]; then pnpm config set registry "${NPM_REGISTRY}"; fi && \
    pnpm install --frozen-lockfile --prod=false --filter @auplc/jupyterlab-runtime-status...

RUN pnpm --filter @auplc/runtime-status run build && \
    pnpm --filter @auplc/jupyterlab-runtime-status run build && \
    test -f /build/runtime/notebook/jupyterlab-runtime-status/auplc_jupyterlab_runtime_status/labextension/package.json && \
    test -n "$(find /build/runtime/notebook/jupyterlab-runtime-status/auplc_jupyterlab_runtime_status/labextension/static -name 'remoteEntry*.js' -print -quit)"

FROM ${BASE_IMAGE}

# Default Shell
SHELL ["/bin/bash", "-c"]

# Create jovyan user with specific UID 1000
ARG NB_USER=jovyan
ARG NB_UID=1000
ARG NB_GID=100

# GPU target and ROCm baseline
#
# Supported GPU_TARGET values:
#   gfx110x       RDNA 3   — gfx1100/1101/1102/1103 (dGPU)
#   gfx1150       RDNA 3.5 — Radeon 890M (Strix Point iGPU)
#   gfx1151       RDNA 3.5 — Radeon 8060S (Strix Halo iGPU)
#   gfx1152       RDNA 3.5 — Ryzen AI 300-series iGPU
#   gfx120x       RDNA 4   — Radeon 9x series (dGPU)
#
# ROCm 7.13 ships TheRock Core SDK packages in the Ubuntu apt repo.
# GPU_TARGET selects the image tag and PyTorch wheel bucket; ROCM_SDK_TARGET
# selects the concrete apt SDK package for size-conscious installs.
ARG GPU_TARGET=gfx1151
ARG ROCM_SDK_TARGET=gfx1151
ARG ROCM_VERSION=7.13.0
ARG PYTORCH_VERSION=2.9.1
ARG TORCHVISION_VERSION=0.24.0
ARG TORCHAUDIO_VERSION=2.9.0
# PyTorch wheel target — only differs from GPU_TARGET for the generic
# buckets (gfx110x → gfx110X-all, gfx120x → gfx120X-all). Specific targets
# like gfx1150/gfx1151/gfx1152 use the same name in both apt and wheel repos.
# See https://repo.amd.com/rocm/whl/ for available targets.
ARG PYTORCH_WHL_TARGET=
ARG PYTORCH_INDEX_URL=

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV ROCM_PATH=/opt/rocm
ENV PATH="${ROCM_PATH}/bin:${PATH}"
ENV LD_LIBRARY_PATH="${ROCM_PATH}/lib"
ENV DEVICE_LIB_PATH=/opt/rocm/lib/llvm/amdgcn/bitcode
ENV HIP_DEVICE_LIB_PATH=/opt/rocm/lib/llvm/amdgcn/bitcode

# Allow pip to install system-wide
RUN printf '[global]\nbreak-system-packages = true\n' > /etc/pip.conf

# Install system dependencies (deduplicated)
RUN apt-get update && apt-get install -y --no-install-recommends \
  sudo \
  ca-certificates \
  curl \
  wget \
  git \
  gnupg \
  build-essential \
  python3-dev \
  python3-pip \
  libnuma-dev \
  libelf1 \
  kmod \
  file \
  locales \
  fonts-liberation \
  run-one \
  libatomic1 \
  unzip \
  libavutil-dev \
  && apt-get clean \
  && rm -rf /var/lib/apt/lists/*

# Download and install GPG key for AMD ROCm repo
RUN mkdir -p /etc/apt/keyrings && \
    wget -q https://repo.amd.com/rocm/packages/gpg/rocm.gpg -O - | \
    gpg --dearmor -o /etc/apt/keyrings/amdrocm.gpg

# Add AMD ROCm apt repository (Ubuntu 24.04)
RUN echo 'deb [arch=amd64 signed-by=/etc/apt/keyrings/amdrocm.gpg] https://repo.amd.com/rocm/packages/ubuntu2404 stable main' \
    > /etc/apt/sources.list.d/rocm.list

# Install ROCm from apt.
# After install, create standard /opt/rocm/{bin,lib,include} symlinks — the amdrocm
# packages install into /opt/rocm/core-<ver>/ without the standard top-level layout,
# so hipcc and device bitcode are not found via PATH / LD_LIBRARY_PATH otherwise.
RUN apt-get update && \
    apt-get install -y --no-install-recommends amdrocm-core-sdk${ROCM_VERSION%.*}-${ROCM_SDK_TARGET} && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/* && \
    CORE_BASE=$(ls -d /opt/rocm/core-[0-9]*.[0-9]* 2>/dev/null | sort -V | tail -1) && \
    if [ -n "$CORE_BASE" ]; then \
        for d in bin lib include libexec share; do \
            if [ ! -e /opt/rocm/$d ] && [ -e "$CORE_BASE/$d" ]; then \
                ln -s "$CORE_BASE/$d" /opt/rocm/$d; \
            fi; \
        done; \
    fi

# Install ROCm PyTorch.
#
# The SDK apt packages accept the same generic RDNA bucket names as GPU_TARGET,
# such as `gfx110x` and `gfx120x`. The pip wheel index at
# https://repo.amd.com/rocm/whl/ uses the "long" `gfx110X-all` /
# `gfx120X-all` paths for those buckets (the short lowercase paths return 403),
# so image targets still map to wheel targets below.
#
# CI passes PYTORCH_WHL_TARGET explicitly via .github/build-config.json. For
# local builds driven by dockerfiles/Makefile or auplc-installer — which only
# set GPU_TARGET — we derive the wheel target here so ad-hoc
# `docker build --build-arg GPU_TARGET=gfx120x` and `make base-rocm
# GPU_TARGET=gfx120x` both Just Work.
RUN WHL_TARGET="${PYTORCH_WHL_TARGET}" && \
    if [ -z "${WHL_TARGET}" ]; then \
        case "${GPU_TARGET}" in \
            gfx110x) WHL_TARGET="gfx110X-all" ;; \
            gfx120x) WHL_TARGET="gfx120X-all" ;; \
            *)       WHL_TARGET="${GPU_TARGET}" ;; \
        esac; \
    fi && \
    INDEX_URL="${PYTORCH_INDEX_URL:-https://repo.amd.com/rocm/whl/${WHL_TARGET}/}" && \
    TORCH_SUFFIX="+rocm${ROCM_VERSION}" && \
    echo "Installing ROCm PyTorch from ${INDEX_URL} (GPU_TARGET=${GPU_TARGET}, WHL_TARGET=${WHL_TARGET})" && \
    python3 -m pip install --no-cache-dir \
    --index-url "${INDEX_URL}" \
    "torch==${PYTORCH_VERSION}${TORCH_SUFFIX}" \
    "torchvision==${TORCHVISION_VERSION}${TORCH_SUFFIX}" \
    "torchaudio==${TORCHAUDIO_VERSION}${TORCH_SUFFIX}" && \
    SDK_INC=$(python3 -c "import _rocm_sdk_core, os; print(os.path.join(os.path.dirname(_rocm_sdk_core.__file__), 'include'))" 2>/dev/null) && \
    if [ -n "$SDK_INC" ] && [ ! -e /opt/rocm/include/hip ]; then \
        ln -s "${SDK_INC}/hip" /opt/rocm/include/hip; \
    fi

# Install JupyterHub and related packages
# jupyterhub must match the Hub image version to ensure singleuser <-> hub
# protocol compatibility. See runtime/chart/Chart.yaml appVersion.
RUN pip3 install --no-cache-dir \
    jupyterhub==5.4.4 \
    jupyterlab==4.5.6 \
    notebook==7.5.5 \
    ipywidgets==8.1.8 \
    ipykernel==6.31.0 \
    matplotlib \
    uv \
    seaborn \
    scikit-learn \
    numpy \
    pandas \
    hiredis \
    pillow \
    tqdm

COPY --from=runtime-status-builder /build/runtime/notebook/jupyterlab-runtime-status /tmp/auplc-jupyterlab-runtime-status

RUN python3 -m pip install --no-cache-dir /tmp/auplc-jupyterlab-runtime-status && \
    rm -rf /tmp/auplc-jupyterlab-runtime-status

# Remove the jenkins user if it exists (user with ID 1000)
RUN if getent passwd 1000 > /dev/null; then \
    userdel -r $(getent passwd 1000 | cut -d: -f1); \
    fi

# Create user and setup permissions
RUN useradd -m -s /bin/bash -N -u $NB_UID -g $NB_GID $NB_USER && \
    echo "$NB_USER ALL=(ALL) NOPASSWD:ALL" >> /etc/sudoers

# Add jovyan to video and render groups for ROCm access
RUN if getent group render; then \
        groupmod -g 992 render; \
    else \
        groupadd -g 992 render; \
    fi
RUN usermod -aG video,render ${NB_USER}

# Create necessary Jupyter directories with correct permissions
RUN mkdir -p /home/$NB_USER/.jupyter && \
    mkdir -p /home/$NB_USER/.local/share/jupyter/runtime && \
    chown -R $NB_UID:$NB_GID /home/$NB_USER/.jupyter && \
    chown -R $NB_UID:$NB_GID /home/$NB_USER/.local

# Create startup script - ensuring it exists in the correct location
RUN echo '#!/bin/bash' > /home/$NB_USER/start-jupyter.sh && \
    echo 'export USER=jovyan' >> /home/$NB_USER/start-jupyter.sh && \
    echo 'export SHELL=/bin/bash' >> /home/$NB_USER/start-jupyter.sh && \
    echo 'exec python3 -m jupyterhub.singleuser --ip=0.0.0.0 --port=8888 "$@"' >> /home/$NB_USER/start-jupyter.sh && \
    chmod +x /home/$NB_USER/start-jupyter.sh && \
    chown $NB_USER:$NB_GID /home/$NB_USER/start-jupyter.sh && \
    # Verify the file exists (will fail build if not)
    ls -la /home/$NB_USER/start-jupyter.sh

# Set proper permissions for ROCm devices
RUN mkdir -p /etc/udev/rules.d && \
    echo 'SUBSYSTEM=="kfd", GROUP="video", MODE="0666"' > /etc/udev/rules.d/70-kfd.rules && \
    echo 'SUBSYSTEM=="dri", GROUP="video", MODE="0666"' > /etc/udev/rules.d/70-dri.rules

# Create entrypoint script to set permissions and start services
RUN echo '#!/bin/bash' > /entrypoint.sh && \
    echo 'chmod 666 /dev/kfd 2>/dev/null || true' >> /entrypoint.sh && \
    echo 'chmod 666 /dev/dri/renderD* 2>/dev/null || true' >> /entrypoint.sh && \
    echo 'export USER=jovyan' >> /entrypoint.sh && \
    echo 'export SHELL=/bin/bash' >> /entrypoint.sh && \
    echo 'exec python3 -m jupyterhub.singleuser --ip=0.0.0.0 --port=8888 "$@"' >> /entrypoint.sh && \
    chmod +x /entrypoint.sh

EXPOSE 8888

USER $NB_UID
WORKDIR /home/jovyan

CMD ["/bin/bash", "/entrypoint.sh"]