address PR comments for import profiler CI

Author: hebaalazzeh

.github/workflows/import-profiler.yml

@@ -31,7 +31,7 @@ jobs:
         python -m pip install nox
     - name: Run import profiler
       env:
-        BUILD_TYPE: presubmit
+        BUILD_TYPE: ${{ contains(github.event.pull_request.labels.*.name, 'import_profile:all_packages') && 'all' || 'presubmit' }}
         TARGET_BRANCH: ${{ github.base_ref || github.event.merge_group.base_ref }}
         TEST_TYPE: import_profile
         PY_VERSION: "3.15"

ci/run_single_test.sh

@@ -106,27 +106,13 @@ case ${TEST_TYPE} in
         esac
         ;;
     import_profile)
-        python -m pip install .
-        PKG_NAME=$(python setup.py --name)
-        MODULE=$(python -c "
-import sys, importlib.util
-pkg = '${PKG_NAME}'
-candidates = [
-    pkg.replace('-', '.'),
-    '.'.join(pkg.split('-')[:-1]) + '_' + pkg.split('-')[-1] if '-' in pkg else pkg,
-    pkg.replace('-', '_')
-]
-for mod in candidates:
-    try:
-        if importlib.util.find_spec(mod):
-            print(mod)
-            sys.exit(0)
-    except Exception:
-        pass
-print(candidates[0])
-")
-        python ${PROJECT_ROOT}/scripts/import_profiler/profiler.py --module "${MODULE}" --iterations 10
-        retval=$?
+        if nox --list-sessions | grep -q "import_profile"; then
+            nox -s import_profile
+            retval=$?
+        else
+            echo "Skipping import_profile as it is not supported by this package yet."
+            retval=0
+        fi
         ;;
     *)
         nox -s ${TEST_TYPE}

inject_import_profile.py

@@ -0,0 +1,89 @@
+import os
+import glob
+import re
+
+def patch_noxfile(path):
+    try:
+        with open(path, 'r', encoding='utf-8') as f:
+            content = f.read()
+    except UnicodeDecodeError:
+        return False
+        
+    if "import_profile" in content:
+        return False
+
+    # 1. Add "import_profile" to nox.options.sessions
+    # It looks like:
+    # nox.options.sessions = [
+    #     "unit",
+    #     ...
+    #     "docs",
+    # ]
+    # We want to add it right before the closing bracket.
+    if "nox.options.sessions =" in content:
+        content = re.sub(
+            r'(\n\s*"docs",\n)(\s*\])',
+            r'\1    "import_profile",\n\2',
+            content
+        )
+        # If "docs", is not the last one, maybe just add it before the closing bracket
+        if "import_profile" not in content:
+            content = re.sub(
+                r'(\n)(\s*\]\n)',
+                r'\1    "import_profile",\n\2',
+                content,
+                count=1,
+                flags=re.DOTALL
+            )
+            # Make sure it only replaced the first matching closing bracket for sessions
+            # wait, it's safer to just replace:
+            # content = content.replace('"docs",\n', '"docs",\n    "import_profile",\n')
+
+    # 2. Append the import_profile session at the end
+    # We need to extract the module namespace. It's usually "google".
+    # Let's extract python-gapic-namespace or python-gapic-name from somewhere,
+    # or just hardcode "google" like the previous agent, which worked for CI.
+    
+    session_text = """
+
+@nox.session(python="3.15")
+def import_profile(session):
+    \"\"\"Ensure import times remain below defined thresholds.\"\"\"
+    profiler_script = (
+        CURRENT_DIRECTORY.parent.parent / "scripts" / "import_profiler" / "profiler.py"
+    )
+    if not profiler_script.exists():
+        session.skip("The import profiler script was not found.")
+
+    session.install(".")
+    session.run(
+        "python",
+        str(profiler_script),
+        "--package",
+        "{pkg}",
+        "--iterations",
+        "10",
+        "--fail-threshold",
+        "5000",
+    )
+"""
+    
+    package_name = path.split('/')[1]
+    session_text = session_text.replace("{pkg}", package_name)
+    
+    content += session_text
+    
+    with open(path, 'w', encoding='utf-8') as f:
+        f.write(content)
+    return True
+
+if __name__ == "__main__":
+    count = 0
+    for path in glob.glob('packages/*/noxfile.py'):
+        # don't modify the goldens here, we use bazel for them
+        if "goldens" in path:
+            continue
+        if patch_noxfile(path):
+            count += 1
+
+    print(f"Patched {count} noxfiles.")

packages/gapic-generator/gapic/templates/noxfile.py.j2

@@ -86,6 +86,7 @@ nox.options.sessions = [
     "lint_setup_py",
     "blacken",
     "docs",
+    "import_profile",
 ]
 
 # Error if a python version is missing
@@ -641,4 +642,30 @@ def core_deps_from_source(session, protobuf_implementation):
             "PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION": protobuf_implementation,
         },
     )
+
+
+@nox.session(python=DEFAULT_PYTHON_VERSION)
+def import_profile(session):
+    """Ensure import times remain below defined thresholds."""
+    profiler_script = (
+        CURRENT_DIRECTORY.parent.parent / "scripts" / "import_profiler" / "profiler.py"
+    )
+    if not profiler_script.exists():
+        session.skip("The import profiler script was not found.")
+
+    session.install(".")
+    session.run(
+        "python",
+        str(profiler_script),
+        "--module",
+        {% if api.naming.module_namespace %}
+        "{{ api.naming.module_namespace[0] }}",
+        {% else %}
+        "{{ api.naming.versioned_module_name }}",
+        {% endif %}
+        "--iterations",
+        "10",
+        "--fail-threshold",
+        "5000",
+    )
 {% endblock %}

