@@ -154,7 +154,7 @@ def _format_stats(title, data, p50, p90, p99, fmt):
154154{ _format_stats ("Physical RSS RAM (MB)" , rss_memories , p50_rss , p90_rss , p99_rss , ".4f" )} """
155155 print (final_output .strip ())
156156
157- def run_master (iterations , target_module , cpu = 0 , csv_path = None , clear_cache = True ):
157+ def run_master (iterations , target_module , cpu = 0 , csv_path = None , clear_cache = True , fail_threshold = None , diff_baseline = None , diff_threshold = None ):
158158 """Orchestrates the benchmark."""
159159 if iterations < 1 :
160160 raise ValueError ("Number of iterations must be at least 1." )
@@ -191,6 +191,7 @@ def run_master(iterations, target_module, cpu=0, csv_path=None, clear_cache=True
191191 times .append (data ["time_ms" ])
192192 memories .append (data ["peak_ram_mb" ])
193193 rss_memories .append (data ["rss_ram_mb" ])
194+ print (f"Iteration { i + 1 } /{ iterations } completed in { data ['time_ms' ]:.2f} ms" )
194195 if i > 0 and loaded_modules_val != data ["loaded_modules" ]:
195196 print (f"WARNING: Non-deterministic import behavior! Iteration { i + 1 } loaded { data ['loaded_modules' ]} modules (expected { loaded_modules_val } )." , file = sys .stderr )
196197 if i > 0 and loaded_lines_val != data ["loaded_lines" ]:
@@ -207,6 +208,13 @@ def run_master(iterations, target_module, cpu=0, csv_path=None, clear_cache=True
207208 print (f"Error in worker process:\n { e .stderr } " , file = sys .stderr )
208209 raise e
209210
211+ if iterations > 1 :
212+ times = times [1 :]
213+ memories = memories [1 :]
214+ rss_memories = rss_memories [1 :]
215+ iterations -= 1
216+ print ("Discarded the first iteration as a cache burn-in run." )
217+
210218 # Write CSV if requested
211219 if csv_path :
212220 with open (csv_path , "w" , newline = "" , encoding = "utf-8" ) as f :
@@ -229,6 +237,68 @@ def run_master(iterations, target_module, cpu=0, csv_path=None, clear_cache=True
229237 rss_memories , p50_rss , p90_rss , p99_rss
230238 )
231239
240+ exit_code = 0
241+ final_messages = []
242+
243+ baseline_p50 = None
244+ if diff_baseline :
245+ if os .path .exists (diff_baseline ):
246+ baseline_times = []
247+ with open (diff_baseline , "r" , encoding = "utf-8" ) as f :
248+ reader = csv .reader (f )
249+ next (reader ) # skip header
250+ for row in reader :
251+ baseline_times .append (float (row [1 ]))
252+ if baseline_times :
253+ baseline_p50 , _ , _ = _calculate_percentiles (baseline_times )
254+
255+ if baseline_p50 is not None :
256+ diff = p50_time - baseline_p50
257+
258+ diff_msg = (
259+ f"--- Diff vs Baseline ---\n "
260+ f"Baseline Median: { baseline_p50 :.2f} ms\n "
261+ f"Current Median: { p50_time :.2f} ms\n "
262+ f"Difference: { diff :+.2f} ms"
263+ )
264+ final_messages .append (diff_msg )
265+
266+ relative_diff_threshold = 0.15 * baseline_p50
267+ if diff > diff_threshold and diff > relative_diff_threshold :
268+ final_messages .append (
269+ f"FAILURE: Import time regression of { diff :.2f} ms exceeds both the absolute threshold ({ diff_threshold } ms) "
270+ f"and the relative threshold ({ relative_diff_threshold :.2f} ms, 15% of baseline Median)."
271+ )
272+ exit_code = 1
273+ else :
274+ if diff > diff_threshold :
275+ final_messages .append (f"SUCCESS: Import time regression of { diff :.2f} ms exceeds absolute threshold ({ diff_threshold } ms) but is within relative threshold ({ relative_diff_threshold :.2f} ms, 15%)." )
276+ else :
277+ final_messages .append ("SUCCESS: Import time diff is within acceptable thresholds." )
278+ else :
279+ final_messages .append (f"WARNING: Baseline CSV { diff_baseline } not found. Skipping diff check." )
280+
281+ if fail_threshold is not None :
282+ if p50_time > fail_threshold :
283+ if baseline_p50 is not None and baseline_p50 > fail_threshold :
284+ final_messages .append (f"WARNING: Median import time ({ p50_time :.2f} ms) exceeds the absolute failure threshold ({ fail_threshold } ms), but the baseline ({ baseline_p50 :.2f} ms) also exceeded it. Bypassing absolute backstop failure." )
285+ else :
286+ final_messages .append (f"FAILURE: Median import time ({ p50_time :.2f} ms) exceeds the failure threshold ({ fail_threshold } ms)." )
287+ exit_code = 1
288+ else :
289+ final_messages .append (f"SUCCESS: Median import time ({ p50_time :.2f} ms) is within the failure threshold ({ fail_threshold } ms)." )
290+
291+ if final_messages :
292+ print ("\n " + "\n " .join (final_messages ))
293+
294+ if exit_code == 0 :
295+ print ("\n Session import_profiler was successful." )
296+ sys .exit (0 )
297+ else :
298+ print ("\n Session import_profiler failed." )
299+ sys .exit (1 )
300+
301+
232302def run_trace (target_module ):
233303 """Generates importtime trace log and writes it to a file."""
234304 trace_file = f"import_trace_{ target_module .replace ('.' , '_' )} .log"
@@ -307,8 +377,59 @@ def validate_module_name(module_name):
307377 raise argparse .ArgumentTypeError (f"'{ module_name } ' is not a valid Python module identifier." )
308378 return module_name
309379
380+ def find_module_from_package (pkg ):
381+ import importlib .metadata
382+ import importlib .util
383+
384+ # 1. Try to use importlib.metadata.files (works for standard installations from PyPI/wheels)
385+ try :
386+ files = importlib .metadata .files (pkg )
387+ if files :
388+ init_files = [str (f ) for f in files if str (f ).endswith ('__init__.py' ) and '__pycache__' not in str (f ) and not str (f ).startswith ('tests/' )]
389+ if init_files :
390+ from pathlib import Path
391+ shortest_init = min (init_files , key = lambda p : len (Path (p ).parts ))
392+ parts = Path (shortest_init ).parent .parts
393+ mod = '.' .join (parts )
394+ if importlib .util .find_spec (mod ):
395+ return mod
396+ except Exception :
397+ pass
398+
399+ # 2. Try setuptools.find_namespace_packages() in current directory (works for editable installs in source trees)
400+ try :
401+ import setuptools
402+ import os
403+ if os .path .exists ('setup.py' ) or os .path .exists ('pyproject.toml' ):
404+ pkgs = setuptools .find_namespace_packages (where = '.' )
405+ for p in sorted (pkgs , key = len ):
406+ if p in ("google" , "google.cloud" ) or p .startswith ("tests" ):
407+ continue
408+ path = p .replace ('.' , os .sep )
409+ if os .path .isfile (os .path .join (path , '__init__.py' )):
410+ if importlib .util .find_spec (p ):
411+ return p
412+ except Exception :
413+ pass
414+
415+ # 3. Fallback to basic string manipulation heuristics
416+ candidates = [
417+ pkg .replace ('-' , '.' ),
418+ '.' .join (pkg .split ('-' )[:- 1 ]) + '_' + pkg .split ('-' )[- 1 ] if '-' in pkg else pkg ,
419+ pkg .replace ('-' , '_' )
420+ ]
421+ for mod in candidates :
422+ try :
423+ if importlib .util .find_spec (mod ):
424+ return mod
425+ except Exception :
426+ pass
427+ return candidates [0 ]
428+
310429 parser = argparse .ArgumentParser (description = "Python SDK Import Profiler" )
311- parser .add_argument ("--module" , type = validate_module_name , default = "google.cloud.compute_v1" , help = "Target module to profile" )
430+ group = parser .add_mutually_exclusive_group (required = True )
431+ group .add_argument ("--module" , type = validate_module_name , help = "Target module to profile" )
432+ group .add_argument ("--package" , help = "Target package name to profile (auto-detects module)" )
312433 parser .add_argument ("--iterations" , type = int , default = 50 , help = "Number of iterations" )
313434 default_cpu = 0 if sys .platform .startswith ("linux" ) else NO_CPU_PINNING
314435 parser .add_argument ("--cpu" , type = int , default = default_cpu , help = "CPU core to pin to (or -1 for no pinning)" )
@@ -317,20 +438,27 @@ def validate_module_name(module_name):
317438 parser .add_argument ("--cprofile" , action = "store_true" , help = "Run cProfile" )
318439 parser .add_argument ("--mprofile" , action = "store_true" , help = "Run tracemalloc memory snapshot" )
319440 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)" )
441+ parser .add_argument ("--fail-threshold" , type = float , help = "Fail the profiling if the Median time exceeds this threshold (in ms)." )
442+ parser .add_argument ("--diff-baseline" , help = "Path to a baseline CSV file to compare against." )
443+ parser .add_argument ("--diff-threshold" , type = float , default = 100.0 , help = "Fail if Median time exceeds baseline Median by this many ms." )
320444 parser .add_argument ("--worker" , action = "store_true" , help = argparse .SUPPRESS )
321445
322446 args = parser .parse_args ()
323447
448+ target_module = args .module
449+ if args .package :
450+ target_module = find_module_from_package (args .package )
451+
324452 if args .worker :
325- run_worker (args . module )
453+ run_worker (target_module )
326454 elif args .trace :
327455 if not args .keep_pycache : clean_bytecode ()
328- run_trace (args . module )
456+ run_trace (target_module )
329457 elif args .cprofile :
330458 if not args .keep_pycache : clean_bytecode ()
331- run_cprofile (args . module )
459+ run_cprofile (target_module )
332460 elif args .mprofile :
333461 if not args .keep_pycache : clean_bytecode ()
334- run_mprofile (args . module )
462+ run_mprofile (target_module )
335463 else :
336- run_master (args .iterations , args . module , args .cpu , args .csv , not args .keep_pycache )
464+ run_master (args .iterations , target_module , args .cpu , args .csv , not args .keep_pycache , args . fail_threshold , args . diff_baseline , args . diff_threshold )
0 commit comments