Merge branch 'ci-improvements-20260702' into chore-update-librarian-2026-06-26

Author: parthea

.github/workflows/gapic-generator-tests.yml

@@ -135,11 +135,13 @@ jobs:
       - name: Install System Deps
         run: sudo apt-get update && sudo apt-get install -y pandoc
       - name: Run Goldens
+        env:
+          LATEST_STABLE_PYTHON: ${{ needs.python_config.outputs.latest_stable_python }}
         run: |
           pip install nox
           cd packages/gapic-generator
           for pkg in credentials eventarc logging redis; do
-            nox -f tests/integration/goldens/$pkg/noxfile.py -s format lint unit-${{ needs.python_config.outputs.latest_stable_python }}
+            nox -f tests/integration/goldens/$pkg/noxfile.py -s format lint unit-${LATEST_STABLE_PYTHON}
           done
       # Run pylint (errors-only) over the goldens so generator regressions
       # like undefined names or import-time breakage that ruff/flake8 do
@@ -214,11 +216,12 @@ jobs:
           # Run fragment for current matrix python
           nox -s fragment-${MATRIX_PYTHON}
           # Run snippetgen only on the latest stable to avoid the "Python not found" error
-          if [ "${MATRIX_PYTHON}" == "${{ needs.python_config.outputs.latest_stable_python }}" ]; then
+          if [ "${MATRIX_PYTHON}" == "${LATEST_STABLE_PYTHON}" ]; then
             nox -s snippetgen
           fi
         env:
           MATRIX_PYTHON: ${{ matrix.python }}
+          LATEST_STABLE_PYTHON: ${{ needs.python_config.outputs.latest_stable_python }}
   integration:
     needs: python_config
     # Only runs if the Gatekeeper passed

.github/workflows/unittest.yml

@@ -12,14 +12,48 @@ name: unittest
 permissions:
   contents: read
 
+# Configurable global environment variables for batching
+env:
+  TEST_ALL_PACKAGES: "false" # 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', '3.15']"
 
 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:
+      batch-indices: ${{ steps.set-matrix.outputs.indices }}
+    steps:
+    - name: Checkout
+      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: |
+        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: [python_config, discover-packages]
     strategy:
-      fail-fast: true
       matrix:
-        python: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
+        python: ${{ fromJSON(needs.python_config.outputs.all_python) }}
+        batch-index: ${{ fromJson(needs.discover-packages.outputs.batch-indices) }}
     steps:
     - name: Checkout
       uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6
@@ -54,9 +88,11 @@ jobs:
         restore-keys: |
           ${{ runner.os }}-uv-wheels-3.15-
     - name: Run unit tests
+      id: run-tests
       env:
-        COVERAGE_FILE: ${{ github.workspace }}/.coverage-${{ matrix.python }}
-        BUILD_TYPE: presubmit
+        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 }}
@@ -65,6 +101,15 @@ jobs:
         # Follow https://github.com/Instagram/LibCST/issues/1445 for updates.
         PYO3_USE_ABI3_FORWARD_COMPATIBILITY: "1"
       run: |
+        UNIQUE_BATCH_PACKAGES=$(python -c 'import sys; sys.path.append("ci"); import get_batches; print(get_batches.get_batch_slice(${{ matrix.batch-index }}))')
+
+        if [ -z "$UNIQUE_BATCH_PACKAGES" ]; then
+          echo "No structural units allocated to this matrix slice."
+          exit 0
+        fi
+
+        echo "Running tests for packages in this weighted batch: ${UNIQUE_BATCH_PACKAGES}"
+
         # Only inject overrides if we are strictly on the 3.15 pre-release
         # This is needed to speed up builds
         if [ "${{ matrix.python }}" = "3.15" ]; then
@@ -74,18 +119,31 @@ jobs:
           # Route uv's wheel and environment cache to our persistent workspace directory
           export UV_CACHE_DIR="${{ github.workspace }}/.uv-wheel-cache"
         fi
