Update benchmarks.yml

Author: juliagmt-google

.github/workflows/benchmarks.yml

@@ -50,72 +50,60 @@ jobs:
       - name: Check TF_CPP_MAX_VLOG_LEVEL
         working-directory: xla
         run: echo "$TF_CPP_MAX_VLOG_LEVEL"
-      - name: Wait For Connection
-        uses: google-ml-infra/actions/ci_connection@main
-        with:
-          halt-dispatch-input: ${{ inputs.halt-for-connection }}
-      - name: Wait For Connection
-        uses: google-ml-infra/actions/ci_connection@main
-        with:
-          halt-dispatch-input: ${{ inputs.halt-for-connection }}
 
-      # - name: Build hlo_runner_main
-      #   working-directory: xla
-      #   run: bazel build -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main
+      - name: Build hlo_runner_main
+        working-directory: xla
+        run: bazel build -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main
 
       # - name: Wait For Connection
       #   uses: google-ml-infra/actions/ci_connection@main
       #   with:
       #     halt-dispatch-input: ${{ inputs.halt-for-connection }}
 
-      # - name: Create gpu_hlo_backend.hlo
-      #   working-directory: xla
-      #   run: |
-      #     cat << EOF > gpu_hlo_backend.hlo
-      #     HloModule module
-      #     // CHECK: is_scheduled=true
+      - name: Create gpu_hlo_backend.hlo
+        working-directory: xla
+        run: |
+          cat << EOF > gpu_hlo_backend.hlo
+          HloModule module
+          // CHECK: is_scheduled=true
           
-      #     ENTRY computation {
-      #         p = f32[5000,6000]{1,0} parameter(0)
-      #         e = f32[5000,6000]{1,0} sqrt(p)
-      #         c = f32[6000,5000] transpose(p), dimensions={1,0}
-      #         r = f32[300,20,5000] reshape(c)
-      #         ROOT out = (f32[5000,6000], f32[300,20,5000]) tuple(e,r)
-      #     }
-      #     EOF
+          ENTRY computation {
+              p = f32[5000,6000]{1,0} parameter(0)
+              e = f32[5000,6000]{1,0} sqrt(p)
+              c = f32[6000,5000] transpose(p), dimensions={1,0}
+              r = f32[300,20,5000] reshape(c)
+              ROOT out = (f32[5000,6000], f32[300,20,5000]) tuple(e,r)
+          }
+          EOF
           
-      # - name: Wait For Connection
-      #   uses: google-ml-infra/actions/ci_connection@main
-      #   with:
-      #     halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@main
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
 
-      # - name: Run an HLO file
-      #   working-directory: xla
-      #   run: |
-      #     nvidia-smi --query-gpu=utilization.gpu --format=csv -l 1 > results/gpu_utilization_gpu_backend.log & ./bazel-bin/xla/tools/multihost_hlo_runner/hlo_runner_main --device_type=gpu --log_output=True --use_spmd_partitioning gpu_hlo_backend.hlo
-
-      #     cd results
-      #     ls
-      #     cat gpu_hlo_backend.hlo.log
+      - name: Run an HLO file
+        working-directory: xla
+        run: |
+          ./bazel-bin/xla/tools/multihost_hlo_runner/hlo_runner_main --device_type=gpu --log_output=True --use_spmd_partitioning gpu_hlo_backend.hlo  &> results/gpu_hlo_backend.log
     
-      # - name: Wait For Connection
-      #   uses: google-ml-infra/actions/ci_connection@main
-      #   with:
-      #     halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@main
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
 
-      # - name: Download parse_xla_logs.py
-      #   working-directory: xla
-      #   run: wget https://raw.githubusercontent.com/juliagmt-google/xla/main/.github/workflows/parse_xla_logs.py  
+      - name: Download parse_xla_logs.py
+        working-directory: xla
+        run: wget https://raw.githubusercontent.com/juliagmt-google/xla/main/.github/workflows/parse_xla_logs.py  
 
-      # - name: Parse XLA logs
-      #   working-directory: xla
-      #   run: python parse_xla_logs.py results/gpu_hlo_backend.hlo.log
+      - name: Parse XLA logs
+        working-directory: xla
+        run: python parse_xla_logs.py results/gpu_hlo_backend.hlo.log
 
-      # - name: Upload Results
-      #   uses: actions/upload-artifact@v4
-      #   with:
-      #     name: gpu-xla-benchmarks
-      #     path: xla/results
+      - name: Upload Results
+        uses: actions/upload-artifact@v4
+        with:
+          name: gpu-xla-benchmarks
+          path: xla/results
   # # jax-build-and-test:
   # #   runs-on: linux-x86-g2-48-l4-4gpu # Use a GPU-enabled runner
   # #   container: