From 032e978cda22c98c6a7b46e833e346d4a5677ccf Mon Sep 17 00:00:00 2001
From: Brian Yang <125406446+gpupuck@users.noreply.github.com>
Date: Thu, 3 Oct 2024 10:46:54 -0700
Subject: [PATCH 1/3] Do not add leading "//tests:" for given test targets
 (#1079)

The original logic is if the given test targets do not start with
`//tests:`, then add this prefix for them.

I'm removing this logic as:
1. We don't actually see any use cases in the workflow
2. It will block us running per-directory suite targets such as
`//tests/mosaic:gpu_tests`
---
 .github/container/test-jax.sh | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/.github/container/test-jax.sh b/.github/container/test-jax.sh
index 80ad5b02f..6afbaace1 100755
--- a/.github/container/test-jax.sh
+++ b/.github/container/test-jax.sh
@@ -109,9 +109,6 @@ else
 fi
 
 for t in $*; do
-    if [[ "$t" != "//tests:"* ]]; then
-        t="//tests:${t}"
-    fi
     BAZEL_TARGET="${BAZEL_TARGET} $t"
 done
 

From d1ff3c88d3ad84c78969232a33fe0a8469a58b49 Mon Sep 17 00:00:00 2001
From: Anna Shors <71393111+ashors1@users.noreply.github.com>
Date: Tue, 8 Oct 2024 08:07:17 -0700
Subject: [PATCH 2/3] Pin etils to fix error with `tfds.load` (#1086)

Signed-off-by: ashors1 <ashors@nvidia.com>
---
 .github/container/Dockerfile.pax.amd64 | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/.github/container/Dockerfile.pax.amd64 b/.github/container/Dockerfile.pax.amd64
index ede475a10..52a7723ab 100644
--- a/.github/container/Dockerfile.pax.amd64
+++ b/.github/container/Dockerfile.pax.amd64
@@ -29,6 +29,9 @@ for src in ${SRC_PATH_PAXML} ${SRC_PATH_PRAXIS}; do
   pushd ${src}
   sed -i "s| @ git+https://github.com/google/flax||g" requirements.in
   sed -i "s| @ git+https://github.com/google/jax||g" requirements.in
+  ## we pin etils because newer etils versions are not compatible with the
+  ## version of TFDS required by Pax
+  sed -i "s/etils/etils==1.7.0/g" requirements.in
   if git diff --quiet; then
       echo "URL specs no longer present in select dependencies for ${src}"
       exit 1

From e8043a5b50bc3b82cf8f91d20f3e2150ce8b01c5 Mon Sep 17 00:00:00 2001
From: Olli Lupton <olupton@nvidia.com>
Date: Thu, 10 Oct 2024 16:18:46 +0200
Subject: [PATCH 3/3] Improve TE test reporting, increase timeout (#1084)

- Increase timeout as A100 tests take >1 hour to run
- Report the test failures/errors, not collection failures/errors
---
 .github/workflows/_ci.yaml        | 5 +++--
 .github/workflows/_test_unit.yaml | 7 ++++++-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml
index 426764323..17ad5f06b 100644
--- a/.github/workflows/_ci.yaml
+++ b/.github/workflows/_ci.yaml
@@ -528,13 +528,14 @@ jobs:
       STATISTICS_SCRIPT: |
         summary_line=$(tail -n1 test-te.log)
         errors=$(echo $summary_line | grep -oE '[0-9]+ error' | awk '{print $1} END { if (!NR) print 0}')
-        passed_tests=$(cat pytest-report.jsonl | jq -r 'select(."$report_type" == "CollectReport" and .outcome == "passed") | .outcome' | wc -l)
-        failed_tests=$(cat pytest-report.jsonl | jq -r 'select(."$report_type" == "CollectReport" and .outcome == "failed") | .outcome' | wc -l)          
+        passed_tests=$(cat pytest-report.jsonl | jq -r 'select(."$report_type" == "TestReport" and .when == "call" and .outcome == "passed") | .outcome' | wc -l)
+        failed_tests=$(cat pytest-report.jsonl | jq -r 'select(."$report_type" == "TestReport" and .when == "call" and .outcome == "failed") | .outcome' | wc -l)
         total_tests=$((failed_tests + passed_tests))
         echo "TOTAL_TESTS=${total_tests}" >> $GITHUB_OUTPUT
         echo "ERRORS=${errors}" >> $GITHUB_OUTPUT
         echo "PASSED_TESTS=${passed_tests}" >> $GITHUB_OUTPUT
         echo "FAILED_TESTS=${failed_tests}" >> $GITHUB_OUTPUT
+      TIMEOUT_MINUTES: 120
       ARTIFACTS: |
         test-te.log
         pytest-report.jsonl
diff --git a/.github/workflows/_test_unit.yaml b/.github/workflows/_test_unit.yaml
index 2edfb2249..fa29557e0 100644
--- a/.github/workflows/_test_unit.yaml
+++ b/.github/workflows/_test_unit.yaml
@@ -19,6 +19,10 @@ on:
         type: string
         description: 'Test artifacts to collect'
         required: false
+      TIMEOUT_MINUTES:
+        type: number
+        description: 'Maximum test runtime, in minutes'
+        default: 60
 
 jobs:
   runner:
@@ -26,7 +30,7 @@ jobs:
     with:
       NAME: "A100"
       LABELS: "A100,${{ github.run_id }}"
-      TIME: "01:00:00"
+      TIME: "${{ inputs.TIMEOUT_MINUTES }}:00"
     secrets: inherit
 
   run-unit-test:
@@ -67,6 +71,7 @@ jobs:
       - name: Run tests
         shell: bash -x -e {0}
         continue-on-error: true
+        timeout-minutes: ${{ inputs.TIMEOUT_MINUTES }}
         run: |
           ${{ inputs.EXECUTE }}