-        ci/run_conditional_tests.sh
+
+        ci/run_conditional_tests.sh ${UNIQUE_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/
+
     - name: Upload coverage results
       uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
       with:
-        name: coverage-artifact-${{ matrix.python }}
-        path: .coverage-${{ matrix.python }}
+        # Appended batch-index to separate parallel coverage uploads cleanly
+        name: coverage-artifact-${{ matrix.python }}-${{ matrix.batch-index }}
+        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
@@ -197,3 +255,52 @@ jobs:
           echo "This usually means the unit tests did not run or failed to upload their coverage files."
           exit 1
         fi
+
+  unittest-runtime-result:
+    name: "unit (${{ matrix.python }})"
+    needs: [python_config, discover-packages, unit]
+    if: always()
+    strategy:
+      fail-fast: false
+      matrix:
+        python: ${{ fromJSON(needs.python_config.outputs.all_python) }}
+    runs-on: ubuntu-latest
+    steps:
+      - name: Download all status footprints for this runtime
+        uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5
+        with:
+          pattern: footprint-${{ matrix.python }}-*
+          path: footprint-results/
+
+      - 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
+            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

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
+
+# Packages that take significantly longer to run tests.
+# These will always be assigned to their own dedicated runner.
+ISOLATED_PACKAGES = [
+    "google-cloud-compute",
+    "google-cloud-compute-v1beta",
+    "google-cloud-dialogflow",
+    "google-cloud-dialogflow-cx",
+    "google-cloud-retail",
+]
+
+def get_batches():
+    """Splits packages into dedicated isolated batches and evenly chunked standard batches."""
+    raw_paths = sorted(glob.glob("packages/*"))
+    package_paths = [p for p in raw_paths if os.path.isdir(p)]
+    
+    batches = []
+    standard_packages = []
+    
+    for path in package_paths:
+        pkg_name = os.path.basename(path)
+        
+        if pkg_name in ISOLATED_PACKAGES:
+            # Put heavy packages into their own standalone batches immediately
+            batches.append([path])
+        else:
+            standard_packages.append(path)
+            
+    # Chunk the remaining standard packages by BATCH_SIZE
+    num_standard_packages = len(standard_packages)
+    num_standard_batches = math.ceil(num_standard_packages / BATCH_SIZE)
+    
+    if num_standard_batches == 0 and not batches:
+        num_standard_batches = 1
+        
+    for i in range(num_standard_batches):
+        start_idx = i * BATCH_SIZE
+        chunk = standard_packages[start_idx : start_idx + BATCH_SIZE]
+        if chunk:
+            batches.append(chunk)
+            
+    return batches
+
+def get_batch_indices():
+    """Returns a JSON string of the array of batch indices for GitHub Actions matrix."""
+    batches = get_batches()
+    return json.dumps(list(range(len(batches))))
+
+def get_batch_slice(batch_index):
+    """Returns a space-separated string of unique package paths for a specific batch index."""
+    batches = get_batches()
+    if batch_index < 0 or batch_index >= len(batches):
+        return ""
+    return " ".join(batches[batch_index])
+
+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])))

ci/run_conditional_tests.sh

@@ -21,6 +21,10 @@
 
 # `TEST_TYPE` and `PY_VERSION` are required by the script `ci/run_single_test.sh`
 
+# Optional Arguments:
+# Pass specific space-separated package paths (e.g., "packages/google-cloud-storage") to only test those directories.
+# If no arguments are provided, the script automatically determines which directories have changed
+#
 # This script will determine which directories have changed
 # under the `packages` folder. For `BUILD_TYPE=="presubmit"`,
 # we'll compare against the `packages` folder in HEAD,
@@ -78,16 +82,43 @@ set -e
 # Now we have a fixed list, but we can change it to autodetect if
 # necessary.
 
