diff --git a/scripts/import_profiler/profiler.py b/scripts/import_profiler/profiler.py index d2734dadce88..3f19bf9425a8 100644 --- a/scripts/import_profiler/profiler.py +++ b/scripts/import_profiler/profiler.py @@ -40,7 +40,7 @@ def get_rss_mb(): except Exception: return 0.0 -def run_worker(target_module): +def run_worker(target_module, skip_line_count=False): """Performs ONE import and returns metrics.""" tracemalloc.start() importlib.invalidate_caches() @@ -66,33 +66,36 @@ def run_worker(target_module): new_modules = modules_after - modules_before loaded_lines = 0 - for m in new_modules: - try: - file_path = sys.modules[m].__file__ - if not file_path: - continue + if not skip_line_count: + for m in new_modules: + try: + file_path = sys.modules[m].__file__ + if not file_path: + continue - if file_path.endswith('.pyc'): - try: - file_path = importlib.util.source_from_cache(file_path) - except ValueError: - # Raised if the .pyc path does not follow standard PEP 3147/488 conventions. - # We pass silently because the unresolved file_path will still end in '.pyc', - # meaning the subsequent '.endswith('.py')' check will fail and safely skip - # trying to count lines in a binary file. - pass - if file_path.endswith('.py'): - try: - with open(file_path, 'r', encoding='utf-8', errors='ignore') as f: - loaded_lines += sum(1 for _ in f) - except OSError as e: - print(f"WARNING: Failed to read lines from {file_path}: {e}", file=sys.stderr) - except KeyError: - # Module disappeared from sys.modules during execution - pass - except AttributeError: - # Module has no __file__ attribute (likely a C-extension or built-in) - pass + if file_path.endswith('.pyc'): + try: + file_path = importlib.util.source_from_cache(file_path) + except ValueError: + # Raised if the .pyc path does not follow standard PEP 3147/488 conventions. + # We pass silently because the unresolved file_path will still end in '.pyc', + # meaning the subsequent '.endswith('.py')' check will fail and safely skip + # trying to count lines in a binary file. + pass + if file_path.endswith('.py'): + try: + with open(file_path, 'r', encoding='utf-8', errors='ignore') as f: + loaded_lines += sum(1 for _ in f) + except OSError as e: + print(f"WARNING: Failed to read lines from {file_path}: {e}", file=sys.stderr) + except KeyError: + # Module disappeared from sys.modules during execution + pass + except AttributeError: + # Module has no __file__ attribute (likely a C-extension or built-in) + pass + else: + loaded_lines = -1 # Output to stdout for the Master to capture metrics = { @@ -106,6 +109,8 @@ def run_worker(target_module): def _run_worker_and_parse(cmd): result = subprocess.run(cmd, capture_output=True, text=True, check=True) + if result.stderr.strip(): + print(result.stderr.strip(), file=sys.stderr) try: lines = result.stdout.strip().splitlines() data = None @@ -185,6 +190,8 @@ def run_master(iterations, target_module, cpu=0, csv_path=None, clear_cache=True cmd += ["taskset", "-c", str(cpu)] cmd += python_exe + [__file__, "--worker", f"--module={target_module}"] + if i > 0: + cmd += ["--skip-line-count"] try: data = _run_worker_and_parse(cmd) @@ -194,11 +201,12 @@ def run_master(iterations, target_module, cpu=0, csv_path=None, clear_cache=True print(f"Iteration {i+1}/{iterations} completed in {data['time_ms']:.2f} ms") if i > 0 and loaded_modules_val != data["loaded_modules"]: print(f"WARNING: Non-deterministic import behavior! Iteration {i+1} loaded {data['loaded_modules']} modules (expected {loaded_modules_val}).", file=sys.stderr) - if i > 0 and loaded_lines_val != data["loaded_lines"]: - print(f"WARNING: Non-deterministic import behavior! Iteration {i+1} loaded {data['loaded_lines']} lines (expected {loaded_lines_val}).", file=sys.stderr) loaded_modules_val = data["loaded_modules"] - loaded_lines_val = data["loaded_lines"] + if data["loaded_lines"] == -1: + data["loaded_lines"] = loaded_lines_val + else: + loaded_lines_val = data["loaded_lines"] except FileNotFoundError as e: if cpu != NO_CPU_PINNING and cmd and cmd[0] == "taskset": print("ERROR: 'taskset' command not found. CPU pinning is enabled but taskset is not installed. " @@ -293,10 +301,9 @@ def run_master(iterations, target_module, cpu=0, csv_path=None, clear_cache=True if exit_code == 0: print("\nSession import_profiler was successful.") - sys.exit(0) else: print("\nSession import_profiler failed.") - sys.exit(1) + return exit_code def run_trace(target_module): @@ -368,63 +375,64 @@ def run_mprofile(target_module): if p.exitcode != 0: print(f"Error generating memory snapshot, process exited with code {p.exitcode}", file=sys.stderr) -if __name__ == "__main__": +def validate_module_name(module_name): + """Validates that the input is a structurally valid Python module identifier to prevent arbitrary code execution.""" import argparse + if not all(part.isidentifier() for part in module_name.split('.')): + raise argparse.ArgumentTypeError(f"'{module_name}' is not a valid Python module identifier.") + return module_name - def validate_module_name(module_name): - """Validates that the input is a structurally valid Python module identifier to prevent arbitrary code execution.""" - if not all(part.isidentifier() for part in module_name.split('.')): - raise argparse.ArgumentTypeError(f"'{module_name}' is not a valid Python module identifier.") - return module_name +def find_module_from_package(pkg): + import importlib.metadata + import importlib.util - def find_module_from_package(pkg): - import importlib.metadata - import importlib.util + # 1. Try to use importlib.metadata.files (works for standard installations from PyPI/wheels) + try: + files = importlib.metadata.files(pkg) + if files: + 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/')] + if init_files: + from pathlib import Path + shortest_init = min(init_files, key=lambda p: len(Path(p).parts)) + parts = Path(shortest_init).parent.parts + mod = '.'.join(parts) + if importlib.util.find_spec(mod): + return mod + except Exception: + pass - # 1. Try to use importlib.metadata.files (works for standard installations from PyPI/wheels) - try: - files = importlib.metadata.files(pkg) - if files: - 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/')] - if init_files: - from pathlib import Path - shortest_init = min(init_files, key=lambda p: len(Path(p).parts)) - parts = Path(shortest_init).parent.parts - mod = '.'.join(parts) - if importlib.util.find_spec(mod): - return mod - except Exception: - pass + # 2. Try setuptools.find_namespace_packages() in current directory (works for editable installs in source trees) + try: + import setuptools + import os + if os.path.exists('setup.py') or os.path.exists('pyproject.toml'): + pkgs = setuptools.find_namespace_packages(where='.') + for p in sorted(pkgs, key=len): + if p in ("google", "google.cloud") or p.startswith("tests"): + continue + path = p.replace('.', os.sep) + if os.path.isfile(os.path.join(path, '__init__.py')): + if importlib.util.find_spec(p): + return p + except Exception: + pass - # 2. Try setuptools.find_namespace_packages() in current directory (works for editable installs in source trees) + # 3. Fallback to basic string manipulation heuristics + candidates = [ + pkg.replace('-', '.'), + '.'.join(pkg.split('-')[:-1]) + '_' + pkg.split('-')[-1] if '-' in pkg else pkg, + pkg.replace('-', '_') + ] + for mod in candidates: try: - import setuptools - import os - if os.path.exists('setup.py') or os.path.exists('pyproject.toml'): - pkgs = setuptools.find_namespace_packages(where='.') - for p in sorted(pkgs, key=len): - if p in ("google", "google.cloud") or p.startswith("tests"): - continue - path = p.replace('.', os.sep) - if os.path.isfile(os.path.join(path, '__init__.py')): - if importlib.util.find_spec(p): - return p + if importlib.util.find_spec(mod): + return mod except Exception: pass + return candidates[0] - # 3. Fallback to basic string manipulation heuristics - 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] +if __name__ == "__main__": + import argparse parser = argparse.ArgumentParser(description="Python SDK Import Profiler") group = parser.add_mutually_exclusive_group(required=True) @@ -442,6 +450,7 @@ def find_module_from_package(pkg): 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 Median time exceeds baseline Median by this many ms.") parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS) + parser.add_argument("--skip-line-count", action="store_true", help=argparse.SUPPRESS) args = parser.parse_args() @@ -450,7 +459,7 @@ def find_module_from_package(pkg): target_module = find_module_from_package(args.package) if args.worker: - run_worker(target_module) + run_worker(target_module, skip_line_count=args.skip_line_count) elif args.trace: if not args.keep_pycache: clean_bytecode() run_trace(target_module) @@ -461,4 +470,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, args.diff_baseline, args.diff_threshold) \ No newline at end of file + sys.exit(run_master(args.iterations, target_module, args.cpu, args.csv, not args.keep_pycache, args.fail_threshold, args.diff_baseline, args.diff_threshold)) \ No newline at end of file diff --git a/scripts/import_profiler/test_profiler.py b/scripts/import_profiler/test_profiler.py new file mode 100644 index 000000000000..db75be16786f --- /dev/null +++ b/scripts/import_profiler/test_profiler.py @@ -0,0 +1,366 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import csv +import json +import os +import sys +from unittest.mock import MagicMock, mock_open, patch + +import pytest + +# Ensure scripts/import_profiler is in sys.path +sys.path.insert(0, os.path.dirname(__file__)) + +import profiler +from profiler import ( + NO_CPU_PINNING, + _run_worker_and_parse, + find_module_from_package, + get_rss_mb, + run_master, + run_worker, +) + +# ===================================================================== +# 1. UTILITY FUNCTIONS TESTS +# ===================================================================== + + +def test_get_rss_mb(): + """Verifies get_rss_mb returns a float representing megabytes if available.""" + rss = get_rss_mb() + assert isinstance(rss, float) + assert rss >= 0.0 + + +@patch("os.walk") +@patch("os.remove") +@patch("shutil.rmtree") +def test_clean_bytecode(mock_rmtree, mock_remove, mock_walk): + """Verifies clean_bytecode successfully cleans bytecode files & caches.""" + clean_bytecode_helper = getattr(profiler, "clean_bytecode", None) + if clean_bytecode_helper is None: + pytest.skip("clean_bytecode is not exported at the module level.") + + mock_walk.return_value = [ + ("/test_dir", ["__pycache__"], ["test.pyc", "test.py"]) + ] + clean_bytecode_helper() + assert mock_remove.called or mock_rmtree.called + + +def test_find_module_from_package_resolves(): + """Verifies resolving of packages to modules if helper is exported.""" + with patch( + "importlib.metadata.packages_distributions", + return_value={"google-cloud-storage": ["google.cloud.storage"]}, + ): + res = find_module_from_package("google-cloud-storage") + assert res == "google.cloud.storage" + + +def test_find_module_from_package_fallback(): + """Verifies fallback transforms work correctly if helper is exported.""" + with patch("importlib.util.find_spec", side_effect=lambda mod: mod == "my_dummy_mod"): + res = find_module_from_package("my-dummy-mod") + assert res == "my_dummy_mod" + + +@patch("subprocess.run") +def test_run_trace(mock_run): + """Verifies that run_trace executes python with -X importtime and writes logs.""" + run_trace_helper = getattr(profiler, "run_trace", None) + if run_trace_helper is None: + pytest.skip("run_trace is not exported at the module level.") + + # Mock subprocess output + mock_run.return_value = MagicMock( + stdout="", stderr="importtime: dummy trace output", returncode=0 + ) + + with patch("builtins.open", mock_open()) as mock_file, patch( + "builtins.print" + ): + run_trace_helper("math") + + # Verify that subprocess was invoked with -X importtime + called_cmd = mock_run.call_args[0][0] + assert "-X" in called_cmd + assert "importtime" in called_cmd + + # Verify that it attempted to write the trace log file + assert mock_file.called + + +# ===================================================================== +# 2. WORKER TESTS (run_worker) +# ===================================================================== + + +def test_run_worker_with_skip_line_count(capsys): + """Verifies worker returns -1 for loaded_lines when flag is active.""" + dummy_module = "sys" + run_worker(dummy_module, skip_line_count=True) + captured = capsys.readouterr() + + assert "__METRICS__:" in captured.out + metrics_line = [ + l for l in captured.out.splitlines() if l.startswith("__METRICS__:") + ][0] + metrics = json.loads(metrics_line.split("__METRICS__:", 1)[1]) + + assert metrics["loaded_lines"] == -1 + + +def test_run_worker_counts_lines_correctly(capsys): + """Verifies worker correctly resolves paths and counts raw lines.""" + dummy_module = "math" + mock_module = MagicMock() + mock_module.__file__ = "/mock/path/dummy_module.py" + dummy_content = "import os\nprint('hello')\nx = 1\n" + + class CustomModulesDict(dict): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._call_count = 0 + + def keys(self): + self._call_count += 1 + if self._call_count == 1: + return {"math"} + return {"math", "dummy_test_mod"} + + fake_modules = CustomModulesDict( + {"math": sys.modules["math"], "dummy_test_mod": mock_module} + ) + + with patch("sys.modules", fake_modules), patch( + "builtins.open", mock_open(read_data=dummy_content) + ), patch("profiler.importlib.invalidate_caches"): + + run_worker(dummy_module, skip_line_count=False) + captured = capsys.readouterr() + + metrics_line = [ + l for l in captured.out.splitlines() if l.startswith("__METRICS__:") + ][0] + metrics = json.loads(metrics_line.split("__METRICS__:", 1)[1]) + + assert metrics["loaded_lines"] == 3 + + +# ===================================================================== +# 3. PARSER TESTS (_run_worker_and_parse) +# ===================================================================== + + +def test_run_worker_and_parse_success(): + """Verifies master extracts metrics JSON block from worker stdout.""" + mock_stdout = ( + '__METRICS__:{"time_ms": 15.5, "peak_ram_mb": 12.0, "rss_ram_mb":' + ' 10.0, "loaded_modules": 22, "loaded_lines": 120}\n' + ) + mock_process = MagicMock(stdout=mock_stdout, stderr="") + + with patch("subprocess.run", return_value=mock_process): + data = _run_worker_and_parse(["python", "profiler.py"]) + assert data["time_ms"] == 15.5 + assert data["loaded_lines"] == 120 + + +def test_run_worker_and_parse_forwards_stderr(capsys): + """Verifies worker stderr warnings are forwarded to master's stderr.""" + mock_stdout = ( + '__METRICS__:{"time_ms": 10.0, "peak_ram_mb": 12.0, "rss_ram_mb":' + ' 10.0, "loaded_modules": 22, "loaded_lines": 120}' + ) + mock_stderr = "DeprecationWarning: some_pkg is deprecated" + mock_process = MagicMock(stdout=mock_stdout, stderr=mock_stderr) + + with patch("subprocess.run", return_value=mock_process): + _run_worker_and_parse(["python", "profiler.py"]) + captured = capsys.readouterr() + assert "DeprecationWarning: some_pkg is deprecated" in captured.err + + +# ===================================================================== +# 4. MASTER COORDINATION & BENCHMARK INTEGRATIONS +# ===================================================================== + + +@patch("profiler._run_worker_and_parse") +def test_run_master_skips_line_count_and_restores_metrics(mock_parse): + """Verifies master appends --skip-line-count and restores metrics.""" + run_data_1 = { + "loaded_modules": 5, + "loaded_lines": 1000, + "time_ms": 50.0, + "peak_ram_mb": 10.0, + "rss_ram_mb": 8.0, + } + run_data_2 = { + "loaded_modules": 5, + "loaded_lines": -1, + "time_ms": 40.0, + "peak_ram_mb": 10.0, + "rss_ram_mb": 8.0, + } + mock_parse.side_effect = [run_data_1, run_data_2] + + with patch("builtins.print"), patch("sys.stderr"): + run_master( + iterations=2, + target_module="dummy_mod", + cpu=NO_CPU_PINNING, + clear_cache=False, + ) + + second_cmd_args = mock_parse.call_args_list[1][0][0] + assert "--skip-line-count" in second_cmd_args + assert run_data_2["loaded_lines"] == 1000 + + +@patch("profiler._run_worker_and_parse") +def test_run_master_checks_non_deterministic_behavior(mock_parse, capsys): + """Verifies warnings are printed upon non-deterministic module loads.""" + mock_parse.side_effect = [ + { + "loaded_modules": 50, + "loaded_lines": 500, + "time_ms": 10.0, + "peak_ram_mb": 1.0, + "rss_ram_mb": 1.0, + }, + { + "loaded_modules": 55, + "loaded_lines": -1, + "time_ms": 10.0, + "peak_ram_mb": 1.0, + "rss_ram_mb": 1.0, + }, + ] + + run_master( + iterations=2, + target_module="dummy_mod", + cpu=NO_CPU_PINNING, + clear_cache=False, + ) + captured = capsys.readouterr() + assert "WARNING: Non-deterministic import behavior!" in captured.err + + +@patch("profiler._run_worker_and_parse") +def test_run_master_with_cpu_pinning(mock_parse): + """Verifies taskset command configuration on Linux platforms.""" + mock_parse.return_value = { + "loaded_modules": 10, + "loaded_lines": 500, + "time_ms": 12.0, + "peak_ram_mb": 1.0, + "rss_ram_mb": 1.0, + } + + with patch("sys.platform", "linux"), patch("builtins.print"), patch( + "sys.stderr" + ): + run_master( + iterations=1, target_module="dummy_mod", cpu=1, clear_cache=False + ) + + args = mock_parse.call_args[0][0] + assert "taskset" in args + assert "1" in args + + +@patch("profiler._run_worker_and_parse") +def test_run_master_writes_csv_and_diffs_baseline(mock_parse, tmp_path): + """Verifies benchmark results CSV writing, reading, and thresholds comparison.""" + mock_parse.return_value = { + "loaded_modules": 10, + "loaded_lines": 500, + "time_ms": 12.0, + "peak_ram_mb": 1.0, + "rss_ram_mb": 1.0, + } + + csv_file = tmp_path / "results.csv" + baseline_file = tmp_path / "baseline.csv" + + # Write mock baseline CSV data + with open(baseline_file, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow( + [ + "iteration", + "time_ms", + "peak_ram_mb", + "rss_ram_mb", + "loaded_modules", + "loaded_lines", + ] + ) + writer.writerow([1, 10.0, 1.0, 1.0, 10, 500]) + + with patch("builtins.print"): + run_master( + iterations=1, + target_module="dummy_mod", + cpu=NO_CPU_PINNING, + csv_path=str(csv_file), + diff_baseline=str(baseline_file), + diff_threshold=50.0, + clear_cache=False, + ) + + assert csv_file.exists() + + +def test_validate_module_name_valid(): + """Verifies validate_module_name passes valid identifiers.""" + from profiler import validate_module_name + assert validate_module_name("google.cloud.storage") == "google.cloud.storage" + + +def test_validate_module_name_invalid(): + """Verifies validate_module_name raises ArgumentTypeError for invalid identifiers.""" + import argparse + from profiler import validate_module_name + with pytest.raises(argparse.ArgumentTypeError): + validate_module_name("google.cloud; rm -rf /") + + +@patch("subprocess.run") +def test_run_cprofile(mock_run): + """Verifies run_cprofile executes cProfile subprocess.""" + from profiler import run_cprofile + mock_run.return_value = MagicMock(returncode=0) + with patch("pstats.Stats"), patch("builtins.print"): + run_cprofile("math") + assert mock_run.called + + +@patch("multiprocessing.get_context") +def test_run_mprofile(mock_context): + """Verifies run_mprofile spawns process for memory snapshot.""" + from profiler import run_mprofile + mock_proc = MagicMock(exitcode=0) + mock_context.return_value.Process.return_value = mock_proc + with patch("builtins.print"): + run_mprofile("math") + assert mock_proc.start.called + assert mock_proc.join.called +