add weighting

parthea · parthea · commit bbd81442073b · 2026-07-06T15:47:37.000Z
diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml
@@ -14,11 +14,19 @@ permissions:
 
 # Configurable global environment variables for batching
 env:
-  BATCH_SIZE: 10
   TEST_ALL_PACKAGES: "true" # Set to "false" to only run tests for packages with a git diff
+  ALL_PYTHON: '["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"]'
 
 jobs:
   # Dynamic package discovery job to calculate required matrix size automatically
+  python_config:
+    runs-on: ubuntu-latest
+    outputs:
+      all_python: ${{ steps.export.outputs.python_list }}
+    steps:
+      - id: export
+        run: echo "python_list=${{ env.ALL_PYTHON }}" >> "$GITHUB_OUTPUT"
+
   discover-packages:
     runs-on: ubuntu-latest
     outputs:
@@ -28,54 +36,23 @@ jobs:
       uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6
       with:
         persist-credentials: false
+    - name: Setup Python
+      uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
+      with:
+        python-version: "3.14"
     - name: Generate Batch Indices
       id: set-matrix
       run: |
-        # Testing all monorepo packages sequentially on a single runner node is 
-        # too slow and creates a severe CI bottleneck as the repository expands.
-        #
-        # To scale efficiently, we chunk the workload into parallel batch slices.
-        # Instead of using a fixed, hardcoded matrix array (which risks silently 
-        # skipping newly added packages if the repository outgrows the array size), 
-        # this step dynamically audits the 'packages/' directory at runtime.
-        #
-        # It calculates exactly how many concurrent runners are required based on
-        # the current repository size and the configured BATCH_SIZE variable, 
-        # ensuring 100% test coverage with zero manual YAML maintenance.
-
-        # 1. Count the number of total directories matching the 'packages/*' pattern.
-        #    Redirect stderr to /dev/null so empty repos do not print unneeded errors.
-        TOTAL_PACKAGES=$(ls -d packages/*/ 2>/dev/null | wc -l | tr -d ' ')
-        
-        # 2. Safety fallback: If no packages are detected, assign a single slice index [0]
-        #    so subsequent matrix-dependent jobs do not break or fail validation on an empty matrix.
-        if [ "$TOTAL_PACKAGES" -eq 0 ]; then
-          echo "indices=[0]" >> "$GITHUB_OUTPUT"
-          exit 0
-        fi
-
-        # 3. Calculate the number of batches required using ceiling division: ceil(TOTAL_PACKAGES / BATCH_SIZE).
-        #    The formula ((A + B - 1) / B) ensures integer division rounds up if there's any remaining package leftover.
-        #    Example: 251 packages with a batch size of 10 gives ((251 + 10 - 1) / 10) = 260 / 10 = 26 batches.
-        NUM_BATCHES=$(( (TOTAL_PACKAGES + ${{ env.BATCH_SIZE }} - 1) / ${{ env.BATCH_SIZE }} ))
-        
-        # 4. Generate a zero-indexed sequence from 0 to (NUM_BATCHES - 1).
-        #    Use jq to securely parse the raw numbers and compile them into a compacted JSON array string.
-        #    Example output format: [0,1,2,3,...,25]
-        INDICES=$(seq 0 $((NUM_BATCHES - 1)) | jq -R . | jq -s -c .)
-        
-        # 5. Output the finished JSON string to the GitHub environment outputs pipeline.
-        #    This will safely feed directly into the execution matrix downstream.
+        INDICES=$(python -c 'import sys; sys.path.append("ci"); import get_batches; print(get_batches.get_batch_indices())')
         echo "indices=${INDICES}" >> "$GITHUB_OUTPUT"
 
   unit:
     name: "unit-run (${{ matrix.python }}, Batch ${{ matrix.batch-index }})"
     runs-on: ubuntu-22.04
-    needs: discover-packages
+    needs: [python_config, discover-packages]
     strategy:
       matrix:
-        python: ['3.9', '3.10', "3.11", "3.12", "3.13", "3.14"]
-        # Dynamically scales to fit every package perfectly without hardcoding array indices
+        python: ${{ fromJSON(needs.python_config.outputs.all_python) }}
         batch-index: ${{ fromJson(needs.discover-packages.outputs.batch-indices) }}
     steps:
     - name: Checkout
@@ -95,43 +72,47 @@ jobs:
         python -m pip install --upgrade setuptools pip wheel
         python -m pip install nox
     - name: Run unit tests
+      id: run-tests
       env:
-        COVERAGE_FILE: ${{ github.workspace }}/.coverage-${{ matrix.python }}
+        COVERAGE_FILE: ${{ github.workspace }}/.coverage-${{ matrix.python }}-${{ matrix.batch-index }}
         # Dynamically set BUILD_TYPE to an empty string to skip the diff calculation if TEST_ALL_PACKAGES is true
         BUILD_TYPE: ${{ env.TEST_ALL_PACKAGES == 'true' && '' || 'presubmit' }}
         TARGET_BRANCH: ${{ github.base_ref || github.event.merge_group.base_ref }}
         TEST_TYPE: unit
         PY_VERSION: ${{ matrix.python }}
       run: |
-        # Gather all packages in alphabetical order
-        ALL_PACKAGES=($(ls -d packages/*/ | sort))
-        TOTAL_PACKAGES=${#ALL_PACKAGES[@]}
+        UNIQUE_BATCH_PACKAGES=$(python -c 'import sys; sys.path.append("ci"); import get_batches; print(get_batches.get_batch_slice(${{ matrix.batch-index }}))')
 
-        # Determine this runner's slice window
-        START_INDEX=$(( ${{ matrix.batch-index }} * ${{ env.BATCH_SIZE }} ))
-        
-        if [ $START_INDEX -ge $TOTAL_PACKAGES ]; then
+        if [ -z "$UNIQUE_BATCH_PACKAGES" ]; then
+          echo "No structural units allocated to this matrix slice."
           exit 0
         fi
 
-        BATCH_PACKAGES=("${ALL_PACKAGES[@]:$START_INDEX:${{ env.BATCH_SIZE }}}")
+        echo "Running tests for packages in this weighted batch: ${UNIQUE_BATCH_PACKAGES}"
+        ci/run_conditional_tests.sh ${UNIQUE_BATCH_PACKAGES}
 
-        # Strip trailing slashes to pass down directly into ci/run_conditional_tests.sh
-        subdirs=("${BATCH_PACKAGES[@]%/}")
+    - name: Save Status Footprint
+      run: |
+        mkdir -p footprints
+        echo "${{ steps.run-tests.outcome }}" > footprints/status.txt
+
+    - name: Upload Status Footprint
+      uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
+      with:
+        name: footprint-${{ matrix.python }}-${{ matrix.batch-index }}
+        path: footprints/
 
-        ci/run_conditional_tests.sh "${subdirs[@]}"
     - name: Upload coverage results
       uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
       with:
         # Appended batch-index to separate parallel coverage uploads cleanly
         name: coverage-artifact-${{ matrix.python }}-${{ matrix.batch-index }}
-        path: .coverage-${{ matrix.python }}
+        path: .coverage-${{ matrix.python }}-${{ matrix.batch-index }}
         include-hidden-files: true
 
   cover:
     runs-on: ubuntu-latest
-    needs:
-      - unit
+    needs: [unit]
     steps:
     - name: Checkout
       uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
@@ -246,19 +227,49 @@ jobs:
 
   unittest-runtime-result:
     name: "unit (${{ matrix.python }})"
-    needs: unit
+    needs: [python_config, discover-packages, unit]
+    if: always()
     strategy:
+      fail-fast: false
       matrix:
-        python: ['3.9', '3.10', "3.11", "3.12", "3.13", "3.14"]
+        python: ${{ fromJSON(needs.python_config.outputs.all_python) }}
     runs-on: ubuntu-latest
     steps:
-      - name: Check unit tests results
-        run: |
-          UNIT_STATUS="${{ needs.unit.result }}"
+      - name: Download all status footprints for this runtime
+        uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5
+        with:
+          pattern: footprint-${{ matrix.python }}-*
+          path: footprint-results/
 
-          if [[ "$UNIT_STATUS" == "success" ]]; then
-            echo "Python ${{ matrix.python }} tests passed."
+      - name: Validate complete batch execution footprint status
+        run: |
+          EXPECTED_BATCHES=$(echo '${{ needs.discover-packages.outputs.batch-indices }}' | jq '. | length')
+          
+          if [ -d "footprint-results" ]; then
+            ACTUAL_BATCHES=$(find footprint-results -type f -name 'status.txt' | wc -l | tr -d ' ')
           else
-            echo "Error: Python ${{ matrix.python }} status is '$UNIT_STATUS'."
+            ACTUAL_BATCHES=0
+          fi
+
+          echo "Validation metrics for Python ${{ matrix.python }}:"
+          echo " -> Expected footprint files: $EXPECTED_BATCHES"
+          echo " -> Downloaded footprint files: $ACTUAL_BATCHES"
+
+          if [ "$ACTUAL_BATCHES" -ne "$EXPECTED_BATCHES" ]; then
+            echo "Error: Footprint count mismatch! Expected $EXPECTED_BATCHES files, but found $ACTUAL_BATCHES."
+            exit 1
+          fi
+
+          FAILED_RUNS=0
+          while read -r STATUS_FILE; do
+            RUN_STATUS=$(cat "$STATUS_FILE" | tr -d '[:space:]')
+            if [[ "$RUN_STATUS" != "success" ]]; then
+              echo "Failure detected in batch profile path: $STATUS_FILE (Status: $RUN_STATUS)"
+              FAILED_RUNS=$((FAILED_RUNS + 1))
+            fi
+          done < <(find footprint-results -type f -name 'status.txt' 2>/dev/null)
+
+          if [ "$FAILED_RUNS" -gt 0 ]; then
+            echo "Error: Validation failed. Found $FAILED_RUNS failing test batches."
             exit 1
           fi
diff --git a/ci/get_batches.py b/ci/get_batches.py
@@ -0,0 +1,81 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#!/usr/bin/env python3
+import os
+import glob
+import math
+import json
+import sys
+
+BATCH_SIZE = 10
+
+# Large SDKs take significantly longer to run tests. Assigning higher weights 
+# ensures these packages occupy more runner slots and load-balance effectively.
+PACKAGE_WEIGHTS = {
+    "google-cloud-compute": 5,
+    "google-cloud-compute-v1beta": 5,
+    "google-cloud-dialogflow": 5,
+    "google-cloud-dialogflow-cx": 5,
+    "google-cloud-retail": 5,
+}
+
+def get_weighted_package_list():
+    """Audits the packages directory and expands the list based on weights."""
+    raw_paths = sorted(glob.glob("packages/*"))
+    package_paths = [p for p in raw_paths if os.path.isdir(p)]
+    
+    weighted_list = []
+    for path in package_paths:
+        pkg_name = os.path.basename(path)
+        weight = PACKAGE_WEIGHTS.get(pkg_name, 1)
+        
+        for _ in range(weight):
+            weighted_list.append(path)
+            
+    return weighted_list
+
+def get_batch_indices():
+    """Returns a JSON string of the array of batch indices for GitHub Actions matrix."""
+    weighted_list = get_weighted_package_list()
+    total_units = len(weighted_list)
+    num_batches = math.ceil(total_units / BATCH_SIZE)
+
+    if num_batches == 0:
+        num_batches = 1
+
+    return json.dumps(list(range(num_batches)))
+
+def get_batch_slice(batch_index):
+    """Returns a space-separated string of unique package paths for a specific batch index."""
+    weighted_list = get_weighted_package_list()
+    total_units = len(weighted_list)
+    start_idx = batch_index * BATCH_SIZE
+    
+    if start_idx >= total_units:
+        return ""
+        
+    slice_window = weighted_list[start_idx : start_idx + BATCH_SIZE]
+    
+    unique_packages = []
+    for pkg in slice_window:
+        if pkg not in unique_packages:
+            unique_packages.append(pkg)
+            
+    return " ".join(unique_packages)
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1 and sys.argv[1] == "--count":
+        print(get_batch_indices())
+    elif len(sys.argv) > 2 and sys.argv[1] == "--slice":
+        print(get_batch_slice(int(sys.argv[2])))
