feat: implement dynamic import profile diff check against baseline

Author: hebaalazzeh

ci/run_single_test.sh

@@ -107,7 +107,27 @@ case ${TEST_TYPE} in
         ;;
     import_profile)
         if nox --list-sessions | grep -q "import_profile"; then
-            nox -s import_profile
+            rm -f /tmp/baseline.csv
+            if [ -n "${TARGET_BRANCH}" ]; then
+                if git rev-parse HEAD^1 >/dev/null 2>&1; then
+                    echo "Checking out HEAD^1 for baseline..."
+                    git checkout HEAD^1
+                    if nox --list-sessions | grep -q "import_profile"; then
+                        nox -s import_profile -- --csv /tmp/baseline.csv
+                    else
+                        echo "No import_profile session on baseline."
+                    fi
+                    git checkout -
+                else
+                    echo "Could not find HEAD^1. Skipping baseline generation."
+                fi
+            fi
+            
+            if [ -f /tmp/baseline.csv ]; then
+                nox -s import_profile -- --diff-baseline /tmp/baseline.csv --diff-threshold 100
+            else
+                nox -s import_profile
+            fi
             retval=$?
         else
             echo "Skipping import_profile as it is not supported by this package yet."

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

@@ -667,5 +667,6 @@ def import_profile(session):
         "10",
         "--fail-threshold",
         "5000",
+        *session.posargs,
     )
 {% endblock %}

patch_goldens.py

@@ -0,0 +1,87 @@
+import os
+import glob
+import re
+
+def get_module_name(golden_dir):
+    setup_py = os.path.join(golden_dir, "setup.py")
+    if not os.path.exists(setup_py):
+        return "google"
+    
+    with open(setup_py, "r", encoding="utf-8") as f:
+        content = f.read()
+    
+    match = re.search(r"'(google/[^']+)/gapic_version.py'", content)
+    if match:
+        return match.group(1).replace("/", ".")
+    
+    match = re.search(r"\"(google/[^\"]+)/gapic_version.py\"", content)
+    if match:
+        return match.group(1).replace("/", ".")
+        
+    return "google"
+
+def patch_noxfile(golden_dir):
+    path = os.path.join(golden_dir, "noxfile.py")
+    if not os.path.exists(path):
+        return False
+        
+    try:
+        with open(path, 'r', encoding='utf-8') as f:
+            content = f.read()
+    except UnicodeDecodeError:
+        return False
+        
+    module_name = get_module_name(golden_dir)
+
+    idx = content.find("@nox.session(python=\"3.15\")\ndef import_profile(session):")
+    if idx != -1:
+        content = content[:idx].strip() + "\n"
+    else:
+        idx = content.find("@nox.session\ndef import_profile")
+        if idx != -1:
+            content = content[:idx].strip() + "\n"
+
+    if "nox.options.sessions =" in content and "\"import_profile\"" not in content:
+        content = re.sub(
+            r'(\n\s*"docs",\n)(\s*\])',
+            r'\1    "import_profile",\n\2',
+            content
+        )
+    
+    session_text = f"""
+@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",
+        "{module_name}",
+        "--fail-threshold",
+        "5000",
+        "--iterations",
+        "10",
+        *session.posargs,
+    )
+"""
+    
+    content += "\n" + session_text
+    
+    with open(path, 'w', encoding='utf-8') as f:
+        f.write(content)
+    return True
+
+count = 0
+for golden_dir in glob.glob('packages/gapic-generator/tests/integration/goldens/*'):
+    if os.path.isdir(golden_dir):
+        if patch_noxfile(golden_dir):
+            count += 1
+
+print(f"Patched {count} goldens.")

scripts/import_profiler/profiler.py

@@ -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, fail_threshold=None):
+def run_master(iterations, target_module, cpu=0, csv_path=None, clear_cache=True, fail_threshold=None, diff_baseline=None, diff_threshold=None):
     """Orchestrates the benchmark."""
     if iterations < 1:
         raise ValueError("Number of iterations must be at least 1.")
@@ -229,12 +229,39 @@ 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).")
+if fail_threshold is not None:
+            if p99_time > fail_threshold:
+                print(f"
+FAILURE: P99 import time ({p99_time:.2f} ms) exceeds the failure threshold ({fail_threshold} ms).", file=sys.stderr)
+                sys.exit(1)
+            else:
+                print(f"
+SUCCESS: P99 import time ({p99_time:.2f} ms) is within the failure threshold ({fail_threshold} ms).")
+    
+        if diff_baseline:
+            if os.path.exists(diff_baseline):
+                baseline_times = []
+                with open(diff_baseline, "r", encoding="utf-8") as f:
+                    reader = csv.reader(f)
+                    next(reader) # skip header
+                    for row in reader:
+                        baseline_times.append(float(row[1]))
+                _, _, baseline_p99 = _calculate_percentiles(baseline_times)
+                diff = p99_time - baseline_p99
+                
+                print(f"
+--- Diff vs Baseline ---")
+                print(f"Baseline P99: {baseline_p99:.2f} ms")
+                print(f"Current P99:  {p99_time:.2f} ms")
+                print(f"Difference:   {diff:+.2f} ms")
+                
+                if diff > diff_threshold:
+                    print(f"FAILURE: Import time regression of {diff:.2f} ms exceeds the allowed threshold of {diff_threshold} ms.", file=sys.stderr)
+                    sys.exit(1)
+                else:
+                    print("SUCCESS: Import time diff is within acceptable thresholds.")
+            else:
+                print(f"WARNING: Baseline CSV {diff_baseline} not found. Skipping diff check.")
 
 
 def run_trace(target_module):
@@ -342,6 +369,8 @@ def find_module_from_package(pkg):
     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("--diff-baseline", help="Path to a baseline CSV file to compare against.")
+    parser.add_argument("--diff-threshold", type=float, default=100.0, help="Fail if P99 time exceeds baseline P99 by this many ms.")
     parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS)
 
     args = parser.parse_args()
@@ -362,4 +391,4 @@ def find_module_from_package(pkg):
         if not args.keep_pycache: clean_bytecode()
         run_mprofile(target_module)
     else:
-        run_master(args.iterations, target_module, args.cpu, args.csv, not args.keep_pycache, args.fail_threshold)
+        run_master(args.iterations, target_module, args.cpu, args.csv, not args.keep_pycache, args.fail_threshold, args.diff_baseline, args.diff_threshold)