chore(ci): use get_package_shards.py in import-profiler for dynamic sharding (#18079)

### Summary
Updates the `import-profiler` workflow to use `ci/get_package_shards.py`
for dynamic sharding matrix generation, bringing it in line with the
sharding pattern used in `lint.yml` and `unittest.yml`.

### Changes
* **`initialize` Job**: Added an initialization job running
`ci/get_package_shards.py` with `MAX_SHARDS: 8` to dynamically calculate
shard distribution based on modified packages.
* **Dynamic Shards**: Replaced static 8-shard matrix (`[0..7]`) with
dynamic shard matrix evaluation. If few or no packages changed, only the
necessary number of shards are spawned instead of always spinning up 8
runners.
* **Code Cleanliness**: Removed manual git diff parsing and modulo
arithmetic from the workflow step, passing `PACKAGE_LIST` directly to
`ci/run_conditional_tests.sh`.
* **Workflow Status Check**: Updated `all-import-profiles` to depend on
`[initialize, import-profile]` and verify both job outcomes.

Author: hebaalazzeh

.github/workflows/import-profiler.yml

@@ -13,19 +13,62 @@ permissions:
   contents: read
 
 jobs:
+  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: 8
+    steps:
+    - name: Checkout
+      uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
+      # Use a fetch-depth of 2 to avoid error `fatal: origin/main...HEAD: no merge base`
+      # See https://github.com/googleapis/google-cloud-python/issues/12013
+      # and https://github.com/actions/checkout#checkout-head.
+      with:
+        fetch-depth: 2
+        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 }}
+        TEST_ALL_PACKAGES: ${{ steps.check-label.outputs.is_full_run }}
+        MAX_SHARDS: ${{ env.MAX_SHARDS }}
+      run: |
+        if [ -n "$TARGET_BRANCH" ]; then
+          git fetch origin "$TARGET_BRANCH" --depth=1 || true
+        fi
+        python3 ci/get_package_shards.py
+
   import-profile:
+    needs: initialize
+    if: needs.initialize.outputs.matrix != '[]' && needs.initialize.outputs.matrix != ''
     runs-on: ubuntu-latest
     timeout-minutes: 60
     strategy:
       fail-fast: false
       matrix:
-        shard: [0, 1, 2, 3, 4, 5, 6, 7] # 8 parallel shards
-    name: import-profile (Shard ${{ matrix.shard }})
+        package_shard: ${{ fromJson(needs.initialize.outputs.matrix) }}
+    name: ${{ matrix.package_shard.is_sharded && format('import-profile ({0})', matrix.package_shard.name) || format('import-profile ({0})', matrix.package_shard.description) }}
     steps:
     - name: Checkout
-      uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6
+      uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
       with:
-        fetch-depth: 0 # Fetch git history to find changed packages
+        fetch-depth: 0 # Fetch git history to find changed packages / baseline
         persist-credentials: false
     - name: Setup Python
       uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
@@ -38,51 +81,30 @@ jobs:
         python -m pip install --upgrade pip
         pip install pytest pytest-cov setuptools
         pytest scripts/import_profiler/test_profiler.py --cov=profiler --cov-report=term-missing --cov-fail-under=100
-    - name: Run import profiler
+    - name: Run import profiler for ${{ matrix.package_shard.description }}
       env:
         BUILD_TYPE: presubmit
         TARGET_BRANCH: ${{ github.base_ref || github.event.merge_group.base_ref }}
         TEST_TYPE: import_profile
         PY_VERSION: "3.15"
         # Workaround: Allows libcst to compile on Python 3.15+ while PyO3 catches up
         PYO3_USE_ABI3_FORWARD_COMPATIBILITY: "1"
-        SHARD_INDEX: ${{ matrix.shard }}
-        TOTAL_SHARDS: 8
+        PACKAGE_LIST: ${{ matrix.package_shard.packages }}
       run: |
-        TARGET_BRANCH=${TARGET_BRANCH:-main}
-        git fetch origin "${TARGET_BRANCH}" --deepen=200 || true
-        
-        # Get unique list of modified packages under packages/
-        modified_packages=$(git diff --name-only origin/"${TARGET_BRANCH}"... | grep '^packages/' | cut -d/ -f1,2 | sort -u)
-        
-        # Filter packages assigned to this specific shard index
-        idx=0
-        packages_to_test=""
-        for pkg in $modified_packages; do
-          if [ -d "$pkg" ]; then
-            if [ "$((idx % TOTAL_SHARDS))" -eq "${SHARD_INDEX}" ]; then
-              packages_to_test="$packages_to_test $pkg"
-            fi
-            idx=$((idx + 1))
-          fi
-        done
-        
-        # Run tests on the assigned packages
-        if [ -n "$packages_to_test" ]; then
-          echo "Shard ${{ matrix.shard }} running packages: $packages_to_test"
-          PACKAGE_LIST="$packages_to_test" ci/run_conditional_tests.sh
-        else
-          echo "No packages assigned to Shard ${{ matrix.shard }}."
-        fi
+        ci/run_conditional_tests.sh
 
   all-import-profiles:
