[DRAFT] experiment: add test sharding

.github/workflows/unittest.yml

@@ -13,11 +13,40 @@ permissions:
   contents: read
 
 jobs:
+  initialize:
+    runs-on: ubuntu-latest
+    outputs:
+      matrix: ${{ steps.set-matrix.outputs.matrix }}
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+    - name: Setup Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.10"
+    - name: Get package shards
+      id: set-matrix
+      env:
+        BUILD_TYPE: presubmit
+        TARGET_BRANCH: ${{ github.base_ref || github.event.merge_group.base_ref }}
+      run: |
+        if [ -n "$TARGET_BRANCH" ]; then
+          git fetch origin "$TARGET_BRANCH" --deepen=200 || true
+        fi
+        echo "matrix=$(python3 ci/get_package_shards.py)" >> $GITHUB_OUTPUT
+
   unit:
+    needs: initialize
+    if: needs.initialize.outputs.matrix != '[]' && needs.initialize.outputs.matrix != ''
     runs-on: ubuntu-22.04
     strategy:
+      fail-fast: true
       matrix:
         python: ['3.9', '3.10', "3.11", "3.12", "3.13", "3.14"]
+        package_shard: ${{ fromJson(needs.initialize.outputs.matrix) }}
+    name: ${{ matrix.package_shard.is_sharded && format('unit ({0}, {1})', matrix.python, matrix.package_shard.name) || format('unit ({0})', matrix.python) }}
     steps:
     - name: Checkout
       uses: actions/checkout@v4
@@ -34,20 +63,40 @@ jobs:
       run: |
         python -m pip install --upgrade setuptools pip wheel
         python -m pip install nox
-    - name: Run unit tests
+    - name: Run unit tests for ${{ matrix.package_shard.description }}
       env:
-        COVERAGE_FILE: .coverage-${{ matrix.python }}
         BUILD_TYPE: presubmit
         TARGET_BRANCH: ${{ github.base_ref || github.event.merge_group.base_ref }}
         TEST_TYPE: unit
         PY_VERSION: ${{ matrix.python }}
+        PACKAGE_LIST: ${{ matrix.package_shard.packages }}
       run: |
         ci/run_conditional_tests.sh
     - name: Upload coverage results
       uses: actions/upload-artifact@v4
       with:
-        name: coverage-artifact-${{ '{{' }} matrix.python {{ '}}' }}
-        path: .coverage-${{ matrix.python }}
+        name: coverage-artifact-${{ matrix.python }}-${{ matrix.package_shard.index }}
+        path: .coverage.${{ matrix.python }}.*
+        include-hidden-files: true
+
+  all-tests:
+    needs: [initialize, unit]
+    if: always()
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check unit test results
+        run: |
+          # 1. Check initialize job
+          if [[ "${{ needs.initialize.result }}" != "success" ]]; then
+            echo "Error: The initialize job status was: ${{ needs.initialize.result }}"
+            exit 1
+          fi
+          # 2. Check unit test shards
+          if [[ "${{ needs.unit.result }}" != "success" && "${{ needs.unit.result }}" != "skipped" ]]; then
+            echo "Unit tests failed"
+            exit 1
+          fi
+          echo "All unit tests passed or were skipped"
 
   cover:
     runs-on: ubuntu-latest
@@ -67,20 +116,84 @@ jobs:
         python-version: "3.10"
     - name: Set number of files changes in packages directory
       id: packages
-      run: echo "::set-output name=num_files_changed::$(git diff HEAD~1 -- packages | wc -l)"
+      run: echo "num_files_changed=$(git diff HEAD~1 -- packages | wc -l)" >> $GITHUB_OUTPUT
     - name: Install coverage
-      if: steps.packages.num_files_changed > 0
+      if: steps.packages.outputs.num_files_changed > 0
       run: |
         python -m pip install --upgrade setuptools pip wheel
         python -m pip install coverage
     - name: Download coverage results
-      if: ${{ steps.date.packages.num_files_changed > 0 }}
+      if: steps.packages.outputs.num_files_changed > 0
       uses: actions/download-artifact@v4
       with:
         path: .coverage-results/
     - name: Report coverage results
-      if: ${{ steps.date.packages.num_files_changed > 0 }}
+      if: steps.packages.outputs.num_files_changed > 0
+      env:
+        # TODO: default to 100% coverage after next gapic-generator release
+        # https://github.com/googleapis/google-cloud-python/issues/17459
+        DEFAULT_FAIL_UNDER: 99
       run: |
