From 2b2d578ab8f0da4e0c67fd9f020692c3eab618e0 Mon Sep 17 00:00:00 2001
From: Tongze Wang <wtongze@gmail.com>
Date: Wed, 11 Sep 2024 21:01:57 -0700
Subject: [PATCH] feat: simplify rocm 4.5+

---
 base.dockerfile       | 8 ++++++++
 pytorch.dockerfile    | 8 --------
 tensorflow.dockerfile | 8 --------
 3 files changed, 8 insertions(+), 16 deletions(-)

diff --git a/base.dockerfile b/base.dockerfile
index 6b2921b..54176cd 100644
--- a/base.dockerfile
+++ b/base.dockerfile
@@ -50,6 +50,14 @@ apt-get clean
 rm -rf /var/lib/apt/lists/*
 EOF
 
+# For ROCm 4.5+
+RUN <<EOF
+sudo apt-get update
+sudo apt-get install -y rocm-ml-sdk rocm-hip-sdk || true
+sudo apt-get clean
+sudo rm -rf /var/lib/apt/lists/*
+EOF
+
 # Create non-root user
 RUN <<EOF
 echo "%sudo ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/sudo-nopasswd
diff --git a/pytorch.dockerfile b/pytorch.dockerfile
index 8889e55..7d73297 100644
--- a/pytorch.dockerfile
+++ b/pytorch.dockerfile
@@ -29,14 +29,6 @@ EOF
 
 RUN python${PYTHON_VERSION} -m pip install typing-extensions wheel
 
-# For ROCm 4.5+
-RUN <<EOF
-sudo apt-get update
-sudo apt-get install -y rocm-ml-sdk rocm-hip-sdk
-sudo apt-get clean
-sudo rm -rf /var/lib/apt/lists/*
-EOF
-
 RUN <<EOF
 sudo apt-get update
 sudo apt-get install -y ninja-build
diff --git a/tensorflow.dockerfile b/tensorflow.dockerfile
index 13c6e81..5af6f84 100644
--- a/tensorflow.dockerfile
+++ b/tensorflow.dockerfile
@@ -18,14 +18,6 @@ RUN sudo wget -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/rel
 WORKDIR /home/rocm
 RUN git clone --depth 1 https://github.com/ROCm/tensorflow-upstream.git -b ${TENSORFLOW_BRANCH} tensorflow
 
-# For ROCm 4.5+
-RUN <<EOF
-sudo apt-get update
-sudo apt-get install -y rocm-ml-sdk rocm-hip-sdk
-sudo apt-get clean
-sudo rm -rf /var/lib/apt/lists/*
-EOF
-
 RUN <<EOF
 sudo apt-get update
 sudo apt install -y software-properties-common