diff --git a/pyperf/_process_time.py b/pyperf/_process_time.py index f18dee20..71e03b49 100644 --- a/pyperf/_process_time.py +++ b/pyperf/_process_time.py @@ -22,12 +22,13 @@ import tempfile import time +# Try importing resource module for memory usage tracking try: import resource except ImportError: resource = None - +# Function to get maximum RSS (Resident Set Size) memory usage def get_max_rss(*, children): if resource is not None: if children: @@ -35,13 +36,14 @@ def get_max_rss(*, children): else: resource_type = resource.RUSAGE_SELF usage = resource.getrusage(resource_type) + # macOS returns bytes, Linux returns kilobytes if sys.platform == 'darwin': return usage.ru_maxrss return usage.ru_maxrss * 1024 else: return 0 - +# Merge two cProfile profiling statistics files def merge_profile_stats_files(src, dst): """ Merging one existing pstats file into another. @@ -52,16 +54,17 @@ def merge_profile_stats_files(src, dst): dst_stats = pstats.Stats(dst) dst_stats.add(src_stats) dst_stats.dump_stats(dst) - os.unlink(src) + os.unlink(src) # Delete source after merging else: - os.rename(src, dst) - + os.rename(src, dst) # If no destination exists, just move source +# Benchmark a process by running it multiple times and collecting timing and memory data def bench_process(loops, args, kw, profile_filename=None): max_rss = 0 range_it = range(loops) start_time = time.perf_counter() + # If profiling is requested, create a temporary profile output file if profile_filename: temp_profile_filename = tempfile.mktemp() args = [args[0], "-m", "cProfile", "-o", temp_profile_filename] + args[1:] @@ -69,9 +72,10 @@ def bench_process(loops, args, kw, profile_filename=None): for _ in range_it: start_rss = get_max_rss(children=True) + # Start the external process proc = subprocess.Popen(args, **kw) with proc: - proc.wait() + proc.wait() # Wait for the process to complete exitcode = proc.returncode if exitcode != 0: @@ -81,9 +85,11 @@ def bench_process(loops, args, kw, profile_filename=None): os.unlink(temp_profile_filename) sys.exit(exitcode) + # Update maximum observed memory usage rss = get_max_rss(children=True) - start_rss max_rss = max(max_rss, rss) + # Merge profiling results if profiling is enabled if profile_filename: merge_profile_stats_files( temp_profile_filename, profile_filename @@ -92,7 +98,7 @@ def bench_process(loops, args, kw, profile_filename=None): dt = time.perf_counter() - start_time return (dt, max_rss) - +# Parse optional hook plugins from the command line arguments def load_hooks(metadata): hook_names = [] while "--hook" in sys.argv: @@ -103,7 +109,7 @@ def load_hooks(metadata): del sys.argv[hook_idx] if len(hook_names): - # Only import pyperf if we know we have hooks + # Import hooks module only if hooks are requested import pyperf._hooks hook_managers = pyperf._hooks.instantiate_selected_hooks(hook_names) @@ -113,31 +119,32 @@ def load_hooks(metadata): return hook_managers - +# Write benchmark results to stdout in a structured format def write_data(dt, max_rss, metadata, out=sys.stdout): - # Write the data that is communicated back to the main orchestration process. - # It is three lines containing: - # - The runtime (in seconds) - # - max_rss (or -1, if not able to compute) - # - The metadata to add to the benchmark entry, as a JSON dictionary + # Three lines output: + # 1. Runtime (seconds) + # 2. Maximum RSS memory (or -1 if unavailable) + # 3. Metadata as a JSON dictionary print(dt, file=out) print(max_rss or -1, file=out) json.dump(metadata, fp=out) print(file=out) - +# Main function for script execution def main(): - # Make sure that the pyperf module wasn't imported + # Prevent users from wrongly running this internal module via -m option if 'pyperf' in sys.modules: print("ERROR: don't run %s -m pyperf._process, run the .py script" % os.path.basename(sys.executable)) sys.exit(1) + # Validate minimal argument count if len(sys.argv) < 3: print("Usage: %s %s loops program [arg1 arg2 ...] [--profile profile]" % (os.path.basename(sys.executable), __file__)) sys.exit(1) + # Check if profiling is requested via command line if "--profile" in sys.argv: profile_idx = sys.argv.index("--profile") profile_filename = sys.argv[profile_idx + 1] @@ -149,10 +156,12 @@ def main(): metadata = {} hook_managers = load_hooks(metadata) + # Extract the number of loops and the target command arguments loops = int(sys.argv[1]) args = sys.argv[2:] kw = {} + # Redirect stdin, stdout, stderr appropriately if hasattr(subprocess, 'DEVNULL'): devnull = None kw['stdin'] = subprocess.DEVNULL @@ -163,6 +172,7 @@ def main(): kw['stdout'] = devnull kw['stderr'] = subprocess.STDOUT + # Handle multiple hooks cleanly using ExitStack with contextlib.ExitStack() as stack: for hook in hook_managers.values(): stack.enter_context(hook) @@ -171,11 +181,12 @@ def main(): if devnull is not None: devnull.close() + # Call hook teardown methods for hook in hook_managers.values(): hook.teardown(metadata) write_data(dt, max_rss, metadata) - +# Entry point check if __name__ == "__main__": main()