-        find .coverage-results -type f -name '*.zip' -exec unzip {} \;
-        coverage combine .coverage-results/**/.coverage*
-        coverage report --show-missing --fail-under=100
+        if [ -d .coverage-results ]; then
+          find .coverage-results -type f -name '.coverage*' -exec mv {} . \;
+          coverage combine
+
+          # Find all modified packages
+          modified_packages=$(git diff --name-only HEAD~1 -- packages | cut -d/ -f1,2 | sort -u)
+
+          failed_packages=()
+          passed_packages=()
+
+          for pkg in ${modified_packages}; do
+            if [ -d "${pkg}" ]; then
+              echo "============================================================"
+              echo "Evaluating coverage for package: ${pkg}"
+              echo "============================================================"
+
+              set +e
+              if [ -f "${pkg}/.coveragerc" ]; then
+                echo "Using package-specific configuration: ${pkg}/.coveragerc"
+                # If fail_under is specified in the package-specific .coveragerc, coverage report
+                # will automatically enforce it. Otherwise, we enforce the default.
+                if grep -q "fail_under" "${pkg}/.coveragerc"; then
+                  coverage report --rcfile="${pkg}/.coveragerc" --include="${pkg}/**"
+                else
+                  echo "No fail_under specified in ${pkg}/.coveragerc, enforcing default"
+                  coverage report --rcfile="${pkg}/.coveragerc" --include="${pkg}/**" --fail-under="${DEFAULT_FAIL_UNDER}"
+                fi
+              else
+                echo "No .coveragerc found for ${pkg}, enforcing default"
+                coverage report --include="${pkg}/**" --fail-under="${DEFAULT_FAIL_UNDER}"
+              fi
+              status=$?
+              set -e
+
+              if [ ${status} -ne 0 ]; then
+                failed_packages+=("${pkg}")
+              else
+                passed_packages+=("${pkg}")
+              fi
+            fi
+          done
+
+          echo "============================================================"
+          echo "Coverage Evaluation Summary"
+          echo "============================================================"
+          if [ ${#passed_packages[@]} -gt 0 ]; then
+            echo "Passed packages:"
+            for pkg in "${passed_packages[@]}"; do
+              echo "  - ${pkg}"
+            done
+          fi
+          if [ ${#failed_packages[@]} -gt 0 ]; then
+            echo "Failed packages:"
+            for pkg in "${failed_packages[@]}"; do
+              echo "  - ${pkg}"
+            done
+            exit 1
+          fi
+        else
+          echo "Error: No coverage results were downloaded from the unit test jobs."
+          echo "This usually means the unit tests did not run or failed to upload their coverage files."
+          exit 1
+        fi

.kokoro/continuous/common.cfg

@@ -4,6 +4,7 @@
 action {
   define_artifacts {
     regex: "**/*sponge_log.xml"
+    regex: "**/*sponge_log.log"
   }
 }
 

.kokoro/presubmit/common.cfg

@@ -4,6 +4,7 @@
 action {
   define_artifacts {
     regex: "**/*sponge_log.xml"
+    regex: "**/*sponge_log.log"
   }
 }
 

.kokoro/system-single.sh

@@ -32,4 +32,10 @@ NOX_FILE_ARG=""
 
 [[ -z "${NOX_FILE}" ]] || NOX_FILE_ARG="-f ${NOX_FILE}"
 
-python3 -m nox ${NOX_SESSION_ARG} $NOX_FILE_ARG
+NOX_ARGS=()
+if [[ -n "${XML_OUTPUT_FILE}" ]]; then
+  # Pass --junitxml as a positional argument to the underlying pytest call
+  NOX_ARGS+=("--junitxml=${XML_OUTPUT_FILE}")
+fi
+
+python3 -m nox ${NOX_SESSION_ARG} $NOX_FILE_ARG -- "${NOX_ARGS[@]}"

.kokoro/system.sh

@@ -38,7 +38,18 @@ pwd
 run_package_test() {
   local package_name=$1
   local package_path="packages/${package_name}"
-
+
+  # Create a dedicated directory for this package's artifacts
+  local artifact_dir="${KOKORO_ARTIFACTS_DIR}/${package_name}"
+  mkdir -p "${artifact_dir}"
+
+  # Define standard output paths for logs and XML results
+  export XML_OUTPUT_FILE="${artifact_dir}/sponge_log.xml"
+  local log_file="${artifact_dir}/sponge_log.log"
+
+  # Isolate gcloud config for parallel execution
+  export CLOUDSDK_CONFIG=$(mktemp -d)
+
   # Declare local overrides to prevent bleeding into the next loop iteration
   local PROJECT_ID
   local GOOGLE_APPLICATION_CREDENTIALS
@@ -48,6 +59,8 @@ run_package_test() {
 
   echo "------------------------------------------------------------"
   echo "Configuring environment for: ${package_name}"
+  echo "Log file: ${log_file}"
+  echo "XML results: ${XML_OUTPUT_FILE}"
   echo "------------------------------------------------------------"
 
   case "${package_name}" in
@@ -82,17 +95,86 @@ run_package_test() {
   # Run the actual test
   pushd "${package_path}" > /dev/null
   set +e
-  "${system_test_script}"
+  # *** CRITICAL: system-single.sh MUST use XML_OUTPUT_FILE to generate the JUnit XML ***
+  "${system_test_script}" > "${log_file}" 2>&1
   local res=$?
   set -e
   popd > /dev/null
-
+
+  # Clean up isolated gcloud config
+  rm -rf "${CLOUDSDK_CONFIG}"
+
   return $res
 }
 
 # A file for running system tests
 system_test_script="${PROJECT_ROOT}/.kokoro/system-single.sh"
+# Ensure KOKORO_ARTIFACTS_DIR is set, fallback to a local directory
+export KOKORO_ARTIFACTS_DIR="${KOKORO_ARTIFACTS_DIR:-${PROJECT_ROOT}/.artifacts}"
+
+# Parallel execution settings
+MAX_PARALLEL=4
+running_pids=()
+declare -A pid_to_pkg
+declare -A pid_to_resfile
+
+# Array to keep track of results for the final summary
+results=()
+handle_finished_job() {
+  local pid=$1
+  local pkg=${pid_to_pkg[$pid]}
+  local resfile=${pid_to_resfile[$pid]}
+  local pkg_log="${KOKORO_ARTIFACTS_DIR}/${pkg}/sponge_log.log"
+
+  # wait $pid might fail if it was already reaped by wait -n, 
+  # so we ignore its exit code and use the resfile.
+  wait "$pid" 2>/dev/null || true
+
+  local res=1
+  if [ -s "$resfile" ]; then
+    res=$(cat "$resfile")
+    rm "$resfile"
+  else
+    # If the file is empty or missing, the subshell crashed early
+    res=1
+  fi
+
+  echo "============================================================"
+  echo "System tests for ${pkg} finished (Exit code: ${res})"
+  echo "Artifacts in: ${KOKORO_ARTIFACTS_DIR}/${pkg}"
+  echo "============================================================"
 
+  # Print the package log to the console
+  if [ -f "${pkg_log}" ]; then
+    echo "--- Start of Logs for ${pkg} ---"
+    cat "${pkg_log}"
+    echo "--- End of Logs for ${pkg} ---"
+  else
+    echo "Console Fallback Warning: Log file not found at ${pkg_log}"
+  fi
+  echo ""
+
+  if [ -z "${res}" ] || [ "${res}" -ne 0 ]; then
+    echo "============================================================"
+    echo "FAIL-FAST: System tests for ${pkg} failed!"
+    echo "Cancelling all remaining running background jobs..."
+    echo "============================================================"
+
+    # Kill all other active background processes
+    for active_pid in "${running_pids[@]}"; do
+      if [ "$active_pid" != "$pid" ] && kill -0 "$active_pid" 2>/dev/null; then
+        pkg_to_cancel=${pid_to_pkg[$active_pid]}
+        echo "Cancelling active system tests for ${pkg_to_cancel} (PID: ${active_pid})..."
+        # Send SIGTERM to allow graceful cleanup of resources if possible, or SIGKILL
+        kill -9 "$active_pid" 2>/dev/null || true
+      fi
+    done
+
+    exit "${res}"
+  else
+    results+=("${pkg}: PASSED")
+  fi
+}
 # Run system tests for each package with directory packages/*/tests/system
 for path in `find 'packages' \
   \( -type d -wholename 'packages/*/tests/system' \) -o \
@@ -140,10 +222,55 @@ for path in `find 'packages' \
   set -e
 
   if [[ "${package_modified}" -gt 0 || "$KOKORO_BUILD_ARTIFACTS_SUBDIR" == *"continuous"* ]]; then
-      # Call the function - its internal exports won't affect the next loop
-      run_package_test "$package_name" || RETVAL=$?
+      # Wait if we have reached MAX_PARALLEL
+      while [[ ${#running_pids[@]} -ge $MAX_PARALLEL ]]; do
+        set +e
+        wait -n
+        set -e
+        # Find which job finished
+        new_pids=()
+        for pid in "${running_pids[@]}"; do
+          if kill -0 "$pid" 2>/dev/null; then
+            new_pids+=("$pid")
+          else
+            handle_finished_job "$pid"
+          fi
+        done
+        running_pids=("${new_pids[@]}")
+      done
+
+      # Start the next test in the background
+      res_file=$(mktemp)
+      (
+        if run_package_test "$package_name"; then
+          echo 0 > "$res_file"
+        else
+          res=$?
+          echo $res > "$res_file"
+        fi
+      ) &
+      pid=$!
+      running_pids+=($pid)
+      pid_to_pkg[$pid]=$package_name
+      pid_to_resfile[$pid]=$res_file
+      echo "Started system tests for ${package_name} (PID: ${pid})"
   else
       echo "No changes in ${package_name} and not a continuous build, skipping."
   fi
 done
+
+# Wait for all remaining jobs
+for pid in "${running_pids[@]}"; do
+  handle_finished_job "$pid"
+done
+
+echo "------------------------------------------------------------"
+echo "System Test Summary"
+echo "------------------------------------------------------------"
+for res in "${results[@]}"; do
+  echo "$res"
+done
+echo "------------------------------------------------------------"
+
 exit ${RETVAL}
+