Prepare for plugin 0.8.2 release

Author: Ruturaj4

.github/workflows/rocm-perf.yml

@@ -170,4 +170,4 @@ jobs:
             --python-version "${{ needs.build-and-test-jax-perf.outputs.python_version }}" \
             --rocm-version "${{ needs.build-and-test-jax-perf.outputs.rocm_version }}" \
             --gfx-version gfx90a \
-            --jax-version 0.8.0
+            --jax-version 0.8.2

BUILDING.md

@@ -136,7 +136,7 @@ python3 build/ci_build test $TEST_IMAGE --test-cmd "pytest jax_rocm_plugin/tests
 We keep unit tests in the `rocm/jax` repository, and you'll need to clone it
 to run the regular JAX unit tests with ROCm,
 ```shell
-git clone --depth 1 --branch rocm-jaxlib-v0.8.0 [email protected]:ROCm/jax.git
+git clone --depth 1 --branch rocm-jaxlib-v0.8.2 [email protected]:ROCm/jax.git
 # Each release of the ROCm plugin has a corresponding branch. You can find
 # more at https://github.com/ROCm/rocm-jax/branches/all?query=rocm-jaxlib
 

build/ci_build

@@ -114,6 +114,7 @@ def dist_wheels(
     rocm_version,
     python_versions,
     xla_source_dir,
+    jax_source_dir,
     rocm_build_job="",
     rocm_build_num="",
     therock_path=None,
@@ -137,37 +138,17 @@ def dist_wheels(
         xla_path = os.path.realpath(os.path.expanduser(xla_source_dir))
         cmd.append("--xla-source-dir=%s" % xla_path)
 
+    if jax_source_dir:
+        jax_path = os.path.realpath(os.path.expanduser(jax_source_dir))
+        cmd.append("--jax-source-dir=%s" % jax_path)
+
     if rbe:
         cmd.append("--rbe")
 
     cmd.append("dist_wheels")
     subprocess.check_call(cmd, cwd=jax_plugin_dir)
 
 
-def _fetch_jax_metadata(xla_path):
-    cmd = ["git", "rev-parse", "HEAD"]
-    jax_commit = subprocess.check_output(cmd)
-    xla_commit = b""
-
-    if xla_path:
-        try:
-            xla_commit = subprocess.check_output(cmd, cwd=xla_path)
-        except Exception as ex:
-            LOG.warning("Exception while retrieving xla_commit: %s" % ex)
-
-    cmd = ["python3", "setup.py", "-V"]
-    env = dict(os.environ)
-    env["JAX_RELEASE"] = "1"
-
-    jax_version = subprocess.check_output(cmd, env=env)
-
-    return {
-        "jax_version": jax_version.decode("utf8").strip(),
-        "jax_commit": jax_commit.decode("utf8").strip(),
-        "xla_commit": xla_commit.decode("utf8").strip(),
-    }
-
-
 def _apply_filters(docker_filters, dockerfile_basename, docker_dir="docker"):
     """
     Collect Dockerfile paths matching a basename prefix and optional substring filters.
@@ -540,7 +521,12 @@ def parse_args():
 
     p.add_argument(
         "--xla-source-dir",
-        help="Path to XLA source to use during jaxlib build, instead of builtin XLA",
+        help="Path to XLA source to use during plugin and jaxlib build, instead of builtin XLA",
+    )
+
+    p.add_argument(
+        "--jax-source-dir",
+        help="Optional JAX source directory. When provided, builds jaxlib wheel and copies to wheelhouse.",
     )
 
     p.add_argument(
@@ -623,6 +609,7 @@ def main():
             rocm_version=args.rocm_version,
             python_versions=args.python_versions,
             xla_source_dir=args.xla_source_dir,
+            jax_source_dir=args.jax_source_dir,
             rocm_build_job=args.rocm_build_job,
             rocm_build_num=args.rocm_build_num,
             therock_path=args.therock_path,

build/requirements.txt

@@ -1,2 +1,2 @@
-jaxlib==0.8.0
-jax==0.8.0
+jaxlib==0.8.2
+jax==0.8.2

ci/Dockerfile.maxtext

@@ -15,4 +15,4 @@ WORKDIR /maxtext
 # Explicitly install jax,jaxlib to avoid pip pulling a newer version (e.g. 0.8.1)
 RUN --mount=type=cache,mode=0755,target=/root/.cache/pip \
     pip install -r requirements.txt && \
-    pip install jax==0.8.0 jaxlib==0.8.0 && pip freeze
+    pip install jax==0.8.2 jaxlib==0.8.2 && pip freeze

ci/jax_rbe/pr_test.sh

@@ -42,7 +42,7 @@ python3 build/build.py build --wheels=jax-rocm-plugin --configure_only --python_
     --config=rocm_rbe \
     --noremote_accept_cached \
     --//jax:build_jaxlib=false \
-    --action_env=TF_ROCM_AMDGPU_TARGETS="gfx906,gfx908,gfx90a,gfx942,gfx950,gfx1030,gfx1100,gfx1101,gfx1200,gfx1201" \
+    --action_env=TF_ROCM_AMDGPU_TARGETS="gfx908,gfx90a,gfx942,gfx950,gfx1030,gfx1100,gfx1101,gfx1200,gfx1201" \
     --test_verbose_timeout_warnings \
     --test_output=errors \
     //tests:core_test_gpu \

docker/Dockerfile.base-ubu22

@@ -4,7 +4,7 @@ FROM ubuntu:22.04
 # The Python version to use to build JAX ROCm plugin and pjrt.
 ARG PY_VERSION=3.11.13
 # The list of target devices to be supported by the JAX ROCm plugin and pjrt.
-ARG GPU_DEVICE_TARGETS="gfx906 gfx908 gfx90a gfx942 gfx950 gfx1030 gfx1100 gfx1101 gfx1200 gfx1201"
+ARG GPU_DEVICE_TARGETS="gfx908 gfx90a gfx942 gfx950 gfx1030 gfx1100 gfx1101 gfx1200 gfx1201"
 # The ROCm version to be used inside the container.
 ARG ROCM_VERSION
 # The installation path for ROCm.

docker/Dockerfile.base-ubu24

@@ -2,7 +2,7 @@ FROM ubuntu:24.04
 
 ### Container Build Arguments:
 # The list of target devices to be supported by the JAX ROCm plugin and pjrt.
-ARG GPU_DEVICE_TARGETS="gfx906 gfx908 gfx90a gfx942 gfx950 gfx1030 gfx1100 gfx1101 gfx1200 gfx1201"
+ARG GPU_DEVICE_TARGETS="gfx908 gfx90a gfx942 gfx950 gfx1030 gfx1100 gfx1101 gfx1200 gfx1201"
 # The ROCm version to be used inside the container.
 ARG ROCM_VERSION
 # The installation path for ROCm.

docker/Dockerfile.jax-ubu22

@@ -56,4 +56,5 @@ LABEL com.amdgpu.jax_version="$JAX_VERSION" \
 RUN --mount=type=cache,mode=0755,target=/root/.cache/pip \
     --mount=type=bind,source=wheelhouse,target=/wheelhouse \
     ls -lah /wheelhouse && \
-    /usr/local/bin/pip3 install -f /wheelhouse --no-deps --no-index "jax_rocm${PLUGIN_NAMESPACE}_plugin" "jax_rocm${PLUGIN_NAMESPACE}_pjrt"
+    /usr/local/bin/pip3 install -f /wheelhouse --no-deps --no-index "jax_rocm${PLUGIN_NAMESPACE}_plugin" "jax_rocm${PLUGIN_NAMESPACE}_pjrt" && \
+    /usr/local/bin/pip3 install -f /wheelhouse --no-deps --no-index --force-reinstall "jaxlib"

docker/Dockerfile.jax-ubu24

@@ -55,4 +55,5 @@ LABEL com.amdgpu.jax_version="$JAX_VERSION" \
 RUN --mount=type=cache,mode=0755,target=/root/.cache/pip \
     --mount=type=bind,source=wheelhouse,target=/wheelhouse \
     ls -lah /wheelhouse && \
-    pip3 install -f /wheelhouse --no-deps --no-index "jax_rocm${PLUGIN_NAMESPACE}_plugin" "jax_rocm${PLUGIN_NAMESPACE}_pjrt"
+    pip3 install -f /wheelhouse --no-deps --no-index "jax_rocm${PLUGIN_NAMESPACE}_plugin" "jax_rocm${PLUGIN_NAMESPACE}_pjrt" && \
+    pip3 install -f /wheelhouse --no-deps --no-index --force-reinstall "jaxlib"