chore: add test sharding to unit tests

.github/workflows/unittest.yml

@@ -13,11 +13,40 @@ permissions:
   contents: read
 
 jobs:
+  discover:
+    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
+        fi
+        echo "matrix=$(python ci/get_package_shards.py)" >> $GITHUB_OUTPUT
+
   unit:
+    needs: discover
+    if: needs.discover.outputs.matrix != '[]' && needs.discover.outputs.matrix != ''
     runs-on: ubuntu-22.04
     strategy:
+      fail-fast: false
       matrix:
         python: ['3.9', '3.10', "3.11", "3.12", "3.13", "3.14"]
+        package_shard: ${{ fromJson(needs.discover.outputs.matrix) }}
+    name: unit (${{ matrix.package_shard.name }}, ${{ matrix.python }})
     steps:
     - name: Checkout
       uses: actions/checkout@v4
@@ -36,17 +65,18 @@ jobs:
         python -m pip install nox
     - name: Run unit tests
       env:
-        COVERAGE_FILE: .coverage-${{ matrix.python }}
+        COVERAGE_FILE: ${{ github.workspace }}/.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 {{ '}}' }}
+        name: coverage-artifact-${{ matrix.python }}-${{ matrix.package_shard.index }}
         path: .coverage-${{ matrix.python }}
 
   cover:
@@ -67,20 +97,19 @@ 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
       run: |
-        find .coverage-results -type f -name '*.zip' -exec unzip {} \;
         coverage combine .coverage-results/**/.coverage*
         coverage report --show-missing --fail-under=100

.kokoro/system.sh