-subdirs=(
-    packages
-)
+if [ $# -gt 0 ]; then
+    subdirs=("$@")
+else
+    subdirs=(
+        packages
+    )
+fi
 
 RETVAL=0
 
-for subdir in ${subdirs[@]}; do
-    for d in `ls -d ${subdir}/*/`; do
+for subdir in "${subdirs[@]}"; do
+    if [ ! -d "${subdir}" ]; then
+        echo "Error: Directory '${subdir}' does not exist." >&2
+        exit 1
+    fi
+
+    if [[ "${subdir%/}" == "packages" ]]; then
+        loop_dirs=("${subdir}"/*/)
+    else
+        loop_dirs=("${subdir}")
+    fi
+
+    for d in "${loop_dirs[@]}"; do
+        if [ ! -d "$d" ]; then
+            continue
+        fi
+        # Ensure the directory path always ends with a trailing slash for git diff safety
+        if [[ "$d" != */ ]]; then
+            d="$d/"
+        fi
         should_test=false
-        if [ -n "${GIT_DIFF_ARG}" ]; then
+        
+        # Override check: Force test if explicitly asked to test all packages
+        if [[ "${TEST_ALL_PACKAGES}" == "true" ]]; then
+            echo "TEST_ALL_PACKAGES is true, forcing execution for ${d}"
+            should_test=true
+        elif [ -n "${GIT_DIFF_ARG}" ]; then
             echo "checking changes with 'git diff --quiet ${GIT_DIFF_ARG} ${d}'"
             set +e
             git diff --quiet ${GIT_DIFF_ARG} ${d}
@@ -119,4 +150,4 @@ for subdir in ${subdirs[@]}; do
     done
 done
 
-exit ${RETVAL}
+exit ${RETVAL}

packages/bigframes/bigframes/ml/llm.py

@@ -59,6 +59,8 @@
 _GEMINI_2P5_PRO_ENDPOINT = "gemini-2.5-pro"
 _GEMINI_2P5_FLASH_ENDPOINT = "gemini-2.5-flash"
 _GEMINI_2P5_FLASH_LITE_ENDPOINT = "gemini-2.5-flash-lite"
+_GEMINI_3P1_FLASH_LITE_ENDPOINT = "gemini-3.1-flash-lite"
+_GEMINI_3P5_FLASH_ENDPOINT = "gemini-3.5-flash"
 
 _GEMINI_ENDPOINTS = (
     _GEMINI_1P5_PRO_PREVIEW_ENDPOINT,
@@ -73,6 +75,8 @@
     _GEMINI_2P5_PRO_ENDPOINT,
     _GEMINI_2P5_FLASH_ENDPOINT,
     _GEMINI_2P5_FLASH_LITE_ENDPOINT,
+    _GEMINI_3P1_FLASH_LITE_ENDPOINT,
+    _GEMINI_3P5_FLASH_ENDPOINT,
 )
 _GEMINI_PREVIEW_ENDPOINTS = (
     _GEMINI_1P5_PRO_PREVIEW_ENDPOINT,
@@ -96,6 +100,8 @@
     _GEMINI_2P5_PRO_ENDPOINT,
     _GEMINI_2P5_FLASH_ENDPOINT,
     _GEMINI_2P5_FLASH_LITE_ENDPOINT,
+    _GEMINI_3P1_FLASH_LITE_ENDPOINT,
+    _GEMINI_3P5_FLASH_ENDPOINT,
 )
 
 _CLAUDE_3_SONNET_ENDPOINT = "claude-3-sonnet"
@@ -440,7 +446,8 @@ class GeminiTextGenerator(base.RetriableRemotePredictor):
             "gemini-1.5-pro-001", "gemini-1.5-pro-002", "gemini-1.5-flash-001",
             "gemini-1.5-flash-002", "gemini-2.0-flash-exp",
             "gemini-2.0-flash-lite-001", "gemini-2.0-flash-001",
-            "gemini-2.5-pro", "gemini-2.5-flash" and "gemini-2.5-flash-lite".
+            "gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-lite",
+            "gemini-3.1-flash-lite" and "gemini-3.5-flash".
             If no setting is provided, "gemini-2.0-flash-001" will be used by
             default and a warning will be issued.
 
@@ -478,6 +485,8 @@ def __init__(
                 "gemini-2.5-pro",
                 "gemini-2.5-flash",
                 "gemini-2.5-flash-lite",
+                "gemini-3.1-flash-lite",
+                "gemini-3.5-flash",
             ]
         ] = None,
         session: Optional[bigframes.Session] = None,