From 19787af876715cd886f0ec3d4b4849abab99f62f Mon Sep 17 00:00:00 2001 From: Tarek Ziade Date: Fri, 27 Mar 2026 09:31:19 +0100 Subject: [PATCH 1/7] investigate modular conversion speedups --- utils/modular_model_converter.py | 6 +- utils/profile_modular_conversion.py | 193 ++++++++++++++++++++++++++++ 2 files changed, 196 insertions(+), 3 deletions(-) create mode 100644 utils/profile_modular_conversion.py diff --git a/utils/modular_model_converter.py b/utils/modular_model_converter.py index 018316680ece..fa39cfecd7c5 100644 --- a/utils/modular_model_converter.py +++ b/utils/modular_model_converter.py @@ -855,7 +855,7 @@ def merge_modular_dependencies(self, classes, functions, assignments, object_map def visit_and_merge_dependencies( cls, module: cst.Module, classes, functions, assignments, object_mapping, start_lines ) -> "ModelFileMapper": - wrapper = MetadataWrapper(module) + wrapper = MetadataWrapper(module, unsafe_skip_copy=True) mapper = cls(module) wrapper.visit(mapper) # Merge dependencies @@ -1159,7 +1159,7 @@ def get_needed_imports(body: dict[str, dict], all_imports: list[cst.CSTNode]) -> Note: we need to use `isinstance` on scope assignments, m.matches apparently does not work here yet! """ new_body = [k[1]["node"] for k in sorted(body.items(), key=lambda x: x[1]["insert_idx"])] - wrapper = MetadataWrapper(cst.Module(body=all_imports + new_body)) + wrapper = MetadataWrapper(cst.Module(body=all_imports + new_body), unsafe_skip_copy=True) scopes = set(wrapper.resolve(ScopeProvider).values()) unused_imports = set() import_ref_count = defaultdict(lambda: 0) @@ -1948,7 +1948,7 @@ def convert_modular_file(modular_file: str, source_library: str | None = "transf if source_library != "transformers": module = module.visit(AbsoluteImportTransformer(relative_path, source_library)) - wrapper = MetadataWrapper(module) + wrapper = MetadataWrapper(module, unsafe_skip_copy=True) cst_transformers = ModularFileMapper(module, model_name, source_library) wrapper.visit(cst_transformers) for file, module in create_modules( diff --git a/utils/profile_modular_conversion.py b/utils/profile_modular_conversion.py new file mode 100644 index 000000000000..b98f974829dc --- /dev/null +++ b/utils/profile_modular_conversion.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python + +""" +Benchmark and profile modular conversion. + +Examples: + python utils/profile_modular_conversion.py \ + src/transformers/models/conditional_detr/modular_conditional_detr.py + + python utils/profile_modular_conversion.py \ + src/transformers/models/conditional_detr/modular_conditional_detr.py \ + --repeat 3 \ + --compare-old-wrapper \ + --verify-equal \ + --snakeviz +""" + +import argparse +import cProfile +import json +import re +import shutil +import statistics +import time +from contextlib import contextmanager +from pathlib import Path +from pstats import Stats + +import modular_model_converter as mmc +from libcst.metadata import MetadataWrapper as LibCSTMetadataWrapper + + +class OldBehaviorMetadataWrapper(LibCSTMetadataWrapper): + def __init__(self, module, unsafe_skip_copy=False, cache=None): + if cache is None: + cache = {} + super().__init__(module, unsafe_skip_copy=False, cache=cache) + + +@contextmanager +def metadata_wrapper_context(wrapper_cls): + original = mmc.MetadataWrapper + mmc.MetadataWrapper = wrapper_cls + try: + yield + finally: + mmc.MetadataWrapper = original + + +def run_conversion(modular_file: str, wrapper_cls) -> dict[str, str]: + with metadata_wrapper_context(wrapper_cls): + return mmc.convert_modular_file(modular_file) + + +def benchmark_conversion(modular_file: str, wrapper_cls, repeat: int) -> list[float]: + durations = [] + for _ in range(repeat): + start = time.perf_counter() + run_conversion(modular_file, wrapper_cls) + durations.append(time.perf_counter() - start) + return durations + + +def profile_conversion(modular_file: str, wrapper_cls, output_path: Path) -> None: + profile = cProfile.Profile() + profile.enable() + run_conversion(modular_file, wrapper_cls) + profile.disable() + profile.dump_stats(str(output_path)) + + +def write_snakeviz_bundle(profile_path: Path, bundle_dir: Path, profile_name: str) -> Path: + try: + import snakeviz + from snakeviz.stats import json_stats, table_rows + from tornado.template import Template + except ImportError as error: + raise RuntimeError("SnakeViz bundle generation requires snakeviz and tornado to be installed.") from error + + if bundle_dir.exists(): + shutil.rmtree(bundle_dir) + bundle_dir.mkdir(parents=True) + + static_dir = bundle_dir / "static" + snakeviz_root = Path(snakeviz.__file__).resolve().parent + shutil.copytree(snakeviz_root / "static", static_dir) + + stats = Stats(str(profile_path)) + template_text = (snakeviz_root / "templates" / "viz.html").read_text(encoding="utf-8") + html = Template(template_text).generate( + profile_name=profile_name, + table_rows=json.dumps(table_rows(stats)), + callees=json.dumps(json_stats(stats)), + ).decode("utf-8") + + static_uri = static_dir.resolve().as_uri() + html = html.replace('href="/static/', f'href="{static_uri}/') + html = html.replace("href='/static/", f"href='{static_uri}/") + html = html.replace('src="/static/', f'src="{static_uri}/') + html = html.replace("src='/static/", f"src='{static_uri}/") + html = html.replace( + 'event.data[\'url\'] + "/static/vendor/immutable.min.js",', + f'"{(static_dir / "vendor" / "immutable.min.js").resolve().as_uri()}",', + ) + html = html.replace( + 'event.data[\'url\'] + "/static/vendor/lodash.min.js");', + f'"{(static_dir / "vendor" / "lodash.min.js").resolve().as_uri()}");', + ) + html = html.replace("'url': window.location.origin", "'url': ''") + + html_path = bundle_dir / "index.html" + html_path.write_text(html, encoding="utf-8") + return html_path + + +def make_output_root(output_dir: str | None, modular_file: str) -> Path: + if output_dir is not None: + return Path(output_dir) + + match = re.search(r"modular_(.*)(?=\.py$)", modular_file) + name = match.group(1) if match is not None else Path(modular_file).stem + return Path("/tmp/codex-profiles") / f"{name}_modular_profile" + + +def summarize(label: str, durations: list[float]) -> str: + mean = statistics.mean(durations) + median = statistics.median(durations) + minimum = min(durations) + maximum = max(durations) + return ( + f"{label}: mean={mean:.3f}s median={median:.3f}s min={minimum:.3f}s " + f"max={maximum:.3f}s runs={len(durations)}" + ) + + +def main() -> int: + parser = argparse.ArgumentParser(description="Benchmark and profile modular conversion.") + parser.add_argument("modular_file", help="Path to a modular_*.py file.") + parser.add_argument("--repeat", type=int, default=3, help="Number of timing runs per mode.") + parser.add_argument( + "--compare-old-wrapper", + action="store_true", + help="Compare against simulated pre-change MetadataWrapper behavior.", + ) + parser.add_argument( + "--verify-equal", + action="store_true", + help="Verify that current and old-wrapper outputs are byte-identical.", + ) + parser.add_argument( + "--snakeviz", + action="store_true", + help="Generate a standalone SnakeViz HTML bundle for the current mode.", + ) + parser.add_argument( + "--output-dir", + help="Directory for profile outputs. Defaults to /tmp/codex-profiles/_modular_profile.", + ) + args = parser.parse_args() + + modular_file = args.modular_file + output_root = make_output_root(args.output_dir, modular_file) + output_root.mkdir(parents=True, exist_ok=True) + + modes = [("current", LibCSTMetadataWrapper)] + if args.compare_old_wrapper: + modes.append(("old_wrapper", OldBehaviorMetadataWrapper)) + + for label, wrapper_cls in modes: + durations = benchmark_conversion(modular_file, wrapper_cls, args.repeat) + print(summarize(label, durations)) + + if args.compare_old_wrapper and args.verify_equal: + current = run_conversion(modular_file, LibCSTMetadataWrapper) + old = run_conversion(modular_file, OldBehaviorMetadataWrapper) + same = current.keys() == old.keys() and all(current[key] == old[key] for key in current) + print(f"byte_identical={same}") + if not same: + return 1 + + profile_path = output_root / "current.prof" + profile_conversion(modular_file, LibCSTMetadataWrapper, profile_path) + print(f"profile={profile_path}") + + if args.snakeviz: + html_path = write_snakeviz_bundle(profile_path, output_root / "snakeviz", modular_file) + print(f"snakeviz={html_path}") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 31c15a39cfbecbbb1e9ffc21291fb8b30bca0c8b Mon Sep 17 00:00:00 2001 From: Tarek Ziade Date: Fri, 27 Mar 2026 10:15:26 +0100 Subject: [PATCH 2/7] second optim / cache --- utils/modular_model_converter.py | 31 ++++++++-- utils/profile_modular_conversion.py | 91 ++++++++++++++++++++--------- 2 files changed, 87 insertions(+), 35 deletions(-) diff --git a/utils/modular_model_converter.py b/utils/modular_model_converter.py index fa39cfecd7c5..1aa9e20a54a5 100644 --- a/utils/modular_model_converter.py +++ b/utils/modular_model_converter.py @@ -49,16 +49,36 @@ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 """ +ENABLE_MODULE_SOURCE_CACHE = True +_MODULE_SOURCE_CACHE = {} -def get_module_source_from_name(module_name: str) -> str: - # Extract the source code from the module name + +def clear_module_source_cache(): + _MODULE_SOURCE_CACHE.clear() + + +def get_module_source_and_tree_from_name(module_name: str) -> tuple[str, cst.Module]: spec = importlib.util.find_spec(module_name) if spec is None or spec.origin is None: raise ValueError(f"Cannot open file associated with {module_name} module.") - with open(spec.origin, "r", encoding="utf-8") as file: + file_path = spec.origin + cache_key = (module_name, file_path) + mtime_ns = os.stat(file_path).st_mtime_ns + cached = _MODULE_SOURCE_CACHE.get(cache_key) + if ENABLE_MODULE_SOURCE_CACHE and cached is not None and cached[0] == mtime_ns: + return cached[1], cached[2] + + with open(file_path, "r", encoding="utf-8") as file: source_code = file.read() - return source_code + tree = cst.parse_module(source_code) + if ENABLE_MODULE_SOURCE_CACHE: + _MODULE_SOURCE_CACHE[cache_key] = (mtime_ns, source_code, tree) + return source_code, tree + + +def get_module_source_from_name(module_name: str) -> str: + return get_module_source_and_tree_from_name(module_name)[0] def preserve_case_replace(text, patterns: dict, default_name: str): @@ -1414,12 +1434,11 @@ def visit_ImportFrom(self, node: cst.ImportFrom) -> None: if not import_module.startswith("transformers"): import_module = "transformers." + import_module try: - source_code = get_module_source_from_name(import_module) + _, tree = get_module_source_and_tree_from_name(import_module) except ModuleNotFoundError as e: raise ModuleNotFoundError( f"Failed to visit import from for: {self.python_module.code_for_node(node)}. Tried to import {import_module} but failed." ) from e - tree = cst.parse_module(source_code) self.model_specific_modules[import_module] = tree imported_object = self.python_module.code_for_node(imported_.name) self.model_specific_imported_objects[imported_object] = import_module diff --git a/utils/profile_modular_conversion.py b/utils/profile_modular_conversion.py index b98f974829dc..da3911445337 100644 --- a/utils/profile_modular_conversion.py +++ b/utils/profile_modular_conversion.py @@ -9,8 +9,9 @@ python utils/profile_modular_conversion.py \ src/transformers/models/conditional_detr/modular_conditional_detr.py \ + src/transformers/models/deepseek_vl/modular_deepseek_vl.py \ --repeat 3 \ - --compare-old-wrapper \ + --compare-old-branch \ --verify-equal \ --snakeviz """ @@ -47,24 +48,36 @@ def metadata_wrapper_context(wrapper_cls): mmc.MetadataWrapper = original -def run_conversion(modular_file: str, wrapper_cls) -> dict[str, str]: - with metadata_wrapper_context(wrapper_cls): - return mmc.convert_modular_file(modular_file) +@contextmanager +def module_cache_context(enabled: bool): + original = mmc.ENABLE_MODULE_SOURCE_CACHE + mmc.ENABLE_MODULE_SOURCE_CACHE = enabled + mmc.clear_module_source_cache() + try: + yield + finally: + mmc.clear_module_source_cache() + mmc.ENABLE_MODULE_SOURCE_CACHE = original + + +def run_conversion(modular_files: list[str], wrapper_cls, cache_enabled: bool) -> dict[str, dict[str, str]]: + with metadata_wrapper_context(wrapper_cls), module_cache_context(cache_enabled): + return {modular_file: mmc.convert_modular_file(modular_file) for modular_file in modular_files} -def benchmark_conversion(modular_file: str, wrapper_cls, repeat: int) -> list[float]: +def benchmark_conversion(modular_files: list[str], wrapper_cls, cache_enabled: bool, repeat: int) -> list[float]: durations = [] for _ in range(repeat): start = time.perf_counter() - run_conversion(modular_file, wrapper_cls) + run_conversion(modular_files, wrapper_cls, cache_enabled) durations.append(time.perf_counter() - start) return durations -def profile_conversion(modular_file: str, wrapper_cls, output_path: Path) -> None: +def profile_conversion(modular_files: list[str], wrapper_cls, cache_enabled: bool, output_path: Path) -> None: profile = cProfile.Profile() profile.enable() - run_conversion(modular_file, wrapper_cls) + run_conversion(modular_files, wrapper_cls, cache_enabled) profile.disable() profile.dump_stats(str(output_path)) @@ -113,12 +126,15 @@ def write_snakeviz_bundle(profile_path: Path, bundle_dir: Path, profile_name: st return html_path -def make_output_root(output_dir: str | None, modular_file: str) -> Path: +def make_output_root(output_dir: str | None, modular_files: list[str]) -> Path: if output_dir is not None: return Path(output_dir) - match = re.search(r"modular_(.*)(?=\.py$)", modular_file) - name = match.group(1) if match is not None else Path(modular_file).stem + if len(modular_files) == 1: + match = re.search(r"modular_(.*)(?=\.py$)", modular_files[0]) + name = match.group(1) if match is not None else Path(modular_files[0]).stem + else: + name = f"batch_{len(modular_files)}_modular_files" return Path("/tmp/codex-profiles") / f"{name}_modular_profile" @@ -135,17 +151,27 @@ def summarize(label: str, durations: list[float]) -> str: def main() -> int: parser = argparse.ArgumentParser(description="Benchmark and profile modular conversion.") - parser.add_argument("modular_file", help="Path to a modular_*.py file.") + parser.add_argument("modular_files", nargs="+", help="One or more paths to modular_*.py files.") parser.add_argument("--repeat", type=int, default=3, help="Number of timing runs per mode.") + parser.add_argument( + "--compare-old-branch", + action="store_true", + help="Compare against simulated pre-change branch behavior (old wrapper + no module cache).", + ) parser.add_argument( "--compare-old-wrapper", action="store_true", - help="Compare against simulated pre-change MetadataWrapper behavior.", + help="Compare against only the pre-change MetadataWrapper behavior.", + ) + parser.add_argument( + "--compare-no-cache", + action="store_true", + help="Compare against current wrapper behavior with module source caching disabled.", ) parser.add_argument( "--verify-equal", action="store_true", - help="Verify that current and old-wrapper outputs are byte-identical.", + help="Verify that current output and all requested comparison modes are byte-identical.", ) parser.add_argument( "--snakeviz", @@ -158,32 +184,39 @@ def main() -> int: ) args = parser.parse_args() - modular_file = args.modular_file - output_root = make_output_root(args.output_dir, modular_file) + modular_files = args.modular_files + output_root = make_output_root(args.output_dir, modular_files) output_root.mkdir(parents=True, exist_ok=True) - modes = [("current", LibCSTMetadataWrapper)] + modes = [("current", LibCSTMetadataWrapper, True)] + if args.compare_old_branch: + modes.append(("old_branch", OldBehaviorMetadataWrapper, False)) if args.compare_old_wrapper: - modes.append(("old_wrapper", OldBehaviorMetadataWrapper)) + modes.append(("old_wrapper", OldBehaviorMetadataWrapper, True)) + if args.compare_no_cache: + modes.append(("no_cache", LibCSTMetadataWrapper, False)) - for label, wrapper_cls in modes: - durations = benchmark_conversion(modular_file, wrapper_cls, args.repeat) + print(f"files={len(modular_files)}") + for label, wrapper_cls, cache_enabled in modes: + durations = benchmark_conversion(modular_files, wrapper_cls, cache_enabled, args.repeat) print(summarize(label, durations)) - if args.compare_old_wrapper and args.verify_equal: - current = run_conversion(modular_file, LibCSTMetadataWrapper) - old = run_conversion(modular_file, OldBehaviorMetadataWrapper) - same = current.keys() == old.keys() and all(current[key] == old[key] for key in current) - print(f"byte_identical={same}") - if not same: - return 1 + if args.verify_equal and len(modes) > 1: + current = run_conversion(modular_files, LibCSTMetadataWrapper, True) + for label, wrapper_cls, cache_enabled in modes[1:]: + other = run_conversion(modular_files, wrapper_cls, cache_enabled) + same = current.keys() == other.keys() and all(current[key] == other[key] for key in current) + print(f"{label}_byte_identical={same}") + if not same: + return 1 profile_path = output_root / "current.prof" - profile_conversion(modular_file, LibCSTMetadataWrapper, profile_path) + profile_conversion(modular_files, LibCSTMetadataWrapper, True, profile_path) print(f"profile={profile_path}") if args.snakeviz: - html_path = write_snakeviz_bundle(profile_path, output_root / "snakeviz", modular_file) + profile_name = modular_files[0] if len(modular_files) == 1 else f"batch:{len(modular_files)} modular files" + html_path = write_snakeviz_bundle(profile_path, output_root / "snakeviz", profile_name) print(f"snakeviz={html_path}") return 0 From e975252def9b815168f630d3fd53aa9a50513d5d Mon Sep 17 00:00:00 2001 From: Tarek Ziade Date: Fri, 27 Mar 2026 10:44:26 +0100 Subject: [PATCH 3/7] Python-only fast path based on symtable plus a lightweight AST pass --- utils/modular_model_converter.py | 172 ++++++++++++++++++++++++---- utils/profile_modular_conversion.py | 62 +++++++--- 2 files changed, 196 insertions(+), 38 deletions(-) diff --git a/utils/modular_model_converter.py b/utils/modular_model_converter.py index 1aa9e20a54a5..c7fb555710a6 100644 --- a/utils/modular_model_converter.py +++ b/utils/modular_model_converter.py @@ -12,12 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. import argparse +import ast import glob import importlib import multiprocessing as mp import os import re import subprocess +import symtable from abc import ABC, abstractmethod from collections import Counter, defaultdict, deque from functools import partial @@ -50,6 +52,7 @@ """ ENABLE_MODULE_SOURCE_CACHE = True +ENABLE_FAST_IMPORT_ANALYSIS = True _MODULE_SOURCE_CACHE = {} @@ -1173,30 +1176,22 @@ def append_new_import_node( imports_to_keep.append(new_node) -def get_needed_imports(body: dict[str, dict], all_imports: list[cst.CSTNode]) -> list[cst.CSTNode]: - """Get all the imports needed in the `body`, from the list of `all_imports`. - `body` is a dict with the following structure `{str: {"insert_idx": int, "node": cst.CSTNode}}`. - Note: we need to use `isinstance` on scope assignments, m.matches apparently does not work here yet! - """ - new_body = [k[1]["node"] for k in sorted(body.items(), key=lambda x: x[1]["insert_idx"])] - wrapper = MetadataWrapper(cst.Module(body=all_imports + new_body), unsafe_skip_copy=True) - scopes = set(wrapper.resolve(ScopeProvider).values()) - unused_imports = set() - import_ref_count = defaultdict(lambda: 0) - for scope in scopes: - for assignment in scope.assignments: - node = assignment.node - if isinstance(assignment, cst.metadata.Assignment) and isinstance(node, (cst.Import, cst.ImportFrom)): - ref_count = len(assignment.references) - name = assignment.name - import_ref_count[name] = max(ref_count, import_ref_count[name]) - # Similar imports may be redefined, and only used between their 1st and 2nd definition so if we already have - # a ref count > 0 at any point, the imports is actually used - unused_imports = {name for name, count in import_ref_count.items() if count <= 0 or name in body} +def _get_import_names(all_imports: list[cst.CSTNode]) -> set[str]: + import_names = set() + for node in all_imports: + statements = node.body.body if isinstance(node, cst.If) else [node] + for statement in statements: + import_node = statement.body[0] + for alias in import_node.names: + import_names.add(alias.evaluated_alias or alias.evaluated_name) + return import_names + +def _build_needed_imports(all_imports: list[cst.CSTNode], unused_imports: set[str]) -> list[cst.CSTNode]: + """Build the final import block after pruning the unused names.""" imports_to_keep = [] - # We need to keep track of which names were already imported, because some import may be duplicated from multiple sources - # or be both protected and unprotected due to inconsistency between models + # We need to keep track of which names were already imported, because some import may be duplicated from multiple + # sources or be both protected and unprotected due to inconsistency between models. added_names = set() existing_protected_statements = set() # str repr of the import nodes - does not work with the nodes directly for node in all_imports: @@ -1219,6 +1214,139 @@ def get_needed_imports(body: dict[str, dict], all_imports: list[cst.CSTNode]) -> return usual_import_nodes + protected_import_nodes +def _collect_referenced_plain_names(scope_table: symtable.SymbolTable) -> set[str]: + """Collect referenced names across all scopes. + + We intentionally do not distinguish between top-level and nested imports here to preserve the current LibCST + behavior, which merges import reference counts for same-named imports across scopes. + """ + referenced_names = {name for name in scope_table.get_identifiers() if scope_table.lookup(name).is_referenced()} + for child in scope_table.get_children(): + referenced_names.update(_collect_referenced_plain_names(child)) + return referenced_names + + +def _get_attribute_chain(node: ast.Attribute) -> list[str] | None: + parts = [] + current = node + while isinstance(current, ast.Attribute): + parts.append(current.attr) + current = current.value + if isinstance(current, ast.Name): + parts.append(current.id) + return list(reversed(parts)) + return None + + +def _collect_referenced_names(tree: ast.AST, plain_name_filter: set[str] | None = None) -> set[str]: + referenced_names = set() + for node in ast.walk(tree): + if isinstance(node, ast.Attribute) and isinstance(node.ctx, ast.Load): + chain = _get_attribute_chain(node) + if chain is None: + continue + if plain_name_filter is not None and chain[0] not in plain_name_filter: + continue + for end in range(1, len(chain) + 1): + referenced_names.add(".".join(chain[:end])) + elif isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load): + if plain_name_filter is None or node.id in plain_name_filter: + referenced_names.add(node.id) + return referenced_names + + +def _iter_annotation_expressions(module_ast: ast.AST): + for node in ast.walk(module_ast): + if isinstance(node, ast.AnnAssign): + yield node.annotation + elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + for arg in (*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs): + if arg.annotation is not None: + yield arg.annotation + if node.args.vararg is not None and node.args.vararg.annotation is not None: + yield node.args.vararg.annotation + if node.args.kwarg is not None and node.args.kwarg.annotation is not None: + yield node.args.kwarg.annotation + if node.returns is not None: + yield node.returns + + +def _collect_string_annotation_names(module_ast: ast.AST) -> set[str]: + """Collect names used inside quoted annotations and type comments.""" + referenced_names = set() + for annotation in _iter_annotation_expressions(module_ast): + if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str): + try: + annotation_tree = ast.parse(annotation.value, mode="eval") + except SyntaxError: + continue + referenced_names.update(_collect_referenced_names(annotation_tree)) + + for node in ast.walk(module_ast): + type_comment = getattr(node, "type_comment", None) + if isinstance(type_comment, str): + try: + type_comment_tree = ast.parse(type_comment, mode="eval") + except SyntaxError: + continue + referenced_names.update(_collect_referenced_names(type_comment_tree)) + + return referenced_names + + +def _has_future_annotations_import(all_imports: list[cst.CSTNode]) -> bool: + for node in all_imports: + if not isinstance(node, cst.SimpleStatementLine): + continue + import_node = node.body[0] + if not isinstance(import_node, cst.ImportFrom): + continue + if not isinstance(import_node.module, cst.Name) or import_node.module.value != "__future__": + continue + for alias in import_node.names: + if alias.evaluated_name == "annotations": + return True + return False + + +def _get_needed_imports_with_scope_provider(body: dict[str, dict], all_imports: list[cst.CSTNode]) -> list[cst.CSTNode]: + """ScopeProvider-based fallback for modules where the fast path is unsafe.""" + new_body = [k[1]["node"] for k in sorted(body.items(), key=lambda x: x[1]["insert_idx"])] + wrapper = MetadataWrapper(cst.Module(body=all_imports + new_body), unsafe_skip_copy=True) + scopes = set(wrapper.resolve(ScopeProvider).values()) + import_ref_count = defaultdict(lambda: 0) + for scope in scopes: + for assignment in scope.assignments: + node = assignment.node + if isinstance(assignment, cst.metadata.Assignment) and isinstance(node, (cst.Import, cst.ImportFrom)): + ref_count = len(assignment.references) + name = assignment.name + import_ref_count[name] = max(ref_count, import_ref_count[name]) + # Similar imports may be redefined, and only used between their 1st and 2nd definition so if we already have + # a ref count > 0 at any point, the imports is actually used + unused_imports = {name for name, count in import_ref_count.items() if count <= 0 or name in body} + return _build_needed_imports(all_imports, unused_imports) + + +def get_needed_imports(body: dict[str, dict], all_imports: list[cst.CSTNode]) -> list[cst.CSTNode]: + """Get all the imports needed in the `body`, from the list of `all_imports`. + `body` is a dict with the following structure `{str: {"insert_idx": int, "node": cst.CSTNode}}`. + """ + if not ENABLE_FAST_IMPORT_ANALYSIS or _has_future_annotations_import(all_imports): + return _get_needed_imports_with_scope_provider(body, all_imports) + + new_body = [k[1]["node"] for k in sorted(body.items(), key=lambda x: x[1]["insert_idx"])] + module = cst.Module(body=all_imports + new_body) + module_ast = ast.parse(module.code, type_comments=True) + referenced_plain_names = _collect_referenced_plain_names(symtable.symtable(module.code, "", "exec")) + referenced_names = _collect_referenced_names(module_ast, referenced_plain_names) + referenced_names.update(_collect_string_annotation_names(module_ast)) + + import_names = _get_import_names(all_imports) + unused_imports = {name for name in import_names if name not in referenced_names or name in body} + return _build_needed_imports(all_imports, unused_imports) + + def _ensure_utils_availability_imports(imports: list[cst.CSTNode], needed: set[str]) -> list[cst.CSTNode]: """Add is_torch_available and/or is_torchvision_available to the utils import if needed.""" if not needed: diff --git a/utils/profile_modular_conversion.py b/utils/profile_modular_conversion.py index da3911445337..7b48e6e2fcaa 100644 --- a/utils/profile_modular_conversion.py +++ b/utils/profile_modular_conversion.py @@ -11,6 +11,7 @@ src/transformers/models/conditional_detr/modular_conditional_detr.py \ src/transformers/models/deepseek_vl/modular_deepseek_vl.py \ --repeat 3 \ + --compare-old-import-analysis \ --compare-old-branch \ --verify-equal \ --snakeviz @@ -60,24 +61,44 @@ def module_cache_context(enabled: bool): mmc.ENABLE_MODULE_SOURCE_CACHE = original -def run_conversion(modular_files: list[str], wrapper_cls, cache_enabled: bool) -> dict[str, dict[str, str]]: - with metadata_wrapper_context(wrapper_cls), module_cache_context(cache_enabled): +@contextmanager +def fast_import_analysis_context(enabled: bool): + original = mmc.ENABLE_FAST_IMPORT_ANALYSIS + mmc.ENABLE_FAST_IMPORT_ANALYSIS = enabled + try: + yield + finally: + mmc.ENABLE_FAST_IMPORT_ANALYSIS = original + + +def run_conversion( + modular_files: list[str], wrapper_cls, cache_enabled: bool, fast_import_analysis_enabled: bool +) -> dict[str, dict[str, str]]: + with ( + metadata_wrapper_context(wrapper_cls), + module_cache_context(cache_enabled), + fast_import_analysis_context(fast_import_analysis_enabled), + ): return {modular_file: mmc.convert_modular_file(modular_file) for modular_file in modular_files} -def benchmark_conversion(modular_files: list[str], wrapper_cls, cache_enabled: bool, repeat: int) -> list[float]: +def benchmark_conversion( + modular_files: list[str], wrapper_cls, cache_enabled: bool, fast_import_analysis_enabled: bool, repeat: int +) -> list[float]: durations = [] for _ in range(repeat): start = time.perf_counter() - run_conversion(modular_files, wrapper_cls, cache_enabled) + run_conversion(modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled) durations.append(time.perf_counter() - start) return durations -def profile_conversion(modular_files: list[str], wrapper_cls, cache_enabled: bool, output_path: Path) -> None: +def profile_conversion( + modular_files: list[str], wrapper_cls, cache_enabled: bool, fast_import_analysis_enabled: bool, output_path: Path +) -> None: profile = cProfile.Profile() profile.enable() - run_conversion(modular_files, wrapper_cls, cache_enabled) + run_conversion(modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled) profile.disable() profile.dump_stats(str(output_path)) @@ -168,6 +189,11 @@ def main() -> int: action="store_true", help="Compare against current wrapper behavior with module source caching disabled.", ) + parser.add_argument( + "--compare-old-import-analysis", + action="store_true", + help="Compare against the current branch with the previous get_needed_imports analysis.", + ) parser.add_argument( "--verify-equal", action="store_true", @@ -188,30 +214,34 @@ def main() -> int: output_root = make_output_root(args.output_dir, modular_files) output_root.mkdir(parents=True, exist_ok=True) - modes = [("current", LibCSTMetadataWrapper, True)] + modes = [("current", LibCSTMetadataWrapper, True, True)] if args.compare_old_branch: - modes.append(("old_branch", OldBehaviorMetadataWrapper, False)) + modes.append(("old_branch", OldBehaviorMetadataWrapper, False, False)) if args.compare_old_wrapper: - modes.append(("old_wrapper", OldBehaviorMetadataWrapper, True)) + modes.append(("old_wrapper", OldBehaviorMetadataWrapper, True, True)) if args.compare_no_cache: - modes.append(("no_cache", LibCSTMetadataWrapper, False)) + modes.append(("no_cache", LibCSTMetadataWrapper, False, True)) + if args.compare_old_import_analysis: + modes.append(("old_import_analysis", LibCSTMetadataWrapper, True, False)) print(f"files={len(modular_files)}") - for label, wrapper_cls, cache_enabled in modes: - durations = benchmark_conversion(modular_files, wrapper_cls, cache_enabled, args.repeat) + for label, wrapper_cls, cache_enabled, fast_import_analysis_enabled in modes: + durations = benchmark_conversion( + modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled, args.repeat + ) print(summarize(label, durations)) if args.verify_equal and len(modes) > 1: - current = run_conversion(modular_files, LibCSTMetadataWrapper, True) - for label, wrapper_cls, cache_enabled in modes[1:]: - other = run_conversion(modular_files, wrapper_cls, cache_enabled) + current = run_conversion(modular_files, LibCSTMetadataWrapper, True, True) + for label, wrapper_cls, cache_enabled, fast_import_analysis_enabled in modes[1:]: + other = run_conversion(modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled) same = current.keys() == other.keys() and all(current[key] == other[key] for key in current) print(f"{label}_byte_identical={same}") if not same: return 1 profile_path = output_root / "current.prof" - profile_conversion(modular_files, LibCSTMetadataWrapper, True, profile_path) + profile_conversion(modular_files, LibCSTMetadataWrapper, True, True, profile_path) print(f"profile={profile_path}") if args.snakeviz: From 361b0a3ab24f499d31b318e25a9fbc54bc339b6f Mon Sep 17 00:00:00 2001 From: Tarek Ziade Date: Fri, 27 Mar 2026 11:22:12 +0100 Subject: [PATCH 4/7] replaced the imported-module analysis path in convert_modular_file() with a metadata-free mapper visit that tracks top-level context and source order directly, eliminating expensive LibCST ParentNodeProvider/PositionProvider passes while preserving byte-identical output --- utils/modular_model_converter.py | 88 +++++++++++++++++++++----- utils/profile_modular_conversion.py | 95 +++++++++++++++++++++-------- 2 files changed, 143 insertions(+), 40 deletions(-) diff --git a/utils/modular_model_converter.py b/utils/modular_model_converter.py index c7fb555710a6..89962a551b8e 100644 --- a/utils/modular_model_converter.py +++ b/utils/modular_model_converter.py @@ -53,6 +53,7 @@ ENABLE_MODULE_SOURCE_CACHE = True ENABLE_FAST_IMPORT_ANALYSIS = True +ENABLE_FAST_MAPPER_VISIT = True _MODULE_SOURCE_CACHE = {} @@ -570,12 +571,35 @@ def __init__(self, python_module: cst.Module): self.current_function = None # this keeps track of the current module-scope function self.current_class = None # this keeps track of the current module-scope class self.current_assignment = None # this keeps track of the current module-scope assignment + self._suite_depth = 0 # tracks whether we are inside an indented/simple statement suite + self._node_order = {} # source order of recorded top-level nodes + self._next_node_order = 0 # this keeps track of objects imported from modeling files (`from .configuration import Config`) -> `Config` should not be a dependency self.objects_imported_from_modeling = set() # regex pattern joining every possible file type self.match_patterns = "|".join(ALL_FILE_TYPES) # fmt: on + def _is_direct_module_child(self) -> bool: + return self._suite_depth == 0 + + def _record_node_order(self, node_name: str) -> None: + if node_name not in self._node_order: + self._node_order[node_name] = self._next_node_order + self._next_node_order += 1 + + def visit_IndentedBlock(self, node): + self._suite_depth += 1 + + def leave_IndentedBlock(self, node): + self._suite_depth -= 1 + + def visit_SimpleStatementSuite(self, node): + self._suite_depth += 1 + + def leave_SimpleStatementSuite(self, node): + self._suite_depth -= 1 + def visit_ImportFrom(self, node): """This keeps track of objects imported from neighbor modeling files (e.g. in `modeling_xxx.py, we have `from .configuration_xxx import Config`, then `Config` should be recorded as it is not a dependency that needs @@ -596,7 +620,6 @@ def visit_SimpleStatementLine(self, node): Global Assigns like `GEMMA_INPUT_DOCSTRING = 'THIS IS THE INPUT'` and all import statements are extracted and saved in their corresponding dict. They are then used when updating dependency mappings. """ - parent_node = self.get_metadata(cst.metadata.ParentNodeProvider, node) simple_top_level_assign_structure = m.SimpleStatementLine( body=[m.Assign(targets=[m.AssignTarget(target=m.Name())])] ) @@ -604,11 +627,17 @@ def visit_SimpleStatementLine(self, node): body=[m.Assign(targets=[m.AssignTarget(target=m.Subscript(value=m.Name()) | m.Attribute(value=m.Name()))])] ) - if m.matches(parent_node, m.Module()): + is_module_level = ( + self._is_direct_module_child() + if ENABLE_FAST_MAPPER_VISIT + else m.matches(self.get_metadata(cst.metadata.ParentNodeProvider, node), m.Module()) + ) + if is_module_level: if m.matches(node, simple_top_level_assign_structure): left_hand_side = node.body[0].targets[0].target.value self.current_assignment = left_hand_side self.assignments[left_hand_side] = node + self._record_node_order(left_hand_side) # This corresponds to a global variable being indexed or having an attribute look-up elif m.matches(node, simple_top_level_variable_indexing): indexed_variable = node.body[0].targets[0].target.value.value @@ -618,6 +647,7 @@ def visit_SimpleStatementLine(self, node): node_name = self.python_module.code_for_node(node) self.assignments[node_name] = node self.object_dependency_mapping[indexed_variable].add(node_name) + self._record_node_order(node_name) elif m.matches(node, m.SimpleStatementLine(body=[m.Import() | m.ImportFrom()])): self.imports.append(node) @@ -627,14 +657,23 @@ def leave_SimpleStatementLine(self, node): self.current_assignment = None def visit_FunctionDef(self, node): - parent_node = self.get_metadata(cst.metadata.ParentNodeProvider, node) - if m.matches(parent_node, m.Module()): + is_module_level = ( + self._is_direct_module_child() + if ENABLE_FAST_MAPPER_VISIT + else m.matches(self.get_metadata(cst.metadata.ParentNodeProvider, node), m.Module()) + ) + if is_module_level: self.current_function = node.name.value self.functions[node.name.value] = node + self._record_node_order(node.name.value) def leave_FunctionDef(self, node): - parent_node = self.get_metadata(cst.metadata.ParentNodeProvider, node) - if m.matches(parent_node, m.Module()): + is_module_level = ( + self._is_direct_module_child() + if ENABLE_FAST_MAPPER_VISIT + else m.matches(self.get_metadata(cst.metadata.ParentNodeProvider, node), m.Module()) + ) + if is_module_level: self.current_function = None def visit_If(self, node): @@ -647,6 +686,7 @@ def visit_If(self, node): def visit_ClassDef(self, node: ClassDef) -> None: """Record class nodes to create their dependencies at the end.""" self.classes[node.name.value] = node + self._record_node_order(node.name.value) self.current_class = node.name.value def leave_ClassDef(self, node): @@ -669,8 +709,12 @@ def leave_Module(self, node): self.global_nodes = {**self.assignments, **self.classes, **self.functions} # now sort the class dependency_mapping based on the position of the nodes self.start_lines = {} - for id, node in self.global_nodes.items(): - self.start_lines[id] = self.get_metadata(cst.metadata.PositionProvider, node).start.line + if ENABLE_FAST_MAPPER_VISIT: + for node_name in self.global_nodes: + self.start_lines[node_name] = self._node_order[node_name] + else: + for node_name, current_node in self.global_nodes.items(): + self.start_lines[node_name] = self.get_metadata(cst.metadata.PositionProvider, current_node).start.line def _restrict_dependencies_to_known_entities(self): """Since we added every Name as part of `self.object_dependency_mapping`, we need to remove those that @@ -878,9 +922,12 @@ def merge_modular_dependencies(self, classes, functions, assignments, object_map def visit_and_merge_dependencies( cls, module: cst.Module, classes, functions, assignments, object_mapping, start_lines ) -> "ModelFileMapper": - wrapper = MetadataWrapper(module, unsafe_skip_copy=True) mapper = cls(module) - wrapper.visit(mapper) + if ENABLE_FAST_MAPPER_VISIT: + module.visit(mapper) + else: + wrapper = MetadataWrapper(module, unsafe_skip_copy=True) + wrapper.visit(mapper) # Merge dependencies mapper.merge_modular_dependencies(classes, functions, assignments, object_mapping, start_lines) # Create the class dependencies graph @@ -1309,7 +1356,9 @@ def _has_future_annotations_import(all_imports: list[cst.CSTNode]) -> bool: return False -def _get_needed_imports_with_scope_provider(body: dict[str, dict], all_imports: list[cst.CSTNode]) -> list[cst.CSTNode]: +def _get_needed_imports_with_scope_provider( + body: dict[str, dict], all_imports: list[cst.CSTNode] +) -> list[cst.CSTNode]: """ScopeProvider-based fallback for modules where the fast path is unsafe.""" new_body = [k[1]["node"] for k in sorted(body.items(), key=lambda x: x[1]["insert_idx"])] wrapper = MetadataWrapper(cst.Module(body=all_imports + new_body), unsafe_skip_copy=True) @@ -1580,7 +1629,6 @@ def visit_SimpleStatementLine(self, node): """If we visit an import statement not previously visited, record it. If we visit a module-scope assignment, simply record it or, if it is `__all__`, split it between files where we should dispatch it. """ - parent_node = self.get_metadata(cst.metadata.ParentNodeProvider, node) simple_top_level_assign_structure = m.SimpleStatementLine( body=[m.Assign(targets=[m.AssignTarget(target=m.Name())])] ) @@ -1588,7 +1636,12 @@ def visit_SimpleStatementLine(self, node): body=[m.Assign(targets=[m.AssignTarget(target=m.Subscript(value=m.Name()) | m.Attribute(value=m.Name()))])] ) - if m.matches(parent_node, m.Module()): + is_module_level = ( + self._is_direct_module_child() + if ENABLE_FAST_MAPPER_VISIT + else m.matches(self.get_metadata(cst.metadata.ParentNodeProvider, node), m.Module()) + ) + if is_module_level: if m.matches(node, m.SimpleStatementLine(body=[m.Import()])): self.imports.append(node) elif m.matches(node, m.SimpleStatementLine(body=[m.ImportFrom()])): @@ -1612,6 +1665,7 @@ def visit_SimpleStatementLine(self, node): else: self.current_assignment = assigned_variable self.assignments[assigned_variable] = node + self._record_node_order(assigned_variable) # This corresponds to a global variable being indexed or having an attribute look-up elif m.matches(node, simple_top_level_variable_indexing): indexed_variable = node.body[0].targets[0].target.value.value @@ -1621,6 +1675,7 @@ def visit_SimpleStatementLine(self, node): node_name = self.python_module.code_for_node(node) self.assignments[node_name] = node self.object_dependency_mapping[indexed_variable].add(node_name) + self._record_node_order(node_name) def leave_Module(self, node): """When we leave the modular file, we do the following in order: @@ -2095,9 +2150,12 @@ def convert_modular_file(modular_file: str, source_library: str | None = "transf if source_library != "transformers": module = module.visit(AbsoluteImportTransformer(relative_path, source_library)) - wrapper = MetadataWrapper(module, unsafe_skip_copy=True) cst_transformers = ModularFileMapper(module, model_name, source_library) - wrapper.visit(cst_transformers) + if ENABLE_FAST_MAPPER_VISIT: + module.visit(cst_transformers) + else: + wrapper = MetadataWrapper(module, unsafe_skip_copy=True) + wrapper.visit(cst_transformers) for file, module in create_modules( cst_transformers, file_path=relative_path, package_name=source_library ).items(): diff --git a/utils/profile_modular_conversion.py b/utils/profile_modular_conversion.py index 7b48e6e2fcaa..6e9c5c23c9c9 100644 --- a/utils/profile_modular_conversion.py +++ b/utils/profile_modular_conversion.py @@ -12,6 +12,7 @@ src/transformers/models/deepseek_vl/modular_deepseek_vl.py \ --repeat 3 \ --compare-old-import-analysis \ + --compare-old-mapper-visit \ --compare-old-branch \ --verify-equal \ --snakeviz @@ -71,34 +72,61 @@ def fast_import_analysis_context(enabled: bool): mmc.ENABLE_FAST_IMPORT_ANALYSIS = original +@contextmanager +def fast_mapper_visit_context(enabled: bool): + original = mmc.ENABLE_FAST_MAPPER_VISIT + mmc.ENABLE_FAST_MAPPER_VISIT = enabled + try: + yield + finally: + mmc.ENABLE_FAST_MAPPER_VISIT = original + + def run_conversion( - modular_files: list[str], wrapper_cls, cache_enabled: bool, fast_import_analysis_enabled: bool + modular_files: list[str], + wrapper_cls, + cache_enabled: bool, + fast_import_analysis_enabled: bool, + fast_mapper_visit_enabled: bool, ) -> dict[str, dict[str, str]]: with ( metadata_wrapper_context(wrapper_cls), module_cache_context(cache_enabled), fast_import_analysis_context(fast_import_analysis_enabled), + fast_mapper_visit_context(fast_mapper_visit_enabled), ): return {modular_file: mmc.convert_modular_file(modular_file) for modular_file in modular_files} def benchmark_conversion( - modular_files: list[str], wrapper_cls, cache_enabled: bool, fast_import_analysis_enabled: bool, repeat: int + modular_files: list[str], + wrapper_cls, + cache_enabled: bool, + fast_import_analysis_enabled: bool, + fast_mapper_visit_enabled: bool, + repeat: int, ) -> list[float]: durations = [] for _ in range(repeat): start = time.perf_counter() - run_conversion(modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled) + run_conversion( + modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled, fast_mapper_visit_enabled + ) durations.append(time.perf_counter() - start) return durations def profile_conversion( - modular_files: list[str], wrapper_cls, cache_enabled: bool, fast_import_analysis_enabled: bool, output_path: Path + modular_files: list[str], + wrapper_cls, + cache_enabled: bool, + fast_import_analysis_enabled: bool, + fast_mapper_visit_enabled: bool, + output_path: Path, ) -> None: profile = cProfile.Profile() profile.enable() - run_conversion(modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled) + run_conversion(modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled, fast_mapper_visit_enabled) profile.disable() profile.dump_stats(str(output_path)) @@ -121,11 +149,15 @@ def write_snakeviz_bundle(profile_path: Path, bundle_dir: Path, profile_name: st stats = Stats(str(profile_path)) template_text = (snakeviz_root / "templates" / "viz.html").read_text(encoding="utf-8") - html = Template(template_text).generate( - profile_name=profile_name, - table_rows=json.dumps(table_rows(stats)), - callees=json.dumps(json_stats(stats)), - ).decode("utf-8") + html = ( + Template(template_text) + .generate( + profile_name=profile_name, + table_rows=json.dumps(table_rows(stats)), + callees=json.dumps(json_stats(stats)), + ) + .decode("utf-8") + ) static_uri = static_dir.resolve().as_uri() html = html.replace('href="/static/', f'href="{static_uri}/') @@ -133,11 +165,11 @@ def write_snakeviz_bundle(profile_path: Path, bundle_dir: Path, profile_name: st html = html.replace('src="/static/', f'src="{static_uri}/') html = html.replace("src='/static/", f"src='{static_uri}/") html = html.replace( - 'event.data[\'url\'] + "/static/vendor/immutable.min.js",', + "event.data['url'] + \"/static/vendor/immutable.min.js\",", f'"{(static_dir / "vendor" / "immutable.min.js").resolve().as_uri()}",', ) html = html.replace( - 'event.data[\'url\'] + "/static/vendor/lodash.min.js");', + "event.data['url'] + \"/static/vendor/lodash.min.js\");", f'"{(static_dir / "vendor" / "lodash.min.js").resolve().as_uri()}");', ) html = html.replace("'url': window.location.origin", "'url': ''") @@ -165,8 +197,7 @@ def summarize(label: str, durations: list[float]) -> str: minimum = min(durations) maximum = max(durations) return ( - f"{label}: mean={mean:.3f}s median={median:.3f}s min={minimum:.3f}s " - f"max={maximum:.3f}s runs={len(durations)}" + f"{label}: mean={mean:.3f}s median={median:.3f}s min={minimum:.3f}s max={maximum:.3f}s runs={len(durations)}" ) @@ -194,6 +225,11 @@ def main() -> int: action="store_true", help="Compare against the current branch with the previous get_needed_imports analysis.", ) + parser.add_argument( + "--compare-old-mapper-visit", + action="store_true", + help="Compare against the current branch with the previous metadata-based module visit path.", + ) parser.add_argument( "--verify-equal", action="store_true", @@ -214,34 +250,43 @@ def main() -> int: output_root = make_output_root(args.output_dir, modular_files) output_root.mkdir(parents=True, exist_ok=True) - modes = [("current", LibCSTMetadataWrapper, True, True)] + modes = [("current", LibCSTMetadataWrapper, True, True, True)] if args.compare_old_branch: - modes.append(("old_branch", OldBehaviorMetadataWrapper, False, False)) + modes.append(("old_branch", OldBehaviorMetadataWrapper, False, False, False)) if args.compare_old_wrapper: - modes.append(("old_wrapper", OldBehaviorMetadataWrapper, True, True)) + modes.append(("old_wrapper", OldBehaviorMetadataWrapper, True, True, True)) if args.compare_no_cache: - modes.append(("no_cache", LibCSTMetadataWrapper, False, True)) + modes.append(("no_cache", LibCSTMetadataWrapper, False, True, True)) if args.compare_old_import_analysis: - modes.append(("old_import_analysis", LibCSTMetadataWrapper, True, False)) + modes.append(("old_import_analysis", LibCSTMetadataWrapper, True, False, True)) + if args.compare_old_mapper_visit: + modes.append(("old_mapper_visit", LibCSTMetadataWrapper, True, True, False)) print(f"files={len(modular_files)}") - for label, wrapper_cls, cache_enabled, fast_import_analysis_enabled in modes: + for label, wrapper_cls, cache_enabled, fast_import_analysis_enabled, fast_mapper_visit_enabled in modes: durations = benchmark_conversion( - modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled, args.repeat + modular_files, + wrapper_cls, + cache_enabled, + fast_import_analysis_enabled, + fast_mapper_visit_enabled, + args.repeat, ) print(summarize(label, durations)) if args.verify_equal and len(modes) > 1: - current = run_conversion(modular_files, LibCSTMetadataWrapper, True, True) - for label, wrapper_cls, cache_enabled, fast_import_analysis_enabled in modes[1:]: - other = run_conversion(modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled) + current = run_conversion(modular_files, LibCSTMetadataWrapper, True, True, True) + for label, wrapper_cls, cache_enabled, fast_import_analysis_enabled, fast_mapper_visit_enabled in modes[1:]: + other = run_conversion( + modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled, fast_mapper_visit_enabled + ) same = current.keys() == other.keys() and all(current[key] == other[key] for key in current) print(f"{label}_byte_identical={same}") if not same: return 1 profile_path = output_root / "current.prof" - profile_conversion(modular_files, LibCSTMetadataWrapper, True, True, profile_path) + profile_conversion(modular_files, LibCSTMetadataWrapper, True, True, True, profile_path) print(f"profile={profile_path}") if args.snakeviz: From 13e7969604aeb7eb23023d39509bd4dc03f0055f Mon Sep 17 00:00:00 2001 From: Tarek Ziade Date: Thu, 23 Apr 2026 10:52:23 +0200 Subject: [PATCH 5/7] revert some speedups --- utils/modular_model_converter.py | 134 +--------------------------- utils/profile_modular_conversion.py | 51 +++-------- 2 files changed, 15 insertions(+), 170 deletions(-) diff --git a/utils/modular_model_converter.py b/utils/modular_model_converter.py index 89962a551b8e..d0b42771039f 100644 --- a/utils/modular_model_converter.py +++ b/utils/modular_model_converter.py @@ -12,14 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. import argparse -import ast import glob import importlib import multiprocessing as mp import os import re import subprocess -import symtable from abc import ABC, abstractmethod from collections import Counter, defaultdict, deque from functools import partial @@ -52,7 +50,6 @@ """ ENABLE_MODULE_SOURCE_CACHE = True -ENABLE_FAST_IMPORT_ANALYSIS = True ENABLE_FAST_MAPPER_VISIT = True _MODULE_SOURCE_CACHE = {} @@ -1223,17 +1220,6 @@ def append_new_import_node( imports_to_keep.append(new_node) -def _get_import_names(all_imports: list[cst.CSTNode]) -> set[str]: - import_names = set() - for node in all_imports: - statements = node.body.body if isinstance(node, cst.If) else [node] - for statement in statements: - import_node = statement.body[0] - for alias in import_node.names: - import_names.add(alias.evaluated_alias or alias.evaluated_name) - return import_names - - def _build_needed_imports(all_imports: list[cst.CSTNode], unused_imports: set[str]) -> list[cst.CSTNode]: """Build the final import block after pruning the unused names.""" imports_to_keep = [] @@ -1261,105 +1247,10 @@ def _build_needed_imports(all_imports: list[cst.CSTNode], unused_imports: set[st return usual_import_nodes + protected_import_nodes -def _collect_referenced_plain_names(scope_table: symtable.SymbolTable) -> set[str]: - """Collect referenced names across all scopes. - - We intentionally do not distinguish between top-level and nested imports here to preserve the current LibCST - behavior, which merges import reference counts for same-named imports across scopes. +def get_needed_imports(body: dict[str, dict], all_imports: list[cst.CSTNode]) -> list[cst.CSTNode]: + """Get all the imports needed in the `body`, from the list of `all_imports`. + `body` is a dict with the following structure `{str: {"insert_idx": int, "node": cst.CSTNode}}`. """ - referenced_names = {name for name in scope_table.get_identifiers() if scope_table.lookup(name).is_referenced()} - for child in scope_table.get_children(): - referenced_names.update(_collect_referenced_plain_names(child)) - return referenced_names - - -def _get_attribute_chain(node: ast.Attribute) -> list[str] | None: - parts = [] - current = node - while isinstance(current, ast.Attribute): - parts.append(current.attr) - current = current.value - if isinstance(current, ast.Name): - parts.append(current.id) - return list(reversed(parts)) - return None - - -def _collect_referenced_names(tree: ast.AST, plain_name_filter: set[str] | None = None) -> set[str]: - referenced_names = set() - for node in ast.walk(tree): - if isinstance(node, ast.Attribute) and isinstance(node.ctx, ast.Load): - chain = _get_attribute_chain(node) - if chain is None: - continue - if plain_name_filter is not None and chain[0] not in plain_name_filter: - continue - for end in range(1, len(chain) + 1): - referenced_names.add(".".join(chain[:end])) - elif isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load): - if plain_name_filter is None or node.id in plain_name_filter: - referenced_names.add(node.id) - return referenced_names - - -def _iter_annotation_expressions(module_ast: ast.AST): - for node in ast.walk(module_ast): - if isinstance(node, ast.AnnAssign): - yield node.annotation - elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): - for arg in (*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs): - if arg.annotation is not None: - yield arg.annotation - if node.args.vararg is not None and node.args.vararg.annotation is not None: - yield node.args.vararg.annotation - if node.args.kwarg is not None and node.args.kwarg.annotation is not None: - yield node.args.kwarg.annotation - if node.returns is not None: - yield node.returns - - -def _collect_string_annotation_names(module_ast: ast.AST) -> set[str]: - """Collect names used inside quoted annotations and type comments.""" - referenced_names = set() - for annotation in _iter_annotation_expressions(module_ast): - if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str): - try: - annotation_tree = ast.parse(annotation.value, mode="eval") - except SyntaxError: - continue - referenced_names.update(_collect_referenced_names(annotation_tree)) - - for node in ast.walk(module_ast): - type_comment = getattr(node, "type_comment", None) - if isinstance(type_comment, str): - try: - type_comment_tree = ast.parse(type_comment, mode="eval") - except SyntaxError: - continue - referenced_names.update(_collect_referenced_names(type_comment_tree)) - - return referenced_names - - -def _has_future_annotations_import(all_imports: list[cst.CSTNode]) -> bool: - for node in all_imports: - if not isinstance(node, cst.SimpleStatementLine): - continue - import_node = node.body[0] - if not isinstance(import_node, cst.ImportFrom): - continue - if not isinstance(import_node.module, cst.Name) or import_node.module.value != "__future__": - continue - for alias in import_node.names: - if alias.evaluated_name == "annotations": - return True - return False - - -def _get_needed_imports_with_scope_provider( - body: dict[str, dict], all_imports: list[cst.CSTNode] -) -> list[cst.CSTNode]: - """ScopeProvider-based fallback for modules where the fast path is unsafe.""" new_body = [k[1]["node"] for k in sorted(body.items(), key=lambda x: x[1]["insert_idx"])] wrapper = MetadataWrapper(cst.Module(body=all_imports + new_body), unsafe_skip_copy=True) scopes = set(wrapper.resolve(ScopeProvider).values()) @@ -1377,25 +1268,6 @@ def _get_needed_imports_with_scope_provider( return _build_needed_imports(all_imports, unused_imports) -def get_needed_imports(body: dict[str, dict], all_imports: list[cst.CSTNode]) -> list[cst.CSTNode]: - """Get all the imports needed in the `body`, from the list of `all_imports`. - `body` is a dict with the following structure `{str: {"insert_idx": int, "node": cst.CSTNode}}`. - """ - if not ENABLE_FAST_IMPORT_ANALYSIS or _has_future_annotations_import(all_imports): - return _get_needed_imports_with_scope_provider(body, all_imports) - - new_body = [k[1]["node"] for k in sorted(body.items(), key=lambda x: x[1]["insert_idx"])] - module = cst.Module(body=all_imports + new_body) - module_ast = ast.parse(module.code, type_comments=True) - referenced_plain_names = _collect_referenced_plain_names(symtable.symtable(module.code, "", "exec")) - referenced_names = _collect_referenced_names(module_ast, referenced_plain_names) - referenced_names.update(_collect_string_annotation_names(module_ast)) - - import_names = _get_import_names(all_imports) - unused_imports = {name for name in import_names if name not in referenced_names or name in body} - return _build_needed_imports(all_imports, unused_imports) - - def _ensure_utils_availability_imports(imports: list[cst.CSTNode], needed: set[str]) -> list[cst.CSTNode]: """Add is_torch_available and/or is_torchvision_available to the utils import if needed.""" if not needed: diff --git a/utils/profile_modular_conversion.py b/utils/profile_modular_conversion.py index 6e9c5c23c9c9..8ab2297e259f 100644 --- a/utils/profile_modular_conversion.py +++ b/utils/profile_modular_conversion.py @@ -11,7 +11,6 @@ src/transformers/models/conditional_detr/modular_conditional_detr.py \ src/transformers/models/deepseek_vl/modular_deepseek_vl.py \ --repeat 3 \ - --compare-old-import-analysis \ --compare-old-mapper-visit \ --compare-old-branch \ --verify-equal \ @@ -62,16 +61,6 @@ def module_cache_context(enabled: bool): mmc.ENABLE_MODULE_SOURCE_CACHE = original -@contextmanager -def fast_import_analysis_context(enabled: bool): - original = mmc.ENABLE_FAST_IMPORT_ANALYSIS - mmc.ENABLE_FAST_IMPORT_ANALYSIS = enabled - try: - yield - finally: - mmc.ENABLE_FAST_IMPORT_ANALYSIS = original - - @contextmanager def fast_mapper_visit_context(enabled: bool): original = mmc.ENABLE_FAST_MAPPER_VISIT @@ -86,13 +75,11 @@ def run_conversion( modular_files: list[str], wrapper_cls, cache_enabled: bool, - fast_import_analysis_enabled: bool, fast_mapper_visit_enabled: bool, ) -> dict[str, dict[str, str]]: with ( metadata_wrapper_context(wrapper_cls), module_cache_context(cache_enabled), - fast_import_analysis_context(fast_import_analysis_enabled), fast_mapper_visit_context(fast_mapper_visit_enabled), ): return {modular_file: mmc.convert_modular_file(modular_file) for modular_file in modular_files} @@ -102,16 +89,13 @@ def benchmark_conversion( modular_files: list[str], wrapper_cls, cache_enabled: bool, - fast_import_analysis_enabled: bool, fast_mapper_visit_enabled: bool, repeat: int, ) -> list[float]: durations = [] for _ in range(repeat): start = time.perf_counter() - run_conversion( - modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled, fast_mapper_visit_enabled - ) + run_conversion(modular_files, wrapper_cls, cache_enabled, fast_mapper_visit_enabled) durations.append(time.perf_counter() - start) return durations @@ -120,13 +104,12 @@ def profile_conversion( modular_files: list[str], wrapper_cls, cache_enabled: bool, - fast_import_analysis_enabled: bool, fast_mapper_visit_enabled: bool, output_path: Path, ) -> None: profile = cProfile.Profile() profile.enable() - run_conversion(modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled, fast_mapper_visit_enabled) + run_conversion(modular_files, wrapper_cls, cache_enabled, fast_mapper_visit_enabled) profile.disable() profile.dump_stats(str(output_path)) @@ -220,11 +203,6 @@ def main() -> int: action="store_true", help="Compare against current wrapper behavior with module source caching disabled.", ) - parser.add_argument( - "--compare-old-import-analysis", - action="store_true", - help="Compare against the current branch with the previous get_needed_imports analysis.", - ) parser.add_argument( "--compare-old-mapper-visit", action="store_true", @@ -250,43 +228,38 @@ def main() -> int: output_root = make_output_root(args.output_dir, modular_files) output_root.mkdir(parents=True, exist_ok=True) - modes = [("current", LibCSTMetadataWrapper, True, True, True)] + modes = [("current", LibCSTMetadataWrapper, True, True)] if args.compare_old_branch: - modes.append(("old_branch", OldBehaviorMetadataWrapper, False, False, False)) + modes.append(("old_branch", OldBehaviorMetadataWrapper, False, False)) if args.compare_old_wrapper: - modes.append(("old_wrapper", OldBehaviorMetadataWrapper, True, True, True)) + modes.append(("old_wrapper", OldBehaviorMetadataWrapper, True, True)) if args.compare_no_cache: - modes.append(("no_cache", LibCSTMetadataWrapper, False, True, True)) - if args.compare_old_import_analysis: - modes.append(("old_import_analysis", LibCSTMetadataWrapper, True, False, True)) + modes.append(("no_cache", LibCSTMetadataWrapper, False, True)) if args.compare_old_mapper_visit: - modes.append(("old_mapper_visit", LibCSTMetadataWrapper, True, True, False)) + modes.append(("old_mapper_visit", LibCSTMetadataWrapper, True, False)) print(f"files={len(modular_files)}") - for label, wrapper_cls, cache_enabled, fast_import_analysis_enabled, fast_mapper_visit_enabled in modes: + for label, wrapper_cls, cache_enabled, fast_mapper_visit_enabled in modes: durations = benchmark_conversion( modular_files, wrapper_cls, cache_enabled, - fast_import_analysis_enabled, fast_mapper_visit_enabled, args.repeat, ) print(summarize(label, durations)) if args.verify_equal and len(modes) > 1: - current = run_conversion(modular_files, LibCSTMetadataWrapper, True, True, True) - for label, wrapper_cls, cache_enabled, fast_import_analysis_enabled, fast_mapper_visit_enabled in modes[1:]: - other = run_conversion( - modular_files, wrapper_cls, cache_enabled, fast_import_analysis_enabled, fast_mapper_visit_enabled - ) + current = run_conversion(modular_files, LibCSTMetadataWrapper, True, True) + for label, wrapper_cls, cache_enabled, fast_mapper_visit_enabled in modes[1:]: + other = run_conversion(modular_files, wrapper_cls, cache_enabled, fast_mapper_visit_enabled) same = current.keys() == other.keys() and all(current[key] == other[key] for key in current) print(f"{label}_byte_identical={same}") if not same: return 1 profile_path = output_root / "current.prof" - profile_conversion(modular_files, LibCSTMetadataWrapper, True, True, True, profile_path) + profile_conversion(modular_files, LibCSTMetadataWrapper, True, True, profile_path) print(f"profile={profile_path}") if args.snakeviz: From adf54500ee9175719a6b5c25bfd13bd2845baf69 Mon Sep 17 00:00:00 2001 From: Tarek Ziade Date: Fri, 24 Apr 2026 09:12:16 +0200 Subject: [PATCH 6/7] removed the profiler and cleanup the converter touse both changes --- utils/modular_model_converter.py | 61 ++----- utils/profile_modular_conversion.py | 274 ---------------------------- 2 files changed, 11 insertions(+), 324 deletions(-) delete mode 100644 utils/profile_modular_conversion.py diff --git a/utils/modular_model_converter.py b/utils/modular_model_converter.py index d0b42771039f..9afc59d018aa 100644 --- a/utils/modular_model_converter.py +++ b/utils/modular_model_converter.py @@ -26,7 +26,7 @@ from create_dependency_mapping import find_priority_list from libcst import ClassDef, CSTVisitor from libcst import matchers as m -from libcst.metadata import MetadataWrapper, ParentNodeProvider, PositionProvider, ScopeProvider +from libcst.metadata import MetadataWrapper, ScopeProvider from modular_integrations import ( EXCLUDED_EXTERNAL_FILES, AbsoluteImportTransformer, @@ -49,15 +49,9 @@ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 """ -ENABLE_MODULE_SOURCE_CACHE = True -ENABLE_FAST_MAPPER_VISIT = True _MODULE_SOURCE_CACHE = {} -def clear_module_source_cache(): - _MODULE_SOURCE_CACHE.clear() - - def get_module_source_and_tree_from_name(module_name: str) -> tuple[str, cst.Module]: spec = importlib.util.find_spec(module_name) if spec is None or spec.origin is None: @@ -67,14 +61,13 @@ def get_module_source_and_tree_from_name(module_name: str) -> tuple[str, cst.Mod cache_key = (module_name, file_path) mtime_ns = os.stat(file_path).st_mtime_ns cached = _MODULE_SOURCE_CACHE.get(cache_key) - if ENABLE_MODULE_SOURCE_CACHE and cached is not None and cached[0] == mtime_ns: + if cached is not None and cached[0] == mtime_ns: return cached[1], cached[2] with open(file_path, "r", encoding="utf-8") as file: source_code = file.read() tree = cst.parse_module(source_code) - if ENABLE_MODULE_SOURCE_CACHE: - _MODULE_SOURCE_CACHE[cache_key] = (mtime_ns, source_code, tree) + _MODULE_SOURCE_CACHE[cache_key] = (mtime_ns, source_code, tree) return source_code, tree @@ -555,8 +548,6 @@ class ModuleMapper(CSTVisitor, ABC): modeling files that will be visited. """ - METADATA_DEPENDENCIES = (ParentNodeProvider, PositionProvider) - def __init__(self, python_module: cst.Module): # fmt: off self.python_module: cst.Module = python_module # original cst.Module being visited @@ -624,11 +615,7 @@ def visit_SimpleStatementLine(self, node): body=[m.Assign(targets=[m.AssignTarget(target=m.Subscript(value=m.Name()) | m.Attribute(value=m.Name()))])] ) - is_module_level = ( - self._is_direct_module_child() - if ENABLE_FAST_MAPPER_VISIT - else m.matches(self.get_metadata(cst.metadata.ParentNodeProvider, node), m.Module()) - ) + is_module_level = self._is_direct_module_child() if is_module_level: if m.matches(node, simple_top_level_assign_structure): left_hand_side = node.body[0].targets[0].target.value @@ -654,22 +641,14 @@ def leave_SimpleStatementLine(self, node): self.current_assignment = None def visit_FunctionDef(self, node): - is_module_level = ( - self._is_direct_module_child() - if ENABLE_FAST_MAPPER_VISIT - else m.matches(self.get_metadata(cst.metadata.ParentNodeProvider, node), m.Module()) - ) + is_module_level = self._is_direct_module_child() if is_module_level: self.current_function = node.name.value self.functions[node.name.value] = node self._record_node_order(node.name.value) def leave_FunctionDef(self, node): - is_module_level = ( - self._is_direct_module_child() - if ENABLE_FAST_MAPPER_VISIT - else m.matches(self.get_metadata(cst.metadata.ParentNodeProvider, node), m.Module()) - ) + is_module_level = self._is_direct_module_child() if is_module_level: self.current_function = None @@ -698,20 +677,14 @@ def visit_Name(self, node: cst.Call): def leave_Module(self, node): """When leaving the module, we store the position of each global scoped node to allow sorting the dependencies - based on their position in the code later. We use the PositionProvider metadata wrapper for this. + based on their position in the code later. We also make sure to update `self.object_dependency_mapping` so that it contains only names recorded in `self.global_nodes`. """ # assign all nodes self.global_nodes = {**self.assignments, **self.classes, **self.functions} # now sort the class dependency_mapping based on the position of the nodes - self.start_lines = {} - if ENABLE_FAST_MAPPER_VISIT: - for node_name in self.global_nodes: - self.start_lines[node_name] = self._node_order[node_name] - else: - for node_name, current_node in self.global_nodes.items(): - self.start_lines[node_name] = self.get_metadata(cst.metadata.PositionProvider, current_node).start.line + self.start_lines = {node_name: self._node_order[node_name] for node_name in self.global_nodes} def _restrict_dependencies_to_known_entities(self): """Since we added every Name as part of `self.object_dependency_mapping`, we need to remove those that @@ -920,11 +893,7 @@ def visit_and_merge_dependencies( cls, module: cst.Module, classes, functions, assignments, object_mapping, start_lines ) -> "ModelFileMapper": mapper = cls(module) - if ENABLE_FAST_MAPPER_VISIT: - module.visit(mapper) - else: - wrapper = MetadataWrapper(module, unsafe_skip_copy=True) - wrapper.visit(mapper) + module.visit(mapper) # Merge dependencies mapper.merge_modular_dependencies(classes, functions, assignments, object_mapping, start_lines) # Create the class dependencies graph @@ -1508,11 +1477,7 @@ def visit_SimpleStatementLine(self, node): body=[m.Assign(targets=[m.AssignTarget(target=m.Subscript(value=m.Name()) | m.Attribute(value=m.Name()))])] ) - is_module_level = ( - self._is_direct_module_child() - if ENABLE_FAST_MAPPER_VISIT - else m.matches(self.get_metadata(cst.metadata.ParentNodeProvider, node), m.Module()) - ) + is_module_level = self._is_direct_module_child() if is_module_level: if m.matches(node, m.SimpleStatementLine(body=[m.Import()])): self.imports.append(node) @@ -2023,11 +1988,7 @@ def convert_modular_file(modular_file: str, source_library: str | None = "transf module = module.visit(AbsoluteImportTransformer(relative_path, source_library)) cst_transformers = ModularFileMapper(module, model_name, source_library) - if ENABLE_FAST_MAPPER_VISIT: - module.visit(cst_transformers) - else: - wrapper = MetadataWrapper(module, unsafe_skip_copy=True) - wrapper.visit(cst_transformers) + module.visit(cst_transformers) for file, module in create_modules( cst_transformers, file_path=relative_path, package_name=source_library ).items(): diff --git a/utils/profile_modular_conversion.py b/utils/profile_modular_conversion.py deleted file mode 100644 index 8ab2297e259f..000000000000 --- a/utils/profile_modular_conversion.py +++ /dev/null @@ -1,274 +0,0 @@ -#!/usr/bin/env python - -""" -Benchmark and profile modular conversion. - -Examples: - python utils/profile_modular_conversion.py \ - src/transformers/models/conditional_detr/modular_conditional_detr.py - - python utils/profile_modular_conversion.py \ - src/transformers/models/conditional_detr/modular_conditional_detr.py \ - src/transformers/models/deepseek_vl/modular_deepseek_vl.py \ - --repeat 3 \ - --compare-old-mapper-visit \ - --compare-old-branch \ - --verify-equal \ - --snakeviz -""" - -import argparse -import cProfile -import json -import re -import shutil -import statistics -import time -from contextlib import contextmanager -from pathlib import Path -from pstats import Stats - -import modular_model_converter as mmc -from libcst.metadata import MetadataWrapper as LibCSTMetadataWrapper - - -class OldBehaviorMetadataWrapper(LibCSTMetadataWrapper): - def __init__(self, module, unsafe_skip_copy=False, cache=None): - if cache is None: - cache = {} - super().__init__(module, unsafe_skip_copy=False, cache=cache) - - -@contextmanager -def metadata_wrapper_context(wrapper_cls): - original = mmc.MetadataWrapper - mmc.MetadataWrapper = wrapper_cls - try: - yield - finally: - mmc.MetadataWrapper = original - - -@contextmanager -def module_cache_context(enabled: bool): - original = mmc.ENABLE_MODULE_SOURCE_CACHE - mmc.ENABLE_MODULE_SOURCE_CACHE = enabled - mmc.clear_module_source_cache() - try: - yield - finally: - mmc.clear_module_source_cache() - mmc.ENABLE_MODULE_SOURCE_CACHE = original - - -@contextmanager -def fast_mapper_visit_context(enabled: bool): - original = mmc.ENABLE_FAST_MAPPER_VISIT - mmc.ENABLE_FAST_MAPPER_VISIT = enabled - try: - yield - finally: - mmc.ENABLE_FAST_MAPPER_VISIT = original - - -def run_conversion( - modular_files: list[str], - wrapper_cls, - cache_enabled: bool, - fast_mapper_visit_enabled: bool, -) -> dict[str, dict[str, str]]: - with ( - metadata_wrapper_context(wrapper_cls), - module_cache_context(cache_enabled), - fast_mapper_visit_context(fast_mapper_visit_enabled), - ): - return {modular_file: mmc.convert_modular_file(modular_file) for modular_file in modular_files} - - -def benchmark_conversion( - modular_files: list[str], - wrapper_cls, - cache_enabled: bool, - fast_mapper_visit_enabled: bool, - repeat: int, -) -> list[float]: - durations = [] - for _ in range(repeat): - start = time.perf_counter() - run_conversion(modular_files, wrapper_cls, cache_enabled, fast_mapper_visit_enabled) - durations.append(time.perf_counter() - start) - return durations - - -def profile_conversion( - modular_files: list[str], - wrapper_cls, - cache_enabled: bool, - fast_mapper_visit_enabled: bool, - output_path: Path, -) -> None: - profile = cProfile.Profile() - profile.enable() - run_conversion(modular_files, wrapper_cls, cache_enabled, fast_mapper_visit_enabled) - profile.disable() - profile.dump_stats(str(output_path)) - - -def write_snakeviz_bundle(profile_path: Path, bundle_dir: Path, profile_name: str) -> Path: - try: - import snakeviz - from snakeviz.stats import json_stats, table_rows - from tornado.template import Template - except ImportError as error: - raise RuntimeError("SnakeViz bundle generation requires snakeviz and tornado to be installed.") from error - - if bundle_dir.exists(): - shutil.rmtree(bundle_dir) - bundle_dir.mkdir(parents=True) - - static_dir = bundle_dir / "static" - snakeviz_root = Path(snakeviz.__file__).resolve().parent - shutil.copytree(snakeviz_root / "static", static_dir) - - stats = Stats(str(profile_path)) - template_text = (snakeviz_root / "templates" / "viz.html").read_text(encoding="utf-8") - html = ( - Template(template_text) - .generate( - profile_name=profile_name, - table_rows=json.dumps(table_rows(stats)), - callees=json.dumps(json_stats(stats)), - ) - .decode("utf-8") - ) - - static_uri = static_dir.resolve().as_uri() - html = html.replace('href="/static/', f'href="{static_uri}/') - html = html.replace("href='/static/", f"href='{static_uri}/") - html = html.replace('src="/static/', f'src="{static_uri}/') - html = html.replace("src='/static/", f"src='{static_uri}/") - html = html.replace( - "event.data['url'] + \"/static/vendor/immutable.min.js\",", - f'"{(static_dir / "vendor" / "immutable.min.js").resolve().as_uri()}",', - ) - html = html.replace( - "event.data['url'] + \"/static/vendor/lodash.min.js\");", - f'"{(static_dir / "vendor" / "lodash.min.js").resolve().as_uri()}");', - ) - html = html.replace("'url': window.location.origin", "'url': ''") - - html_path = bundle_dir / "index.html" - html_path.write_text(html, encoding="utf-8") - return html_path - - -def make_output_root(output_dir: str | None, modular_files: list[str]) -> Path: - if output_dir is not None: - return Path(output_dir) - - if len(modular_files) == 1: - match = re.search(r"modular_(.*)(?=\.py$)", modular_files[0]) - name = match.group(1) if match is not None else Path(modular_files[0]).stem - else: - name = f"batch_{len(modular_files)}_modular_files" - return Path("/tmp/codex-profiles") / f"{name}_modular_profile" - - -def summarize(label: str, durations: list[float]) -> str: - mean = statistics.mean(durations) - median = statistics.median(durations) - minimum = min(durations) - maximum = max(durations) - return ( - f"{label}: mean={mean:.3f}s median={median:.3f}s min={minimum:.3f}s max={maximum:.3f}s runs={len(durations)}" - ) - - -def main() -> int: - parser = argparse.ArgumentParser(description="Benchmark and profile modular conversion.") - parser.add_argument("modular_files", nargs="+", help="One or more paths to modular_*.py files.") - parser.add_argument("--repeat", type=int, default=3, help="Number of timing runs per mode.") - parser.add_argument( - "--compare-old-branch", - action="store_true", - help="Compare against simulated pre-change branch behavior (old wrapper + no module cache).", - ) - parser.add_argument( - "--compare-old-wrapper", - action="store_true", - help="Compare against only the pre-change MetadataWrapper behavior.", - ) - parser.add_argument( - "--compare-no-cache", - action="store_true", - help="Compare against current wrapper behavior with module source caching disabled.", - ) - parser.add_argument( - "--compare-old-mapper-visit", - action="store_true", - help="Compare against the current branch with the previous metadata-based module visit path.", - ) - parser.add_argument( - "--verify-equal", - action="store_true", - help="Verify that current output and all requested comparison modes are byte-identical.", - ) - parser.add_argument( - "--snakeviz", - action="store_true", - help="Generate a standalone SnakeViz HTML bundle for the current mode.", - ) - parser.add_argument( - "--output-dir", - help="Directory for profile outputs. Defaults to /tmp/codex-profiles/_modular_profile.", - ) - args = parser.parse_args() - - modular_files = args.modular_files - output_root = make_output_root(args.output_dir, modular_files) - output_root.mkdir(parents=True, exist_ok=True) - - modes = [("current", LibCSTMetadataWrapper, True, True)] - if args.compare_old_branch: - modes.append(("old_branch", OldBehaviorMetadataWrapper, False, False)) - if args.compare_old_wrapper: - modes.append(("old_wrapper", OldBehaviorMetadataWrapper, True, True)) - if args.compare_no_cache: - modes.append(("no_cache", LibCSTMetadataWrapper, False, True)) - if args.compare_old_mapper_visit: - modes.append(("old_mapper_visit", LibCSTMetadataWrapper, True, False)) - - print(f"files={len(modular_files)}") - for label, wrapper_cls, cache_enabled, fast_mapper_visit_enabled in modes: - durations = benchmark_conversion( - modular_files, - wrapper_cls, - cache_enabled, - fast_mapper_visit_enabled, - args.repeat, - ) - print(summarize(label, durations)) - - if args.verify_equal and len(modes) > 1: - current = run_conversion(modular_files, LibCSTMetadataWrapper, True, True) - for label, wrapper_cls, cache_enabled, fast_mapper_visit_enabled in modes[1:]: - other = run_conversion(modular_files, wrapper_cls, cache_enabled, fast_mapper_visit_enabled) - same = current.keys() == other.keys() and all(current[key] == other[key] for key in current) - print(f"{label}_byte_identical={same}") - if not same: - return 1 - - profile_path = output_root / "current.prof" - profile_conversion(modular_files, LibCSTMetadataWrapper, True, True, profile_path) - print(f"profile={profile_path}") - - if args.snakeviz: - profile_name = modular_files[0] if len(modular_files) == 1 else f"batch:{len(modular_files)} modular files" - html_path = write_snakeviz_bundle(profile_path, output_root / "snakeviz", profile_name) - print(f"snakeviz={html_path}") - - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) From 82ba8df8783f774ef317e2fb1166a665625647da Mon Sep 17 00:00:00 2001 From: Tarek Ziade Date: Mon, 27 Apr 2026 08:32:57 +0200 Subject: [PATCH 7/7] make the change closer to main --- utils/modular_model_converter.py | 48 +++++++++++++++----------------- 1 file changed, 23 insertions(+), 25 deletions(-) diff --git a/utils/modular_model_converter.py b/utils/modular_model_converter.py index 9afc59d018aa..a6652d4acc3a 100644 --- a/utils/modular_model_converter.py +++ b/utils/modular_model_converter.py @@ -1189,11 +1189,30 @@ def append_new_import_node( imports_to_keep.append(new_node) -def _build_needed_imports(all_imports: list[cst.CSTNode], unused_imports: set[str]) -> list[cst.CSTNode]: - """Build the final import block after pruning the unused names.""" +def get_needed_imports(body: dict[str, dict], all_imports: list[cst.CSTNode]) -> list[cst.CSTNode]: + """Get all the imports needed in the `body`, from the list of `all_imports`. + `body` is a dict with the following structure `{str: {"insert_idx": int, "node": cst.CSTNode}}`. + Note: we need to use `isinstance` on scope assignments, m.matches apparently does not work here yet! + """ + new_body = [k[1]["node"] for k in sorted(body.items(), key=lambda x: x[1]["insert_idx"])] + wrapper = MetadataWrapper(cst.Module(body=all_imports + new_body), unsafe_skip_copy=True) + scopes = set(wrapper.resolve(ScopeProvider).values()) + unused_imports = set() + import_ref_count = defaultdict(lambda: 0) + for scope in scopes: + for assignment in scope.assignments: + node = assignment.node + if isinstance(assignment, cst.metadata.Assignment) and isinstance(node, (cst.Import, cst.ImportFrom)): + ref_count = len(assignment.references) + name = assignment.name + import_ref_count[name] = max(ref_count, import_ref_count[name]) + # Similar imports may be redefined, and only used between their 1st and 2nd definition so if we already have + # a ref count > 0 at any point, the imports is actually used + unused_imports = {name for name, count in import_ref_count.items() if count <= 0 or name in body} + imports_to_keep = [] - # We need to keep track of which names were already imported, because some import may be duplicated from multiple - # sources or be both protected and unprotected due to inconsistency between models. + # We need to keep track of which names were already imported, because some import may be duplicated from multiple sources + # or be both protected and unprotected due to inconsistency between models added_names = set() existing_protected_statements = set() # str repr of the import nodes - does not work with the nodes directly for node in all_imports: @@ -1216,27 +1235,6 @@ def _build_needed_imports(all_imports: list[cst.CSTNode], unused_imports: set[st return usual_import_nodes + protected_import_nodes -def get_needed_imports(body: dict[str, dict], all_imports: list[cst.CSTNode]) -> list[cst.CSTNode]: - """Get all the imports needed in the `body`, from the list of `all_imports`. - `body` is a dict with the following structure `{str: {"insert_idx": int, "node": cst.CSTNode}}`. - """ - new_body = [k[1]["node"] for k in sorted(body.items(), key=lambda x: x[1]["insert_idx"])] - wrapper = MetadataWrapper(cst.Module(body=all_imports + new_body), unsafe_skip_copy=True) - scopes = set(wrapper.resolve(ScopeProvider).values()) - import_ref_count = defaultdict(lambda: 0) - for scope in scopes: - for assignment in scope.assignments: - node = assignment.node - if isinstance(assignment, cst.metadata.Assignment) and isinstance(node, (cst.Import, cst.ImportFrom)): - ref_count = len(assignment.references) - name = assignment.name - import_ref_count[name] = max(ref_count, import_ref_count[name]) - # Similar imports may be redefined, and only used between their 1st and 2nd definition so if we already have - # a ref count > 0 at any point, the imports is actually used - unused_imports = {name for name, count in import_ref_count.items() if count <= 0 or name in body} - return _build_needed_imports(all_imports, unused_imports) - - def _ensure_utils_availability_imports(imports: list[cst.CSTNode], needed: set[str]) -> list[cst.CSTNode]: """Add is_torch_available and/or is_torchvision_available to the utils import if needed.""" if not needed: