Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 39 additions & 0 deletions .github/workflows/import-profiler.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
name: import-profiler

on:
pull_request:
branches:
- main
- preview
# Trigger workflow on GitHub merge queue events
merge_group:
types: [checks_requested]

permissions:
contents: read

jobs:
import-profile:
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v6
with:
fetch-depth: 2
- name: Setup Python
uses: actions/setup-python@v6
with:
python-version: "3.15"
allow-prereleases: true
- name: Install nox
run: |
python -m pip install --upgrade setuptools pip wheel
python -m pip install nox
- name: Run import profiler
env:
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"
run: |
ci/run_conditional_tests.sh
42 changes: 42 additions & 0 deletions ci/run_single_test.sh
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,48 @@ case ${TEST_TYPE} in
;;
esac
;;
import_profile)
if [ -f setup.py ]; then
echo "Creating temporary virtualenv for import profile..."
python3 -m venv .venv-profiler
source .venv-profiler/bin/activate
pip install -e .

PACKAGE_NAME=$(basename $(pwd))
PROFILER_SCRIPT="../../scripts/import_profiler/profiler.py"

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 [ -f setup.py ]; then
pip install -e .
python ${PROFILER_SCRIPT} --package ${PACKAGE_NAME} --iterations 10 --csv /tmp/baseline.csv
else
echo "setup.py not found on baseline. Skipping baseline generation."
fi
git checkout -
# Re-install the current branch to ensure we profile the latest code
pip install -e .
else
echo "Could not find HEAD^1. Skipping baseline generation."
fi
fi

if [ -f /tmp/baseline.csv ]; then
python ${PROFILER_SCRIPT} --package ${PACKAGE_NAME} --iterations 10 --fail-threshold 5000 --diff-baseline /tmp/baseline.csv --diff-threshold 100
else
python ${PROFILER_SCRIPT} --package ${PACKAGE_NAME} --iterations 10 --fail-threshold 5000
fi
retval=$?
deactivate
rm -rf .venv-profiler
else
echo "Skipping import_profile as this does not appear to be a Python package (no setup.py)."
retval=0
fi
;;
*)
nox -s ${TEST_TYPE}
retval=$?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This PR doesn't touch any packages, so it's hard to tell what the action will look like. Can you temporarily touch some files to test it out?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should also see how long the new check would take if all packages were updated. In #17438, I added a new unit_test:all_packages tag, that will run against all packages when added. Maybe we should support that here too

It's possible we would only want to profile a few key packages instead of all of them. Most generated libraries should be very similar

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

addressed!

Expand Down
68 changes: 61 additions & 7 deletions scripts/import_profiler/profiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, diff_baseline=None, diff_threshold=None):
"""Orchestrates the benchmark."""
if iterations < 1:
raise ValueError("Number of iterations must be at least 1.")
Expand Down Expand Up @@ -229,6 +229,37 @@ 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 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("\n--- 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):
"""Generates importtime trace log and writes it to a file."""
trace_file = f"import_trace_{target_module.replace('.', '_')}.log"
Expand Down Expand Up @@ -307,8 +338,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)")
Expand All @@ -317,20 +364,27 @@ 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("--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()

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, args.diff_baseline, args.diff_threshold)
Loading