chore: add sharding to docs test

.github/workflows/docs.yml

@@ -15,9 +15,81 @@ name: docs
 permissions:
   contents: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+env:
+  NOX_REUSE_EXISTING_VIRTUALENVS: "true"
+  NOX_ENVDIR: "/tmp/shared_nox_envs"
+  NOX_DEFAULT_VENV_BACKEND: "uv"
+  UV_VENV_SEED: "1"
+  # Run checks in parallel using 4 cores
+  PARALLEL_WORKERS: "4"
+
 jobs:
-  docs:
+  initialize:
     runs-on: ubuntu-latest
+    outputs:
+      matrix: ${{ steps.set-matrix.outputs.matrix }}
+      is_full_run: ${{ steps.check-label.outputs.is_full_run }}
+    env:
+      MAX_SHARDS: 5
+      MIN_SHARD_WEIGHT: 25
+      # Define weights for packages with large autodoc workloads to balance shard execution time
+      # Default for unset packages is 1
+      PACKAGE_WEIGHTS: |
+        google-ads-admanager: 9
+        google-cloud-compute: 50
+        google-cloud-compute-v1beta: 100
+        google-cloud-dialogflow: 9
+        google-cloud-dialogflow-cx: 9
+        google-cloud-discoveryengine: 15
+        google-cloud-network-security: 9
+        google-cloud-retail: 9
+    steps:
+    - name: Checkout
+      uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
+      with:
+        fetch-depth: 0
+        filter: blob:none
+        persist-credentials: false
+    - name: Check for unit_test:all_packages label
+      id: check-label
+      run: |
+        if [[ "${{ contains(github.event.pull_request.labels.*.name, 'unit_test:all_packages') }}" == "true" ]]; then
+          echo "is_full_run=true" >> $GITHUB_OUTPUT
+        else
+          echo "is_full_run=false" >> $GITHUB_OUTPUT
+        fi
+    - name: Setup Python
+      uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
+      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 || github.ref_name }}
+        TEST_ALL_PACKAGES: ${{ steps.check-label.outputs.is_full_run }}
+        MAX_SHARDS: ${{ env.MAX_SHARDS }}
+        MIN_SHARD_WEIGHT: ${{ env.MIN_SHARD_WEIGHT }}
+        PACKAGE_WEIGHTS: ${{ env.PACKAGE_WEIGHTS }}
+      run: |
+        if [ -n "$TARGET_BRANCH" ]; then
+          git fetch --no-tags --quiet origin "$TARGET_BRANCH" || true
+        fi
+        python3 ci/get_package_shards.py
+
+  docs-shard:
+    needs: initialize
+    if: needs.initialize.outputs.matrix != '[]' && needs.initialize.outputs.matrix != ''
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: true
+      matrix:
+        package_shard: ${{ fromJson(needs.initialize.outputs.matrix) }}
+    name: ${{ matrix.package_shard.is_sharded && format('docs ({0})', matrix.package_shard.name) || format('docs ({0})', matrix.package_shard.description) }}
     steps:
     - name: Checkout
       uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
@@ -34,16 +106,36 @@ jobs:
     - name: Install nox
       run: |
         python -m pip install --upgrade setuptools pip wheel
-        python -m pip install nox
-    - name: Run docs
+        python -m pip install nox uv
+    - name: Run docs for ${{ matrix.package_shard.description }}
       env:
         BUILD_TYPE: presubmit
         TARGET_BRANCH: ${{ github.base_ref || github.event.merge_group.base_ref || github.ref_name }}
         TEST_TYPE: docs
+        PACKAGE_LIST: ${{ matrix.package_shard.packages }}
         # TODO(https://github.com/googleapis/google-cloud-python/issues/13775): Specify `PY_VERSION` rather than relying on the default python version of the nox session.
         PY_VERSION: "unused"
       run: |
         ci/run_conditional_tests.sh
+
+  docs:
+    if: always()
+    needs: [initialize, docs-shard]
+    runs-on: ubuntu-latest
+    name: docs
+    steps:
+    - name: Check docs job status
+      run: |
+        if [[ "${{ needs.initialize.result }}" != "success" ]]; then
+          echo "Error: The initialize job status was: ${{ needs.initialize.result }}"
+          exit 1
+        fi
+        if [[ "${{ needs['docs-shard'].result }}" != "success" && "${{ needs['docs-shard'].result }}" != "skipped" ]]; then
+          echo "docs failed with result: ${{ needs['docs-shard'].result }}"
+          exit 1
+        fi
+        echo "All docs shards passed or were skipped!"
+
   docfx:
     if: github.event_name == 'push' && github.ref == 'refs/heads/main'
     runs-on: ubuntu-latest
@@ -63,7 +155,7 @@ jobs:
     - name: Install nox
       run: |
         python -m pip install --upgrade setuptools pip wheel
-        python -m pip install nox
+        python -m pip install nox uv
     - name: Run docfx
       env:
         BUILD_TYPE: presubmit

ci/get_package_shards.py

@@ -228,7 +228,8 @@ def group_packages(packages_map):
 
     # Dynamically determine target weight to balance across max shards.
     max_shards = int(os.environ.get("MAX_SHARDS", 16))
-    target_weight = max(10, math.ceil(total_weight / max_shards))
+    min_shard_weight = int(os.environ.get("MIN_SHARD_WEIGHT", 10))
+    target_weight = max(min_shard_weight, math.ceil(total_weight / max_shards))
 
     shards_list = []
     current_shard_items = []