From bfcace4933315e0678366e7dc5fe9909440907ff Mon Sep 17 00:00:00 2001
From: Ruturaj4 <Ruturaj.Vaidya@amd.com>
Date: Tue, 26 Nov 2024 15:33:47 -0600
Subject: [PATCH] [ROCm] ci build and dockerfile changes

---
 .../build_wheels/Dockerfile.manylinux_2_28_x86_64.rocm     | 3 +++
 build/rocm/ci_build                                        | 7 +++++++
 build/rocm/tools/get_rocm.py                               | 3 +++
 3 files changed, 13 insertions(+)

diff --git a/build/rocm/build_wheels/Dockerfile.manylinux_2_28_x86_64.rocm b/build/rocm/build_wheels/Dockerfile.manylinux_2_28_x86_64.rocm
index a67a7ecb2e22..6e610e711c77 100644
--- a/build/rocm/build_wheels/Dockerfile.manylinux_2_28_x86_64.rocm
+++ b/build/rocm/build_wheels/Dockerfile.manylinux_2_28_x86_64.rocm
@@ -4,6 +4,9 @@ ARG ROCM_VERSION=6.1.1
 ARG ROCM_BUILD_JOB
 ARG ROCM_BUILD_NUM
 
+# Install system GCC and C++ libraries.
+RUN yum install -y gcc-c++.x86_64
+
 RUN --mount=type=cache,target=/var/cache/dnf \
     --mount=type=bind,source=build/rocm/tools/get_rocm.py,target=get_rocm.py \
     python3 get_rocm.py --rocm-version=$ROCM_VERSION --job-name=$ROCM_BUILD_JOB --build-num=$ROCM_BUILD_NUM
diff --git a/build/rocm/ci_build b/build/rocm/ci_build
index a32f502f377a..849c082dca38 100755
--- a/build/rocm/ci_build
+++ b/build/rocm/ci_build
@@ -47,6 +47,13 @@ def dist_wheels(
     # create manylinux image with requested ROCm installed
     image = "jax-manylinux_2_28_x86_64_rocm%s" % rocm_version.replace(".", "")
 
+    # Try removing the Docker image.
+    try:
+        subprocess.run(["docker", "rmi", image], check=True)
+        print(f"Image {image} removed successfully.")
+    except subprocess.CalledProcessError as e:
+        print(f"Failed to remove Docker image {image}: {e}")
+
     cmd = [
         "docker",
         "build",
diff --git a/build/rocm/tools/get_rocm.py b/build/rocm/tools/get_rocm.py
index d9cb9ea3f25a..06a31b753f3f 100644
--- a/build/rocm/tools/get_rocm.py
+++ b/build/rocm/tools/get_rocm.py
@@ -75,6 +75,9 @@ def install_packages(self, package_specs):
         if self.pkgbin == "apt":
             env["DEBIAN_FRONTEND"] = "noninteractive"
 
+        # Update indexes.
+        subprocess.check_call(["apt-get", "update"])
+
         LOG.info("Running %r" % cmd)
         subprocess.check_call(cmd, env=env)