packages/google-cloud-asset/noxfile.py

@@ -88,6 +88,7 @@
     "lint_setup_py",
     "blacken",
     "docs",
+    "import_profile",
 ]
 
 # Error if a python version is missing
@@ -637,3 +638,25 @@ def core_deps_from_source(session, protobuf_implementation):
             "PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION": protobuf_implementation,
         },
     )
+
+
+@nox.session(python="3.15")
+def import_profile(session):
+    """Ensure import times remain below defined thresholds."""
+    profiler_script = (
+        CURRENT_DIRECTORY.parent.parent / "scripts" / "import_profiler" / "profiler.py"
+    )
+    if not profiler_script.exists():
+        session.skip("The import profiler script was not found.")
+
+    session.install(".")
+    session.run(
+        "python",
+        str(profiler_script),
+        "--package",
+        "google-cloud-asset",
+        "--iterations",
+        "10",
+        "--fail-threshold",
+        "5000",
+    )

packages/google-cloud-asset/test_trigger.txt

@@ -154,7 +154,7 @@ def _format_stats(title, data, p50, p90, p99, fmt):
 {_format_stats("Physical RSS RAM (MB)", rss_memories, p50_rss, p90_rss, p99_rss, ".4f")}"""
     print(final_output.strip())
 
-def run_master(iterations, target_module, cpu=0, csv_path=None, clear_cache=True):
+def run_master(iterations, target_module, cpu=0, csv_path=None, clear_cache=True, fail_threshold=None):
     """Orchestrates the benchmark."""
     if iterations < 1:
         raise ValueError("Number of iterations must be at least 1.")
@@ -229,6 +229,14 @@ def run_master(iterations, target_module, cpu=0, csv_path=None, clear_cache=True
         rss_memories, p50_rss, p90_rss, p99_rss
     )
 
+    if fail_threshold is not None:
+        if p99_time > fail_threshold:
+            print(f"\nFAILURE: P99 import time ({p99_time:.2f} ms) exceeds the failure threshold ({fail_threshold} ms).", file=sys.stderr)
+            sys.exit(1)
+        else:
+            print(f"\nSUCCESS: P99 import time ({p99_time:.2f} ms) is within the failure threshold ({fail_threshold} ms).")
+
+
 def run_trace(target_module):
     """Generates importtime trace log and writes it to a file."""
     trace_file = f"import_trace_{target_module.replace('.', '_')}.log"
@@ -307,8 +315,24 @@ def validate_module_name(module_name):
             raise argparse.ArgumentTypeError(f"'{module_name}' is not a valid Python module identifier.")
         return module_name
 
+    def find_module_from_package(pkg):
+        candidates = [
+            pkg.replace('-', '.'),
+            '.'.join(pkg.split('-')[:-1]) + '_' + pkg.split('-')[-1] if '-' in pkg else pkg,
+            pkg.replace('-', '_')
+        ]
+        for mod in candidates:
+            try:
+                if importlib.util.find_spec(mod):
+                    return mod
+            except Exception:
+                pass
+        return candidates[0]
+
     parser = argparse.ArgumentParser(description="Python SDK Import Profiler")
-    parser.add_argument("--module", type=validate_module_name, default="google.cloud.compute_v1", help="Target module to profile")
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("--module", type=validate_module_name, help="Target module to profile")
+    group.add_argument("--package", help="Target package name to profile (auto-detects module)")
     parser.add_argument("--iterations", type=int, default=50, help="Number of iterations")
     default_cpu = 0 if sys.platform.startswith("linux") else NO_CPU_PINNING
     parser.add_argument("--cpu", type=int, default=default_cpu, help="CPU core to pin to (or -1 for no pinning)")
@@ -317,20 +341,25 @@ def validate_module_name(module_name):
     parser.add_argument("--cprofile", action="store_true", help="Run cProfile")
     parser.add_argument("--mprofile", action="store_true", help="Run tracemalloc memory snapshot")
     parser.add_argument("--keep-pycache", action="store_true", help="Preserve __pycache__ and allow bytecode execution (Default: False, script automatically sweeps __pycache__ for true cold-starts)")
+    parser.add_argument("--fail-threshold", type=float, help="Fail the profiling if the P99 time exceeds this threshold (in ms).")
     parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS)
 
     args = parser.parse_args()
 
+    target_module = args.module
+    if args.package:
+        target_module = find_module_from_package(args.package)
+    
     if args.worker:
-        run_worker(args.module)
+        run_worker(target_module)
     elif args.trace:
         if not args.keep_pycache: clean_bytecode()
-        run_trace(args.module)
+        run_trace(target_module)
     elif args.cprofile:
         if not args.keep_pycache: clean_bytecode()
-        run_cprofile(args.module)
+        run_cprofile(target_module)
     elif args.mprofile:
         if not args.keep_pycache: clean_bytecode()
-        run_mprofile(args.module)
+        run_mprofile(target_module)
     else:
-        run_master(args.iterations, args.module, args.cpu, args.csv, not args.keep_pycache)
+        run_master(args.iterations, target_module, args.cpu, args.csv, not args.keep_pycache, args.fail_threshold)