Update benchmarks.yml

Author: juliagmt-google

.github/workflows/benchmarks.yml

@@ -16,33 +16,33 @@ on:
         - 'no'
 
 jobs:
-  jax-build-and-test:
-    runs-on: linux-x86-g2-48-l4-4gpu # Use a GPU-enabled runner
-    container:
-      image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
+  # jax-build-and-test:
+  #   runs-on: linux-x86-g2-48-l4-4gpu # Use a GPU-enabled runner
+  #   container:
+  #     image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
 
-    env:
-      JAXCI_HERMETIC_PYTHON_VERSION: 3.11
+  #   env:
+  #     JAXCI_HERMETIC_PYTHON_VERSION: 3.11
 
-    steps:
-      - name: Checkout JAX Fork
-        uses: actions/checkout@v3
-        with:
-          repository: 'google-ml-infra/jax-fork'
-          path: jax-fork
+  #   steps:
+  #     - name: Checkout JAX Fork
+  #       uses: actions/checkout@v3
+  #       with:
+  #         repository: 'google-ml-infra/jax-fork'
+  #         path: jax-fork
 
-      - name: Install JAX Dependencies
-        working-directory: jax-fork
-        run: |
-          python -m pip install --upgrade pip
-          pip install pytest
-          pip install absl-py
-          pip install "jax[cuda12_pip]"  # Adjust CUDA version if needed
-          pip install google-benchmark
-      - name: Run JAX Multiprocess GPU Test
-        working-directory: jax-fork
-        continue-on-error: true
-        run: python -m pytest tests/multiprocess_gpu_test.py
+  #     - name: Install JAX Dependencies
+  #       working-directory: jax-fork
+  #       run: |
+  #         python -m pip install --upgrade pip
+  #         pip install pytest
+  #         pip install absl-py
+  #         pip install "jax[cuda12_pip]"  # Adjust CUDA version if needed
+  #         pip install google-benchmark
+  #     - name: Run JAX Multiprocess GPU Test
+  #       working-directory: jax-fork
+  #       continue-on-error: true
+  #       run: python -m pytest tests/multiprocess_gpu_test.py
 
   build-xla-gpu-and-test:
     runs-on: linux-x86-g2-48-l4-4gpu # Use a GPU-enabled runner