mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 20:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

[ROCm]: Dockerfile and build script updates

Browse Source 
Add hipblaslt in Dockerfile
	Update docker file to default to ROCm5.6
 	CI scripts update to handle multiple ROCm versions
This commit is contained in:
Rahul Batra 2023-07-28 20:33:04 +00:00
parent 1c99dd5447
commit ef79c19093
2 changed files with 19 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
build/rocm/Dockerfile.ms
View File

@ -2,10 +2,10 @@
FROM rocm/dev-ubuntu-20.04:5.4-complete as rt_build
MAINTAINER Rahul Batra<rahbatra@amd.com>
################################################################################
ARG ROCM_DEB_REPO=http://repo.radeon.com/rocm/apt/5.5/
ARG ROCM_DEB_REPO=http://repo.radeon.com/rocm/apt/5.6/
ARG ROCM_BUILD_NAME=ubuntu
ARG ROCM_BUILD_NUM=main
ARG ROCM_PATH=/opt/rocm-5.5.0
ARG ROCM_PATH=/opt/rocm-5.6.0
ARG DEBIAN_FRONTEND=noninteractive
ARG PYTHON_VERSION=3.9.0
@ -50,6 +50,7 @@ RUN apt-get update --allow-insecure-repositories && DEBIAN_FRONTEND=noninteracti
libxmlsec1-dev \
libffi-dev \
liblzma-dev \
hipblaslt-dev \
wget && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
@ -101,6 +102,6 @@ ENV PATH $PYENV_ROOT/shims:$PYENV_ROOT/bin:$PATH
RUN pyenv install $PYTHON_VERSION
RUN eval "$(pyenv init -)" && pyenv local ${PYTHON_VERSION} && pip3 install --upgrade --force-reinstall setuptools pip==22.0 && pip install numpy==1.21.0 setuptools build wheel six auditwheel scipy pytest pytest-rerunfailures matplotlib absl-py
RUN eval "$(pyenv init -)" && pyenv local ${PYTHON_VERSION} && pip3 install --upgrade --force-reinstall setuptools pip && pip install numpy setuptools build wheel six auditwheel scipy pytest pytest-rerunfailures matplotlib absl-py

16
build/rocm/ci_build.sh
View File

@ -26,6 +26,10 @@
#
# COMMAND: Command to be executed in the docker container
#
# ROCM_DEB_REPO_VERSION: ROCm debian repo version
#
# ROCM_PATH: ROCM path in the docker container
#
# Environment variables read by this script
# WORKSPACE
# XLA_REPO
@ -44,6 +48,8 @@ DOCKERFILE_PATH="${SCRIPT_DIR}/Dockerfile.ms"
DOCKER_CONTEXT_PATH="${SCRIPT_DIR}"
KEEP_IMAGE="--rm"
KEEP_CONTAINER="--rm"
ROCM_DEB_REPO_VERSION="5.6" #default for now is 5.6
ROCM_PATH="/opt/rocm-5.6.0"
POSITIONAL_ARGS=()
RUNTIME_FLAG=1
@ -71,6 +77,14 @@ while [[ $# -gt 0 ]]; do
KEEP_CONTAINER=""
shift 1
;;
--rocm_deb_repo_version)
ROCM_DEB_REPO_VERSION="$2"
shift 2
;;
--rocm_path)
ROCM_PATH="$2"
shift 2
;;
*)
POSITIONAL_ARGS+=("$1")
@ -119,7 +133,7 @@ echo "Python Version (${PYTHON_VERSION})"
if [[ "${RUNTIME_FLAG}" -eq 1 ]]; then
echo "Building (runtime) container (${DOCKER_IMG_NAME}) with Dockerfile($DOCKERFILE_PATH)..."
docker build --target rt_build --tag ${DOCKER_IMG_NAME} \
--build-arg PYTHON_VERSION=$PYTHON_VERSION \
--build-arg PYTHON_VERSION=$PYTHON_VERSION --build-arg ROCM_DEB_REPO="http://repo.radeon.com/rocm/apt/"$ROCM_DEB_REPO_VERSION --build-arg ROCM_PATH=$ROCM_PATH\
-f "${DOCKERFILE_PATH}" "${DOCKER_CONTEXT_PATH}"
else
echo "Building (CI) container (${DOCKER_IMG_NAME}) with Dockerfile($DOCKERFILE_PATH)..."