@@ -39,6 +39,9 @@ run_package_test() {
   local package_name=$1
   local package_path="packages/${package_name}"
 
+  # 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
@@ -87,12 +90,52 @@ run_package_test() {
   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"
 
+# Parallel execution settings
+MAX_PARALLEL=4
+running_pids=()
+declare -A pid_to_pkg
+declare -A pid_to_log
+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 log=${pid_to_log[$pid]}
+  local resfile=${pid_to_resfile[$pid]}
+
+  # 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=$(cat "$resfile")
+  rm "$resfile"
+
+  echo "------------------------------------------------------------"
+  echo "System tests for ${pkg} finished (Exit code: ${res})"
+  echo "------------------------------------------------------------"
+  cat "$log"
+  rm "$log"
+
+  if [ "${res}" -ne 0 ]; then
+    RETVAL=${res}
+    results+=("${pkg}: FAILED")
+  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 +183,51 @@ 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
+        wait -n
+        # 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
+      log_file=$(mktemp)
+      res_file=$(mktemp)
+      (
+        run_package_test "$package_name" > "$log_file" 2>&1
+        echo $? > "$res_file"
+      ) &
+      pid=$!
+      running_pids+=($pid)
+      pid_to_pkg[$pid]=$package_name
+      pid_to_log[$pid]=$log_file
+      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}
+

ci/get_package_shards.py

@@ -0,0 +1,94 @@
+import os
+import subprocess
+import json
+import math
+import sys
+
+def get_packages():
+    subdirs = ['packages']
+    packages = []
+    for subdir in subdirs:
+        if not os.path.exists(subdir):
+            continue
+        # Use the same sorting as the shell script
+        pkg_dirs = [os.path.join(subdir, d) + '/' for d in os.listdir(subdir) if os.path.isdir(os.path.join(subdir, d))]
+        packages.extend(sorted(pkg_dirs))
+    return packages
+
+def get_packages_to_test():
+    build_type = os.environ.get('BUILD_TYPE', 'presubmit')
+    target_branch = os.environ.get('TARGET_BRANCH', 'main')
+
+    all_packages = get_packages()
+
+    if build_type == 'presubmit':
+        git_diff_arg = f"origin/{target_branch}..."
+    elif build_type == 'continuous':
+        git_diff_arg = "HEAD~.."
+    else:
+        return all_packages
+
+    # Check if ci/ changed
+    try:
+        subprocess.check_call(['git', 'diff', '--quiet', git_diff_arg, 'ci'])
+        ci_changed = False
+    except subprocess.CalledProcessError:
+        ci_changed = True
+
+    if ci_changed:
+        return all_packages
+
+    try:
+        res = subprocess.check_output(['git', 'diff', '--name-only', git_diff_arg]).decode('utf-8')
+        changed_files = res.splitlines()
+    except subprocess.CalledProcessError:
+        return all_packages
+
+    to_test = []
+    for pkg in all_packages:
+        # Check if any changed file starts with the package path
+        if any(f.startswith(pkg) for f in changed_files):
+            to_test.append(pkg)
+
+    return to_test
+
+def group_packages(packages, max_packages_per_shard=25, max_total_shards=20):
+    if not packages:
+        return []
+
+    num_packages = len(packages)
+
+    # Calculate number of shards based on packages per shard
+    num_shards = math.ceil(num_packages / max_packages_per_shard)
+
+    # Cap the total number of shards
+    num_shards = min(num_shards, max_total_shards)
+
+    # Recalculate shard size to be as even as possible given the capped shards
+    shard_size = math.ceil(num_packages / num_shards)
+
+    shards = []
+    for i in range(num_shards):
+        start = i * shard_size
+        end = min((i + 1) * shard_size, num_packages)
+        if start >= num_packages:
+            break
+
+        shard_packages = packages[start:end]
+        if len(shard_packages) == 1:
+            name = shard_packages[0].strip('/').split('/')[-1]
+        else:
+            name = f"{shard_packages[0].strip('/').split('/')[-1]}...{shard_packages[-1].strip('/').split('/')[-1]}"
+
+        shards.append({
+            "name": name,
+            "index": i + 1,
+            "packages": " ".join(shard_packages)
+        })
+    return shards
+
+if __name__ == "__main__":
+    packages = get_packages_to_test()
+    # Shard into groups of ~25 libraries, up to 20 parallel jobs
+    shards = group_packages(packages, max_packages_per_shard=25, max_total_shards=20)
+    print(json.dumps(shards))

ci/run_conditional_tests.sh

@@ -82,10 +82,42 @@ subdirs=(
     packages
 )
 
+if [ -n "${PACKAGE_LIST}" ]; then
+    echo "Using provided PACKAGE_LIST"
+    to_test=(${PACKAGE_LIST})
+    RETVAL=0
+    for d in ${to_test[@]}; do
+        echo "running test in ${d}"
+        pushd ${d}
+        set +e
+        ${test_script}
+        ret=$?
+        set -e
+        if [ ${ret} -ne 0 ]; then
+            RETVAL=${ret}
+        fi
+        popd
+    done
+    exit ${RETVAL}
+fi
+
+# Sharding logic (fallback for manual runs)
+TOTAL_SHARDS="${TOTAL_SHARDS:-1}"
+SHARD_INDEX="${SHARD_INDEX:-1}"
+count=0
+
 RETVAL=0
 
 for subdir in ${subdirs[@]}; do
-    for d in `ls -d ${subdir}/*/`; do
+    # Sort the directories to ensure consistent sharding across jobs
+    for d in `ls -d ${subdir}/*/ | sort`; do
+        # Sharding logic: only process directories that belong to this shard
+        if (( count % TOTAL_SHARDS != SHARD_INDEX - 1 )); then
+            ((count++))
+            continue
+        fi
+        ((count++))
+
         should_test=false
         if [ -n "${GIT_DIFF_ARG}" ]; then
             echo "checking changes with 'git diff --quiet ${GIT_DIFF_ARG} ${d}'"

packages/django-google-spanner/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-appoptimize/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-artifact-registry/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-asset/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-assured-workloads/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-audit-log/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-auditmanager/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-automl/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-backupdr/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-bare-metal-solution/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-batch/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-beyondcorp-appconnections/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-beyondcorp-appconnectors/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-beyondcorp-appgateways/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-beyondcorp-clientconnectorservices/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-beyondcorp-clientgateways/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-biglake-hive/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-biglake/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-bigquery-analyticshub/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-bigquery-biglake/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-bigquery-connection/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-bigquery-data-exchange/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-bigquery-datapolicies/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-bigquery-datatransfer/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-bigquery-logging/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-bigquery-migration/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-bigquery-reservation/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-bigquery-storage/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding