Merge branch 'main' into vkozlov/fix-levanter-tests-3

Author: yhtang

.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

.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
 

.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

.github/workflows/_test_unit.yaml

@@ -19,14 +19,18 @@ on:
         type: string
         description: 'Test artifacts to collect'
         required: false
+      TIMEOUT_MINUTES:
+        type: number
+        description: 'Maximum test runtime, in minutes'
+        default: 60
 
 jobs:
   runner:
     uses: ./.github/workflows/_runner_ondemand_slurm.yaml
     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 }}