Update benchmarks.yml

Author: juliagmt-google

.github/workflows/benchmarks.yml

@@ -48,41 +48,42 @@ jobs:
       #   working-directory: xla
       #   continue-on-error: true
       #   run: bazel test -c opt --config=cuda //xla/service/gpu/tests:gpu_atomic_test
-  xla-gpu-ci:
-    runs-on: linux-x86-g2-48-l4-4gpu  # Or another suitable GPU runner
+
+  build-xla-gpu:
+    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"
+      image: nvidia/cuda:12.0-base-ubuntu20.04-devel  # Choose appropriate CUDA version and base image
+      options: --gpus all --privileged  # Might need privileged mode, use with caution
 
     steps:
-      - name: Checkout XLA with submodules
+      - name: Checkout XLA
         uses: actions/checkout@v3
         with:
-          repository: openxla/xla
+          repository: openxla/xla  # Replace with your fork if needed
           path: xla
-          submodules: recursive
-          
-      - name: Install Docker
-        run: |
-          apt-get update -y
-          apt-get install -y \
-            ca-certificates \
-            curl \
-            gnupg \
-            lsb-release
 
-          ## Add this line to import the Docker GPG key:
-          curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
+      - name: Install dependencies
+        working-directory: xla
+        run: |
+          sudo apt-get update -y
+          sudo apt-get install -y python3 python3-dev python3-pip  # Add other required packages
 
-          echo \
-            "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu \
-            $(lsb_release -cs) stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null
+      - name: Install Python dependencies
+        working-directory: xla
+        run: |
+          python3 -m pip install --upgrade pip
+          pip3 install wheel setuptools numpy  # Add other required Python packages
 
-          apt-get update -y
-          apt-get install -y docker-ce docker-ce-cli containerd.io
-          # Add any necessary Docker configuration here, if required
+      - name: Configure Bazel
+        working-directory: xla
+        run: |
+          echo "build --action_env PYTHON_BIN_PATH=\"/usr/bin/python3\"" >> ~/.bazelrc
+          echo "build --python_path=\"/usr/bin/python3\"" >> ~/.bazelrc
+          # Add other Bazel configuration options if needed
 
-      - name: Pull Docker image
-        run: docker run --name xla_gpu -w /xla -it -d --rm -v $PWD:/xla tensorflow/build:latest-python3.9 bash
+      - name: Build XLA with GPU support
+        working-directory: xla
+        run: bazel build -c opt --config=cuda //xla/...