-    needs: import-profile
+    needs: [initialize, import-profile]
     if: always()
     runs-on: ubuntu-latest
     steps:
       - name: Check import profile results
         run: |
-          if [[ "${{ needs.import-profile.result }}" != "success" && "${{ needs.import-profile.result }}" != "skipped" ]]; then
+          if [[ "${{ needs.initialize.result }}" != "success" ]]; then
+            echo "Error: The initialize job status was: ${{ needs.initialize.result }}"
+            exit 1
+          fi
+          if [[ "${{ needs['import-profile'].result }}" != "success" && "${{ needs['import-profile'].result }}" != "skipped" ]]; then
             echo "Import profiles failed"
             exit 1
           fi

scripts/import_profiler/profiler.py

@@ -385,7 +385,9 @@ def find_module_from_package(pkg):
     try:
         files = importlib.metadata.files(pkg)
         if files:
-            ignored_parts = ('tests', 'testing', 'samples', 'examples', 'benchmark', 'benchmarks')
+            ignored_parts = {'tests', 'testing', 'samples', 'examples', 'benchmark', 'benchmarks', 'third_party', 'test_utils', 'docs', 'build', 'dist', 'bin', 'ci', 'scripts', 'cloudbuild', 'notebooks', 'assets', 'scratch', 'specs'}
+            if pkg == "google-cloud-testutils":
+                ignored_parts.discard('test_utils')
             init_files = [str(f) for f in files if str(f).endswith('__init__.py') and '__pycache__' not in str(f) and not any(part in ignored_parts for part in str(f).replace('\\', '/').split('/'))]
             if init_files:
                 from pathlib import Path
@@ -406,13 +408,20 @@ def find_module_from_package(pkg):
         import os
         if os.path.exists('setup.py') or os.path.exists('pyproject.toml'):
             where_dir = "src" if os.path.isdir("src") else "."
+            abs_where_dir = os.path.abspath(where_dir)
+            if abs_where_dir not in sys.path:
+                sys.path.insert(0, abs_where_dir)
             pkgs = setuptools.find_namespace_packages(where=where_dir)
             ignored_prefixes = ("tests", "samples", "examples", "benchmark", "benchmarks", "third_party", "testing", "test_utils", "docs", "build", "dist", "bin", "ci", "scripts", "cloudbuild", "notebooks", "assets", "scratch", "specs")
+            ignored_starts = ("test_", "tests_", "sample_", "samples_", "bench_", "benchmarks_", "example_", "examples_", "doc_", "docs_", "notebook_", "notebooks_")
 
             filtered = []
             for p in pkgs:
                 top = p.split(".")[0]
-                if top in ignored_prefixes or top.startswith(("test_", "sample_", "bench_", "example_", "doc_", "notebook_")) or p in ("google", "google.cloud"):
+                is_ignored_top = top in ignored_prefixes or top.startswith(ignored_starts)
+                if is_ignored_top and pkg == "google-cloud-testutils" and top == "test_utils":
+                    is_ignored_top = False
+                if is_ignored_top or p in ("google", "google.cloud"):
                     continue
                 filtered.append(p)
 

scripts/import_profiler/test_profiler.py

@@ -641,6 +641,13 @@ def test_find_module_from_package_metadata_init():
         assert res == "foo.bar"
 
 
+def test_find_module_from_package_metadata_test_utils():
+    with patch("importlib.metadata.files", return_value=["test_utils/__init__.py"]), \
+         patch("importlib.util.find_spec", return_value=True):
+        res = find_module_from_package("google-cloud-testutils")
+        assert res == "test_utils"
+
+
 def test_find_module_from_package_setuptools():
     sys.modules.setdefault("setuptools", MagicMock())
     with patch("importlib.metadata.files", side_effect=Exception), \
@@ -652,6 +659,19 @@ def test_find_module_from_package_setuptools():
         assert res == "my_pkg"
 
 
+def test_find_module_from_package_setuptools_test_utils():
+    sys.modules.setdefault("setuptools", MagicMock())
+    with patch("importlib.metadata.files", side_effect=Exception), \
+         patch("profiler.os.path.exists", return_value=True), \
+         patch("profiler.os.path.isdir", return_value=True), \
+         patch("setuptools.find_namespace_packages", return_value=["test_utils", "tests"]) as mock_find, \
+         patch("profiler.os.path.isfile", return_value=True), \
+         patch("importlib.util.find_spec", return_value=True):
+        res = find_module_from_package("google-cloud-testutils")
+        assert res == "test_utils"
+        mock_find.assert_called_once_with(where="src")
+
+
 def test_find_module_from_package_setuptools_not_file_and_exception():
     sys.modules.setdefault("setuptools", MagicMock())
     def mock_isfile(path):