[ROCm] ci build and dockerfile changes

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

build/rocm/ci_build

@@ -47,6 +47,14 @@ def dist_wheels(
     # create manylinux image with requested ROCm installed
     image = "jax-manylinux_2_28_x86_64_rocm%s" % rocm_version.replace(".", "")
 
+    # Remove 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}")
+        exit(1)
+
     cmd = [
         "docker",
         "build",

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)