mirror of
https://github.com/ROCm/ROCm-docker.git
synced 2025-12-25 06:41:45 +01:00
28 lines
851 B
Docker
28 lines
851 B
Docker
FROM ubuntu:24.04
|
|
|
|
ARG ROCM_VERSION=6.4
|
|
ARG AMDGPU_VERSION=6.4.60400
|
|
ARG JAX_VERSION=0.4.35
|
|
ARG PYTHON_VERSION=cp312-cp312
|
|
|
|
#Prequisite packages
|
|
RUN apt update && apt install -y wget gnupg2 software-properties-common python3-pip python3-venv
|
|
|
|
#install ROCm
|
|
RUN wget https://repo.radeon.com/amdgpu-install/$ROCM_VERSION/ubuntu/noble/amdgpu-install_$AMDGPU_VERSION-1_all.deb
|
|
RUN apt install -y ./*.deb
|
|
RUN amdgpu-install --usecase=rocm -y && rm *.deb
|
|
|
|
#virtual environment
|
|
RUN python3 -m venv /opt/venv
|
|
|
|
#install JAX
|
|
RUN /opt/venv/bin/pip install --upgrade pip
|
|
RUN /opt/venv/bin/pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-6.3/jaxlib-0.4.31-cp312-cp312-manylinux_2_28_x86_64.whl
|
|
|
|
# Adjust final path
|
|
ENV PATH="/opt/venv/bin:$PATH:/opt/rocm/bin/"
|
|
|
|
# Set the default shell
|
|
SHELL ["/bin/bash", "-c"]
|