diff --git a/.ci/scripts/tests/test_wheel_test_modules.py b/.ci/scripts/tests/test_wheel_test_modules.py new file mode 100644 index 00000000000..39c1d3f250f --- /dev/null +++ b/.ci/scripts/tests/test_wheel_test_modules.py @@ -0,0 +1,867 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Unit tests for the test modules the full wheel drops. + +The wheel used to carry every test file in the repository, about 9.7 MB of Python that nothing +in an installed wheel can reach. A test case is only ever loaded by pytest from a path in the +checkout, never through the installed name, so shipping it buys nothing. + +Shared helpers are the opposite. The suites here import each other by installed name, for +example `from executorch.backends.arm.test.tester.test_pipeline import TosaPipelineFP`, so a +helper has to ship or collection breaks. That is why the keep set is computed from the import +graph and not from file names: `test_pipeline.py` and `test_add.py` are indistinguishable by +name and only one of them can go. + +setup.py is read rather than imported. It calls setup() at module scope, so importing it under a +test runner hands setup() the runner's own arguments and the session dies on an invalid command +name. +""" + +import ast +import functools +import os +import re +import shutil +import subprocess +import tempfile +import unittest +from pathlib import Path +from typing import Dict, FrozenSet, List, Optional, Set, Tuple + +from setuptools import find_namespace_packages + +SETUP_PY = Path(__file__).resolve().parents[3] / "setup.py" +REPO_ROOT = SETUP_PY.parent + +# Enough of setup.py to exercise the keep set, and nothing that builds anything. +_WANTED = ( + "_WALK_SKIP_DIRS", + "_TEST_DIR_NAMES", + "_CI_ENTRY_POINTS", + "_CI_ENTRY_POINT_DIRS", + "_SHADER_TEMPLATE_MARKERS", + "_is_shader_template", + "_VENDORED_DIR_NAMES", + "_VENDORED_SUBMODULE_FALLBACK", + "_top_level_package_dirs", + "_first_party_module", + "_is_test_module", + "_module_name", + "_import_targets", + "_scan_imports", + "_GENERATED_DIR_NAMES", + "_unshipped_directories", + "_import_graph", + "_reachable_test_modules", + "_vendored_prefixes", + "_is_vendored_path", + "_full_packages", + "_minimal_packages", +) + + +def _setup_py_module() -> ast.Module: + # Name the encoding: these tests are collected on Windows too (pytest-windows.ini line 19), + # where the default is cp1252, and setup.py holds a non-ascii apostrophe that would decode to + # the wrong characters without a word rather than raising. + return ast.parse(SETUP_PY.read_text(encoding="utf-8")) + + +def _load_from_setup_py() -> Dict[str, object]: + """Run only the named definitions from setup.py, not its build logic.""" + selected: List[ast.stmt] = [] + found: Set[str] = set() + for node in _setup_py_module().body: + if isinstance(node, ast.FunctionDef) and node.name in _WANTED: + selected.append(node) + found.add(node.name) + elif isinstance(node, ast.Assign): + names = { + target.id + for target in node.targets + if isinstance(target, ast.Name) and target.id in _WANTED + } + if names: + selected.append(node) + found |= names + + assert found == set( + _WANTED + ), f"setup.py no longer defines {sorted(set(_WANTED) - found)}, so this test checks nothing" + + namespace: Dict[str, object] = { + "__file__": str(SETUP_PY), + "ast": ast, + "os": os, + "Path": Path, + "functools": functools, + "subprocess": subprocess, + "Dict": Dict, + "FrozenSet": FrozenSet, + "List": List, + "Optional": Optional, + "Set": Set, + "Tuple": Tuple, + "find_namespace_packages": find_namespace_packages, + } + exec( + compile(ast.Module(body=selected, type_ignores=[]), str(SETUP_PY), "exec"), + namespace, + ) + return namespace + + +_NAMESPACE = _load_from_setup_py() +_is_test_module = _NAMESPACE["_is_test_module"] +_import_graph = _NAMESPACE["_import_graph"] +_reachable_test_modules = _NAMESPACE["_reachable_test_modules"] +_CI_ENTRY_POINTS = _NAMESPACE["_CI_ENTRY_POINTS"] + + +@functools.lru_cache(maxsize=None) +def _graph() -> Tuple[Set[str], Dict[str, Set[str]], Set[str]]: + return _import_graph(REPO_ROOT / "src" / "executorch") + + +def _a_dropped_test_module() -> str: + """A real test module the keep set excludes, so the wiring tests assert on real data. + + Must be a leaf inside a test package, because find_package_modules is only given a chance to + drop something when the package it is asked about is itself under a test directory. + """ + modules, _edges, _dynamic = _graph() + keep = _reachable_test_modules() + dropped = sorted( + name + for name in modules + if name not in keep + and _is_test_module(name.rsplit(".", 1)[0]) + and name.rsplit(".", 1)[1] != "__init__" + ) + assert dropped, "nothing is dropped, so the wiring tests would be vacuous" + return dropped[0] + + +_UNREACHABLE_TEST_MODULE = _a_dropped_test_module() + + +@functools.lru_cache(maxsize=None) +def _reachable_from_imports_only() -> FrozenSet[str]: + """The keep set the import graph produces on its own, with no directory entries applied. + + Used to tell a load-bearing directory entry from a redundant one: a module the graph already + reaches would ship whether or not its directory is listed. + """ + modules, edges, dynamic = _graph() + referenced = set(dynamic) + referenced.update(_CI_ENTRY_POINTS) + for targets in edges.values(): + referenced.update(targets) + return frozenset(name for name in referenced if _is_test_module(name)) & modules + + +@functools.lru_cache(maxsize=None) +def _entry_point_dir_drivers() -> Dict[str, str]: + """Why each `_CI_ENTRY_POINT_DIRS` entry exists, as the file that drives it. + + These directories cannot be re-derived from the source, which is the whole reason they are + listed by hand: each is walked by something that never spells out a module name, so there is + no import to find and no literal to grep for. What CAN be checked is that the thing doing the + walking still exists and still refers to the directory. If a driver is deleted or stops + mentioning its directory, the entry has outlived its reason and this pairing fails. + + Keyed by the dotted prefix, valued by a repository-relative path. + """ + return { + "executorch.backends.mlx.custom_kernel_ops": ".github/workflows/mlx.yml", + "executorch.backends.webgpu.test": "backends/webgpu/scripts/test_webgpu_native_ci.sh", + "executorch.backends.test.suite": "backends/test/suite/runner.py", + "executorch.examples.models.llava.test": "examples/models/llava/README.md", + } + + +@functools.lru_cache(maxsize=None) +def _tracked_shell_scripts() -> Tuple[str, ...]: + """Shell scripts this repository actually owns, as repository-relative paths. + + Asked of git rather than found by walking. CI checks other repositories out INSIDE this one, + for example a `pytorch/` sibling clone, and a walk cannot tell those files from ours. It found + `pytorch/.ci/pytorch/test.sh` and reported a module belonging to a different project, so the + walk failed on CI while passing in every local checkout. + + An archive with no git available yields nothing, which makes this test vacuous rather than + wrong. It is a drift guard, so silence in an environment that cannot check is the safe way to + fail. + """ + try: + listed = subprocess.run( + ["git", "-C", str(REPO_ROOT), "ls-files", "-z", "*.sh"], + capture_output=True, + text=True, + check=False, + ) + except OSError: + # No git on PATH, as in an unpacked source archive. + return () + if listed.returncode: + return () + return tuple(name for name in listed.stdout.split("\0") if name) + + +def _fake_prune(build_lib, source_root): + """A CustomBuildPy whose prune runs, with build_lib and the source tree given directly. + + The real method resolves the source tree from setup.py's own location, so the lifted body is + bound to a stand-in whose __file__ points at the fixture instead. + """ + classes = [ + node + for node in _setup_py_module().body + if isinstance(node, ast.ClassDef) and node.name == "CustomBuildPy" + ] + assert len(classes) == 1, "setup.py no longer defines CustomBuildPy" + bodies = [ + node + for node in classes[0].body + if isinstance(node, ast.FunctionDef) and node.name == "_prune_unstaged_files" + ] + assert len(bodies) == 1, "the stale-file prune is gone" + + namespace = { + "os": os, + "Path": Path, + "__file__": str(source_root.parent / "setup.py"), + } + exec(compile(ast.unparse(bodies[0]), "prune", "exec"), namespace) + + class Stub: + editable_mode = False + packages = ["executorch", "executorch.pkg"] + + def __init__(self): + self.build_lib = str(build_lib) + + def find_all_modules(self): + # Deliberately omits stale.py, which is what marks it unwanted. + return [("executorch", "__init__", ""), ("executorch.pkg", "__init__", "")] + + def get_package_dir(self, package): + return str(source_root / Path(*package.split("."))) + + def find_data_files(self, package, src_dir): + return [] + + Stub._prune_unstaged_files = namespace["_prune_unstaged_files"] + return Stub() + + +def _fake_build_py(): + """A CustomBuildPy whose overrides run, without configuring a real distribution. + + The overrides are lifted from setup.py and bound to a stand-in so they can be CALLED. The + point is to exercise the real bodies: a test that only reads their syntax passes on code that + never runs, which is the hole this helper exists to close. + """ + classes = [ + node + for node in _setup_py_module().body + if isinstance(node, ast.ClassDef) and node.name == "CustomBuildPy" + ] + assert len(classes) == 1, "setup.py no longer defines CustomBuildPy" + wanted = ("find_package_modules", "find_data_files") + overrides = [ + node + for node in classes[0].body + if isinstance(node, ast.FunctionDef) and node.name in wanted + ] + assert {node.name for node in overrides} == set( + wanted + ), f"CustomBuildPy no longer overrides {sorted(set(wanted) - {n.name for n in overrides})}" + + package = _UNREACHABLE_TEST_MODULE.rsplit(".", 1)[0] + leaf = _UNREACHABLE_TEST_MODULE.rsplit(".", 1)[1] + modules = [(package, "__init__", "x"), (package, leaf, "y")] + + class Stub: + editable_mode = False + + def __init__(self) -> None: + self._data_files_to_return: List[str] = [] + + # Stands in for build_py's own implementations, which need a configured distribution. + def _super_find_package_modules(self, _package, _package_dir): + return list(modules) + + def _super_find_data_files(self, _package, _src_dir): + return list(self._data_files_to_return) + + namespace = dict(_NAMESPACE) + namespace["os"] = os + # `super()` needs a real base, so give the lifted bodies one that returns the fixtures above. + source = "\n".join( + ast.unparse(node) + .replace( + "super().find_package_modules(package, package_dir)", + "self._super_find_package_modules(package, package_dir)", + ) + .replace( + "super().find_data_files(package, src_dir)", + "self._super_find_data_files(package, src_dir)", + ) + for node in overrides + ) + exec(compile(source, "overrides", "exec"), namespace) + for name in wanted: + setattr(Stub, name, namespace[name]) + return Stub(), package, modules + + +class TestDroppedTestModules(unittest.TestCase): + def test_something_is_actually_dropped(self) -> None: + """The rule removes a substantial number of modules. + + Without this, every assertion below is vacuous on a keep set that happens to contain + everything, and the whole change could be reverted with the suite still green. + """ + modules, _, _ = _graph() + tests = {name for name in modules if _is_test_module(name)} + keep = _reachable_test_modules() + self.assertGreater(len(tests), 500, "no test modules discovered at all") + self.assertLess( + len(keep), + len(tests) // 2, + f"keeping {len(keep)} of {len(tests)} test modules, so almost nothing is dropped", + ) + + def test_shared_helpers_are_kept(self) -> None: + """Modules the suites import by installed name still ship. + + These are the ones whose removal breaks collection rather than a single test. Each is + imported from outside its own directory, which is what makes the installed name matter. + """ + keep = _reachable_test_modules() + for helper in ( + "executorch.backends.arm.test.tester.test_pipeline", + "executorch.backends.xnnpack.test.tester.tester", + "executorch.backends.test.harness.stages", + "executorch.backends.test.graph_builder", + "executorch.exir.backend.test.op_partitioner_demo", + ): + self.assertIn(helper, keep) + + def test_leaf_cases_are_dropped(self) -> None: + """A test case nothing imports does not ship. + + Chosen from different suites, because one backend getting this right says nothing about + the others. + """ + keep = _reachable_test_modules() + modules, _, _ = _graph() + for leaf in ( + "executorch.backends.arm.test.ops.test_add", + "executorch.backends.xnnpack.test.ops.test_bilinear2d", + ): + self.assertIn( + leaf, modules, f"{leaf} no longer exists, pick another example" + ) + self.assertNotIn(leaf, keep) + + def test_relative_imports_are_followed(self) -> None: + """A submodule reached only by a relative import is kept. + + backends/test/harness/stages/__init__.py does `from .export import Export`, so treating + a relative import as reaching nothing new drops stages.export and breaks every importer + of that package. This is a regression guard: it failed exactly that way once. + """ + self.assertIn( + "executorch.backends.test.harness.stages.export", _reachable_test_modules() + ) + + def test_dynamic_imports_are_followed(self) -> None: + """A module named only as a string to importlib is kept. + + backends/mlx/test/run_all_tests.py does + `importlib.import_module(".test_ops", package=__package__)`, which an import scan that + only reads import statements cannot see. Note test_ops is also named like a leaf, so a + file name rule would drop it. + """ + self.assertIn( + "executorch.backends.mlx.test.test_ops", _reachable_test_modules() + ) + + def test_ci_entry_points_are_kept(self) -> None: + """The modules only a workflow names are kept.""" + keep = _reachable_test_modules() + for name in _CI_ENTRY_POINTS: + self.assertIn(name, keep) + + def test_ci_entry_points_still_match_the_workflows(self) -> None: + """The hand-written CI list has not drifted from what the workflows actually run. + + The list is explicit rather than scanned at build time, because a source distribution + carries no .github directory and a scan there would silently keep nothing. The cost of + being explicit is drift, so it is checked here instead. + """ + pattern = re.compile(r"executorch(?:\.[A-Za-z0-9_]+)+") + referenced: Set[str] = set() + # The whole tree, not just .github and .ci. A workflow often calls a script that lives + # beside the backend it tests, and those name modules too: + # backends/webgpu/scripts/test_webgpu_native_ci.sh runs six of them by dotted name. + skip = { + ".git", + "pip-out", + "cmake-out", + "third-party", + "third_party", + "__pycache__", + } + for dirpath, dirnames, filenames in os.walk(REPO_ROOT, followlinks=False): + dirnames[:] = [d for d in dirnames if d not in skip] + for filename in filenames: + # Markdown too: a README documenting `python -m executorch.x.test.y` is a + # promise to users, and dropping that module breaks the documented command. + # Python as well, because several modules document their own `python -m` + # invocation in a docstring rather than in a README, and that is the same + # promise written somewhere else. + if not filename.endswith( + (".yml", ".yaml", ".sh", ".ps1", ".md", ".py") + ): + continue + path = Path(dirpath) / filename + if path.resolve() == Path(__file__).resolve(): + # This file names dropped modules as examples of what the rule removes, so + # reading itself would report them as promised and contradict its own tests. + continue + text = path.read_text(encoding="utf-8", errors="replace") + referenced.update(pattern.findall(text)) + + modules, _, _ = _graph() + + # A reference like executorch.a.test.b.SomeClass.some_method is one dotted run to the + # regex, and it is not a module, so trim each match back to its longest real module + # prefix. Without this the class-suffixed entries silently drop out of the comparison + # and the guard protects fewer names than it appears to. + def longest_module(name: str) -> str: + parts = name.split(".") + while parts: + candidate = ".".join(parts) + if candidate in modules: + return candidate + parts.pop() + return name + + expected = { + trimmed + for trimmed in (longest_module(name) for name in referenced) + if _is_test_module(trimmed) and trimmed in modules + } + missing = sorted(expected - set(_CI_ENTRY_POINTS) - _reachable_test_modules()) + self.assertEqual( + missing, + [], + f"a workflow names these test modules but nothing keeps them: {missing}", + ) + + def test_parent_packages_of_kept_modules_are_kept(self) -> None: + """Every kept module's package chain is kept, or the dotted path cannot resolve.""" + keep = _reachable_test_modules() + modules, _, _ = _graph() + for name in keep: + parts = name.split(".") + for end in range(2, len(parts)): + parent = ".".join(parts[:end]) + if _is_test_module(parent) and parent in modules: + self.assertIn(parent, keep, f"{parent} missing but {name} is kept") + + def test_shader_templates_do_not_ship(self) -> None: + """Shader codegen inputs are dropped, and op definitions are not. + + The cmake build expands these into SPIR-V and WGSL headers, so the wheel already carries + the compiled result. Matched on content, so the two examples below are the real + distinction: one is a template, the other is read at run time through + importlib.resources and must survive. + """ + is_template = _NAMESPACE["_is_shader_template"] + self.assertTrue( + is_template("backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml") + ) + self.assertTrue( + is_template("backends/webgpu/runtime/ops/binary_op/binary_op.yaml") + ) + for needed in ( + "exir/dialects/edge/edge.yaml", + "kernels/portable/functions.yaml", + "backends/cadence/aot/functions.yaml", + ): + self.assertFalse(is_template(needed), f"{needed} would stop shipping") + + def test_build_py_is_wired_to_the_custom_class(self) -> None: + """setup() receives CustomBuildPy, not the stock build_py. + + Every other test here exercises the class directly, so all of them stay green when the + cmdclass entry is pointed back at setuptools' own build_py. That single edit disables + the module filter, the data file filter and the prune at once, and the wheel then ships + everything again. + """ + assignments = [ + node + for node in ast.walk(_setup_py_module()) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "setup" + ] + self.assertEqual(len(assignments), 1, "expected exactly one setup() call") + + mapping = [kw.value for kw in assignments[0].keywords if kw.arg == "cmdclass"] + self.assertEqual(len(mapping), 1, "setup() no longer passes cmdclass") + wired = { + key.value: value.id + for key, value in zip(mapping[0].keys, mapping[0].values) + if isinstance(key, ast.Constant) and isinstance(value, ast.Name) + } + self.assertEqual( + wired.get("build_py"), + "CustomBuildPy", + "build_py is not wired to CustomBuildPy, so none of the filters run", + ) + + def test_both_package_lists_are_anchored_on_this_file(self) -> None: + """Neither package list depends on the working directory. + + A cwd-relative `where` returns nothing when the build runs from anywhere but the + repository root, and an empty package list makes the prune treat every staged file as + unwanted. The full list was anchored for this reason; the minimal one has to match. + + Both lists are CALLED from a directory that is not the repository root, because reading + the syntax of the `where=` argument only proves it is not a literal. Swapping the anchor + for `Path.cwd()` leaves the syntax test green and breaks every build started elsewhere. + """ + original = os.getcwd() + os.chdir(tempfile.gettempdir()) + try: + full = _NAMESPACE["_full_packages"]() + minimal = _NAMESPACE["_minimal_packages"]() + finally: + os.chdir(original) + self.assertIn("executorch", full) + self.assertGreater( + len(full), 100, "the full list collapsed when built from another directory" + ) + self.assertIn("executorch", minimal) + self.assertGreater( + len(minimal), + 1, + "the minimal list collapsed when built from another directory", + ) + + def test_stale_staged_files_are_pruned(self) -> None: + """A rebuild removes what an earlier build staged and this one does not want. + + build_py only copies, so without this a second build into the same directory keeps + every file the first one put there and the wheel packages it. The failure is silent: + the build succeeds and the wheel quietly contains the dropped files. + + The prune is CALLED against a real staging directory, because checking that the method + and its call site exist leaves an early `return` inside the body undetected, and the + prune then does nothing while this test stays green. + """ + staging = Path(tempfile.mkdtemp(prefix="prunetest-")) + self.addCleanup(shutil.rmtree, staging, ignore_errors=True) + source = staging / "src" + (source / "executorch" / "pkg").mkdir(parents=True) + for name in ("executorch/__init__.py", "executorch/pkg/__init__.py"): + (source / name).write_text("") + # Exists in the source tree and is NOT in build_py's file list, so the prune wants it + # gone. That is the whole contract. + (source / "executorch" / "pkg" / "stale.py").write_text( + "# left by an earlier build\n" + ) + build_lib = staging / "lib" + shutil.copytree(source, build_lib) + # Generated by a later build command, absent from src/, and must survive. + (build_lib / "executorch" / "pkg" / "generated.py").write_text( + "# from a template\n" + ) + + command = _fake_prune(build_lib, source) + command._prune_unstaged_files() + + remaining = sorted(p.name for p in (build_lib / "executorch" / "pkg").iterdir()) + self.assertNotIn( + "stale.py", + remaining, + "the prune left a file the current build does not want", + ) + self.assertIn( + "generated.py", + remaining, + "the prune deleted a file another command generated", + ) + self.assertIn("__init__.py", remaining, "the prune deleted a wanted module") + + def test_build_py_applies_the_keep_set(self) -> None: + """The drop is actually wired into the build, checked by CALLING the override. + + An earlier version of this test read the override's syntax tree instead. That passes on + code that is present but never runs, so an early `return modules` at the top of the + override left the filter dead with every assertion here still true. Build a real command + and look at what it returns. + """ + command, package, modules = _fake_build_py() + result = command.find_package_modules(package, "unused") + returned = {entry[1] for entry in result} + offered = {entry[1] for entry in modules} + self.assertIn("__init__", returned, "a kept package must still import") + self.assertTrue( + offered - returned, + "find_package_modules returned everything it was offered, so nothing is dropped", + ) + self.assertNotIn( + _UNREACHABLE_TEST_MODULE.rsplit(".", 1)[1], + returned, + "an unreachable test module was not dropped", + ) + + def test_shader_filter_is_wired_into_find_data_files(self) -> None: + """The shader classifier is actually CALLED, not merely correct. + + test_shader_templates_do_not_ship above checks the predicate. That is not the same thing: + deleting the filtering line in find_data_files leaves the predicate perfect and unused, + and every shader template ships again. + """ + command, _package, _modules = _fake_build_py() + template = "backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml" + needed = "kernels/portable/functions.yaml" + root = str(REPO_ROOT) + command._data_files_to_return = [ + os.path.join(root, template), + os.path.join(root, needed), + ] + kept = command.find_data_files("executorch", root) + self.assertNotIn( + os.path.join(root, template), + kept, + "find_data_files does not drop shader templates, so the filter is not wired in", + ) + self.assertIn( + os.path.join(root, needed), + kept, + "find_data_files dropped a yaml the runtime reads", + ) + + def test_ci_entry_point_dirs_are_all_load_bearing(self) -> None: + """Every listed directory still has a driver, and still keeps something. + + Nothing referenced `_CI_ENTRY_POINT_DIRS`, so an entry could be deleted with the whole + suite green: removing the backend suite line silently stopped 86 modules shipping. Two + checks close that. Each entry must be paired with the file that walks it, which fails when + an entry is added or removed without updating the pairing, and each entry must keep modules + the import graph cannot reach on its own, which fails when an entry becomes dead weight. + """ + listed = set(_NAMESPACE["_CI_ENTRY_POINT_DIRS"]) + self.assertTrue(listed, "the list is empty, so nothing is protected") + + drivers = _entry_point_dir_drivers() + self.assertEqual( + listed, + set(drivers), + "_CI_ENTRY_POINT_DIRS and its list of drivers disagree. Add the new entry with the " + "file that walks it, or drop the driver for the entry that went away", + ) + + for prefix, driver in sorted(drivers.items()): + path = REPO_ROOT / driver + self.assertTrue( + path.is_file(), + f"{prefix} is kept for {driver}, which no longer exists, so the entry may be " + "obsolete", + ) + tail = prefix.split(".")[-1] + self.assertIn( + tail, + path.read_text(encoding="utf-8", errors="replace"), + f"{driver} no longer mentions {tail}, so it may have stopped driving {prefix}", + ) + + # And the other direction: an entry that keeps nothing new is dead weight. + reached_anyway = _reachable_from_imports_only() + keep = _reachable_test_modules() + for entry in sorted(listed): + covered = { + name for name in keep if name == entry or name.startswith(f"{entry}.") + } + self.assertTrue( + covered - reached_anyway, + f"{entry} keeps nothing the import graph does not already reach, so the entry " + "is redundant and should be removed", + ) + + def test_ci_entry_points_cover_constructed_module_names(self) -> None: + """A runner that BUILDS a dotted name is covered too. + + The drift test above searches for a literal dotted name, so it cannot see a script that + assembles one, and a directory whose tests are only reached that way would be dropped + with nothing to warn about. + + A script that runs from a checkout by design is exempt, and says so in its own header. + `backends/apple/coreai/run_all_tests.sh` is the current example: it cds to the repository + root, so it always finds the files on disk and never needs them installed. + """ + pattern = re.compile(r"find\s+([A-Za-z0-9_./-]+)\s+-name\s+'?test_\*\.py'?") + keep = _reachable_test_modules() + listed = _NAMESPACE["_CI_ENTRY_POINT_DIRS"] + unprotected = [] + for relative in _tracked_shell_scripts(): + path = REPO_ROOT / relative + text = path.read_text(encoding="utf-8", errors="replace") + walked = pattern.findall(text) + if not walked: + continue + if "not a landing artifact" in text: + continue + for entry in walked: + dotted = "executorch." + entry.strip("./").replace("/", ".") + covered = any( + dotted == prefix or dotted.startswith(f"{prefix}.") + for prefix in listed + ) or any(name.startswith(f"{dotted}.") for name in keep) + if not covered: + unprotected.append(f"{relative} -> {dotted}") + self.assertEqual( + unprotected, + [], + "a script discovers test modules under these paths by building dotted names, and " + "nothing keeps them. Either add the directory to _CI_ENTRY_POINT_DIRS in setup.py, " + "or say in the script's header that it is not a landing artifact if it only ever " + f"runs from a checkout: {unprotected}", + ) + + def test_unprefixed_first_party_imports_count_as_references(self) -> None: + """`from backends.x import y` keeps y, the same as the prefixed spelling. + + This repository imports itself both ways: most code says `executorch.backends.x`, but the + Arm suites say `backends.arm.test...`, which resolves because the repository root is on + sys.path. Both name the same file. Following only the prefixed spelling dropped two shared + helpers with eight importers between them, which is the invariant this change exists to + preserve. + """ + first_party = _NAMESPACE["_first_party_module"] + self.assertEqual( + first_party("backends.arm.test.common"), + "executorch.backends.arm.test.common", + ) + self.assertEqual( + first_party("executorch.exir.tests.common"), "executorch.exir.tests.common" + ) + # A third-party module whose first component is not one of ours stays out. + self.assertIsNone(first_party("torch.nn.functional")) + self.assertIsNone(first_party("numpy")) + + keep = _reachable_test_modules() + for helper in ( + "executorch.backends.arm.test._custom_vgf_test_utils", + "executorch.backends.arm.test.runtime._vgf_runtime_test_utils", + ): + self.assertIn( + helper, + keep, + f"{helper} is imported without the executorch prefix and must still ship", + ) + + def test_importers_outside_the_shipped_tree_are_followed(self) -> None: + """A file the wheel does not carry can still import one that it does. + + `src/executorch` is a subset of the checkout, so an importer in a directory that is never + packaged is invisible to a walk of the shipped tree alone. Its imports still have to keep + their targets: test/end2end/test_end2end.py imports two model helpers out of exir/tests. + """ + keep = _reachable_test_modules() + importer = REPO_ROOT / "test" / "end2end" / "test_end2end.py" + self.assertTrue( + importer.is_file(), "this test needs a different example importer" + ) + for helper in ( + "executorch.exir.tests.dynamic_shape_models", + "executorch.exir.tests.transformer", + ): + self.assertIn( + helper, + keep, + f"{helper} is imported from outside the shipped tree and must still ship", + ) + + def test_vendored_trees_are_not_read_as_import_evidence(self) -> None: + """A vendored submodule's own imports do not keep anything. + + The package list excludes vendored trees, so nothing in one ships. The import scan has to + agree, or the two disagree about the same directory: a submodule checked out under an + ordinary name, rather than under `third-party`, was read as first-party and its imports + kept test modules the wheel never carries. + + Skipping by directory name alone is not enough, which is why this asserts on the scan's + output rather than on the skip list. + """ + modules, _edges, _dynamic = _graph() + is_vendored = _NAMESPACE["_is_vendored_path"] + vendored = sorted( + name for name in modules if is_vendored(name.replace(".", "/")) + ) + self.assertEqual( + vendored, + [], + "the import scan read these vendored modules as first-party, so their imports can " + f"keep test modules nothing shipped reaches: {vendored[:5]}", + ) + + def test_generated_directories_are_not_read_as_import_evidence(self) -> None: + """A build tree or an in-tree virtualenv does not vote on what ships. + + Those hold an INSTALLED copy of this package, so reading one lets the last wheel decide + what the next carries: a file that shipped once keeps itself alive. A clean checkout has + none of them, so the guard is exercised here by creating one. + """ + unshipped = _NAMESPACE["_unshipped_directories"] + root = REPO_ROOT / "src" / "executorch" + planted = REPO_ROOT / ".venv" + created = not planted.exists() + if created: + (planted / "lib").mkdir(parents=True) + self.addCleanup(shutil.rmtree, planted, ignore_errors=True) + walked = {entry.name for entry in unshipped(root)} + self.assertNotIn( + ".venv", + walked, + "a generated directory is read as import evidence, so an installed copy of this " + "package can keep test modules alive across builds", + ) + self.assertIn( + "test", walked, "the guard also dropped a real unshipped directory" + ) + + def test_editable_installs_are_left_alone(self) -> None: + """An editable install still exposes every test module. + + It maps the package root to a directory, so the suites resolve from the source tree + whatever is listed, and dropping modules there would only make the two install modes + disagree for no benefit. + """ + classes = [ + node + for node in _setup_py_module().body + if isinstance(node, ast.ClassDef) and node.name == "CustomBuildPy" + ] + overrides = [ + node + for node in classes[0].body + if isinstance(node, ast.FunctionDef) and node.name == "find_package_modules" + ] + source = ast.unparse(overrides[0]) + self.assertIn("editable_mode", source) + + +if __name__ == "__main__": + unittest.main() diff --git a/.ci/scripts/tests/test_wheel_vendored_packages.py b/.ci/scripts/tests/test_wheel_vendored_packages.py new file mode 100644 index 00000000000..a3d81a940b0 --- /dev/null +++ b/.ci/scripts/tests/test_wheel_vendored_packages.py @@ -0,0 +1,494 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Unit tests for the packages the full wheel publishes. + +The wheel used to carry the Python files and codegen scripts of every vendored third-party +checkout, because the full build passed no `packages` list and setuptools then discovered +everything under src/executorch. Those files exist to build the C++ targets, so nothing in +an installed wheel imports them. + +Asserting on the discovery result rather than on a built wheel, because the behaviour under +test is a pure function of the source tree plus the exclude patterns, and a full build takes +minutes to exercise one filter. `.ci/scripts/test_minimal_wheel.sh` already covers the +built-artifact side for the minimal wheel. + +setup.py is read rather than imported. It calls setup() at module scope, so importing it under +a test runner hands setup() the runner's own arguments and the session dies on an invalid +command name. +""" + +import ast +import functools +import os +import subprocess +import tempfile +import unittest +from pathlib import Path +from typing import Dict, FrozenSet, List, Set, Tuple + +from setuptools import find_namespace_packages +from setuptools.command.build_py import build_py + +SETUP_PY = Path(__file__).resolve().parents[3] / "setup.py" +REPO_ROOT = SETUP_PY.parent +# Discovery is anchored on this file's location, not on the working directory, so the result +# does not depend on where the runner was started. +PACKAGE_ROOT = str(SETUP_PY.parent / "src") + + +# The helpers this test drives, shared by both loaders below. +_HELPERS = ( + # _VENDORED_DIR_NAMES is not used directly below, but _is_vendored_path closes over it. + "_VENDORED_DIR_NAMES", + "_VENDORED_SUBMODULE_FALLBACK", + "_vendored_prefixes", + "_is_vendored_path", + # CustomBuildPy calls this, so the class cannot be exec'd without it. + "_SHADER_TEMPLATE_MARKERS", + "_is_shader_template", + "_full_packages", +) + + +def _setup_py_module() -> ast.Module: + # Name the encoding: these tests are collected on Windows too (pytest-windows.ini line 19), + # where the default is cp1252, and setup.py holds a non-ascii apostrophe that would decode to + # the wrong characters without a word rather than raising. + return ast.parse(SETUP_PY.read_text(encoding="utf-8")) + + +def _load_from_setup_py(root: Path = None) -> Dict[str, object]: + """The vendored-path helpers and the package list builder, from setup.py's source. + + Only those definitions are executed, so none of setup.py's module level build logic runs. + """ + wanted = _HELPERS + + selected: List[ast.stmt] = [] + found = set() + for node in _setup_py_module().body: + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in wanted: + selected.append(node) + found.add(node.name) + elif isinstance(node, ast.Assign): + names = { + target.id + for target in node.targets + if isinstance(target, ast.Name) and target.id in wanted + } + if names: + selected.append(node) + found |= names + + assert found == set( + wanted + ), f"setup.py no longer defines {sorted(set(wanted) - found)}, so this test checks nothing" + + namespace: Dict[str, object] = { + "__file__": str((root or SETUP_PY.parent) / "setup.py"), + "Path": Path, + "List": List, + "Tuple": Tuple, + "functools": functools, + "subprocess": subprocess, + "find_namespace_packages": find_namespace_packages, + } + exec( + compile(ast.Module(body=selected, type_ignores=[]), str(SETUP_PY), "exec"), + namespace, + ) + return namespace + + +@functools.lru_cache(maxsize=None) +def _load_build_py() -> Dict[str, object]: + """CustomBuildPy plus the helpers it calls, so analyze_manifest can be driven directly. + + Only the class body and those helpers run. Its methods reference names from setup.py's own + imports, so the ones analyze_manifest touches are supplied here. + """ + wanted = {"CustomBuildPy"} | set(_HELPERS) + selected: List[ast.stmt] = [] + for node in _setup_py_module().body: + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in wanted: + selected.append(node) + elif isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id in wanted + for target in node.targets + ): + selected.append(node) + + namespace: Dict[str, object] = { + "__file__": str(SETUP_PY), + "os": os, + "ast": ast, + "Path": Path, + "functools": functools, + "subprocess": subprocess, + "build_py": build_py, + "Dict": Dict, + "FrozenSet": FrozenSet, + "List": List, + "Set": Set, + "Tuple": Tuple, + "find_namespace_packages": find_namespace_packages, + } + exec( + compile(ast.Module(body=selected, type_ignores=[]), str(SETUP_PY), "exec"), + namespace, + ) + return namespace + + +_NAMESPACE = _load_from_setup_py() +_vendored_prefixes = _NAMESPACE["_vendored_prefixes"] +_is_vendored_path = _NAMESPACE["_is_vendored_path"] +_full_packages = _NAMESPACE["_full_packages"] + + +def _discovered_packages() -> List[str]: + """Everything setuptools finds, before any of this change's filtering.""" + return sorted( + find_namespace_packages( + where=PACKAGE_ROOT, include=["executorch", "executorch.*"] + ) + ) + + +def _vendored(packages: List[str]) -> List[str]: + return [ + package for package in packages if _is_vendored_path(package.replace(".", "/")) + ] + + +class TestFullWheelPackages(unittest.TestCase): + def test_the_tree_has_vendored_packages_to_exclude(self) -> None: + """Fail rather than skip when there is nothing to exclude. + + Every other test here is vacuous on a tree with no vendored checkouts: an empty + package list contains no vendored package, so the exclusion would look correct even + if it had been deleted. Assert the premise instead of quietly passing on it. + """ + discovered = _discovered_packages() + self.assertNotEqual( + discovered, [], f"no packages discovered under {PACKAGE_ROOT}" + ) + self.assertNotEqual( + _vendored(discovered), + [], + "no vendored third-party packages in this tree, so the exclusion below cannot " + "be shown to do anything. Initialize the submodules before running this.", + ) + + def test_no_vendored_package_ships(self) -> None: + """No package from another repository is published. + + Compares against what discovery finds rather than re-filtering the helper's own output. + Filtering the result with the same predicate the helper already applied is a tautology: + it is empty whatever the helper did, so it would pass even with the exclusion removed. + """ + discovered = set(_discovered_packages()) + shipped = set(_full_packages()) + dropped = discovered - shipped + + leaked = sorted(shipped & set(_vendored(discovered))) + # Only the count and a few names, because a regression here leaks hundreds of + # packages and the default diff would bury the message. + self.assertEqual( + len(leaked), + 0, + f"the wheel would publish {len(leaked)} vendored packages, " + f"e.g. {leaked[:3]}", + ) + # And the helper really removed them, rather than discovery never having found them. + self.assertEqual( + dropped, + set(_vendored(discovered)), + "the set the helper drops is not the set of vendored packages on disk", + ) + + def test_the_exclusion_is_load_bearing(self) -> None: + """Discovery without the exclusion finds the packages the exclusion removes.""" + self.assertLess( + len(_full_packages()), + len(_discovered_packages()), + "the exclusion dropped nothing, so it is no longer doing any work", + ) + + def test_setup_passes_the_package_list(self) -> None: + """The helper is actually wired into the full build. + + Without this, every test above still passes when the assignment that hands the list + to setuptools is deleted, which is the whole of the change. The sibling wheel test + asserts its own wiring the same way and for the same reason. + + The search is limited to the else branch of the minimal-build check, because an + unrestricted walk also matches an assignment that can never run: moved into the + minimal branch it is overwritten by the next line, and wrapped in a false condition + it is dead, and both of those leave the full wheel discovering everything. + """ + minimal_checks = [ + node + for node in _setup_py_module().body + if isinstance(node, ast.If) + and isinstance(node.test, ast.Call) + and isinstance(node.test.func, ast.Name) + and node.test.func.id == "_is_minimal_build" + ] + self.assertEqual( + len(minimal_checks), + 1, + "expected exactly one module level `if _is_minimal_build():`", + ) + + assigned = [ + node + for node in minimal_checks[0].orelse + if isinstance(node, ast.Assign) + for target in node.targets + if isinstance(target, ast.Subscript) + and isinstance(target.value, ast.Name) + and target.value.id == "setup_kwargs" + and isinstance(target.slice, ast.Constant) + and target.slice.value == "packages" + and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) + and node.value.func.id == "_full_packages" + ] + self.assertEqual( + len(assigned), + 1, + "setup.py does not assign _full_packages() to setup_kwargs['packages'], " + "so the full build falls back to discovering every package", + ) + + def test_first_party_packages_still_ship(self) -> None: + """A named first-party package survives the exclusion. + + Every other test here asks whether unwanted packages left. This one asks whether + wanted ones stayed, which is the failure mode a too-greedy filter produces and the + one nothing else would notice. + """ + packages = _full_packages() + for package in ( + "executorch.exir", + "executorch.backends.xnnpack", + "executorch.extension.pybindings", + "executorch.devtools", + ): + self.assertIn(package, packages) + + def test_only_submodule_sections_are_read(self) -> None: + """A `path` line outside a submodule section is not an exclusion prefix. + + Written against a file with a stray entry rather than by comparing git's output to git's + own output. That comparison holds for any reader on today's clean file, so it would pass + just as well for a line scanner that accepts `path =` from any section, which is the + failure this is meant to catch: one stray line silently removes a real package. + """ + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / ".gitmodules").write_text( + '[submodule "real"]\n' + "\tpath = extension/llm/tokenizers\n" + "[core]\n" + "\tpath = executorch/exir\n" + ) + self.assertEqual( + _load_from_setup_py(root)["_vendored_prefixes"](), + ("extension/llm/tokenizers",), + "a path line outside a submodule section became an exclusion prefix", + ) + + def test_prefixes_are_normalized(self) -> None: + """A legal but unusual spelling in .gitmodules still matches the real directory. + + git treats a trailing slash, a leading ./ and a doubled separator as the same path, + so storing the raw text would silently disable the exclusion for that entry. Asserted + against a written file rather than against today's values, because today's are already + tidy and would pass either way. + """ + for spelling in ( + "extension/llm/tokenizers/", + "./extension/llm/tokenizers", + "extension//llm/tokenizers", + ): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + # Only the path is read, so the entry needs no url. + (root / ".gitmodules").write_text( + f'[submodule "t"]\n\tpath = {spelling}\n' + ) + # The helper reads .gitmodules beside its own setup.py, so it is loaded + # against the temporary tree rather than the real one. + prefixes = _load_from_setup_py(root)["_vendored_prefixes"]() + self.assertEqual( + prefixes, + ("extension/llm/tokenizers",), + f"{spelling!r} did not normalize", + ) + + def test_the_fallback_matches_gitmodules(self) -> None: + """The hardcoded fallback still lists the same submodules the file does. + + It is only used when .gitmodules cannot be read, which is the case in a source + distribution, so nothing else would notice it drifting out of date. + """ + self.assertEqual( + _vendored_prefixes(), _NAMESPACE["_VENDORED_SUBMODULE_FALLBACK"] + ) + + # And it is actually returned when the file is missing, which is the only case it + # exists for. Without this the fallback could be replaced by an empty tuple and the + # comparison above would still hold. + with tempfile.TemporaryDirectory() as tmp: + self.assertEqual( + _load_from_setup_py(Path(tmp))["_vendored_prefixes"](), + _NAMESPACE["_VENDORED_SUBMODULE_FALLBACK"], + "with no .gitmodules the submodule exclusion silently does nothing", + ) + + def test_a_submodule_name_with_a_space_is_read(self) -> None: + """A submodule whose NAME contains a space still yields its path. + + git prints " " and permits spaces in the name, so splitting on the first + space truncates the key and leaves a value that matches no directory, turning the + exclusion off for that entry. + """ + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / ".gitmodules").write_text( + '[submodule "my module"]\n\tpath = extension/llm/tokenizers\n' + ) + self.assertEqual( + _load_from_setup_py(root)["_vendored_prefixes"](), + ("extension/llm/tokenizers",), + ) + + def test_a_broken_gitmodules_falls_back(self) -> None: + """An unreadable .gitmodules reaches the fallback rather than excluding nothing. + + git exits non-zero with empty output on a bad section header or on conflict markers. + Reading that as "this repository has no submodules" would turn the exclusion off with + no warning, which is the one failure the fallback exists to prevent. + """ + for broken in ( + '[submodule "x"\n\tpath = extension/llm/tokenizers\n', + '<<<<<<< HEAD\n[submodule "x"]\n\tpath = a/b\n=======\n', + "", + ): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / ".gitmodules").write_text(broken) + prefixes = _load_from_setup_py(root)["_vendored_prefixes"]() + self.assertEqual( + prefixes, + _NAMESPACE["_VENDORED_SUBMODULE_FALLBACK"], + f"a broken .gitmodules ({broken[:20]!r}) silently excluded nothing", + ) + + def test_manifest_filter_actually_drops_vendored_data(self) -> None: + """The data-file half of the fix removes files, through the real build code path. + + `packages` only governs Python modules. Non-Python files arrive through the + package_data manifest, and setuptools attributes a file under an unlisted directory to + its nearest listed parent, so vendored data returns unless the manifest is filtered too. + + Drives CustomBuildPy.analyze_manifest itself rather than reimplementing the filter here. + Checking the predicate in isolation is not enough: inverting the editable guard or + short-circuiting the condition leaves the predicate correct and the build unfiltered, + and both of those left an earlier version of this test green. + """ + namespace = _load_build_py() + build_py_class = namespace["CustomBuildPy"] + + vendored = ( + "src/executorch/backends/xnnpack/third-party/generate-cpuinfo-wrappers.py" + ) + # A shader template goes through this same filter, and it needs its own example here: + # deleting the shader line leaves the vendored assertions below green, so the manifest + # half of the shader fix was unprotected. + shader = "src/executorch/backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml" + ordinary = "setup.py" + # All of them have to exist on disk, because the filter also drops anything that is not a + # file, and a missing path would be removed for that reason instead of this one. + for relative in (vendored, shader, ordinary): + self.assertTrue( + (REPO_ROOT / relative).is_file(), + f"{relative} is gone, so this test needs a different example", + ) + + # analyze_manifest calls up into setuptools first, which needs the full command + # machinery. Only the filtering after that call is under test, so the parent's method + # is replaced with a no-op for the duration and the manifest seeded directly. This runs + # the shipped code path rather than a copy of it, which is the point: a filter that has + # been turned off still reads correctly in the source. + parent = build_py_class.__mro__[1] + original = parent.analyze_manifest + parent.analyze_manifest = lambda self: None + try: + stub = build_py_class.__new__(build_py_class) + stub.editable_mode = False + stub.manifest_files = {"executorch": [vendored, shader, ordinary]} + stub.analyze_manifest() + kept = stub.manifest_files["executorch"] + finally: + parent.analyze_manifest = original + + self.assertNotIn( + vendored, kept, "a vendored data file survived the manifest filter" + ) + self.assertNotIn(shader, kept, "a shader template survived the manifest filter") + self.assertIn(ordinary, kept, "the filter dropped an ordinary file") + + def test_is_vendored_path_matches_whole_components(self) -> None: + """The filter matches a path component, not a substring.""" + self.assertTrue( + _is_vendored_path( + "src/executorch/backends/xnnpack/third-party/XNNPACK/a.py" + ) + ) + self.assertTrue(_is_vendored_path("src/executorch/x/third_party/y.yaml")) + self.assertFalse(_is_vendored_path("src/executorch/exir/program/_program.py")) + # "third-party" as part of a longer name is a different directory. + self.assertFalse(_is_vendored_path("src/executorch/x/third-party-tools/y.py")) + + def test_submodules_outside_a_vendored_dir_are_recognized(self) -> None: + """A submodule checked out under an ordinary name is still another repository. + + These are not matched by the directory name, so they are read from .gitmodules. Their + nested copies also cannot satisfy the imports the code uses: the FACTO helper imports + facto.specdb from the top level, and the tokenizers ship as a declared dependency. + """ + prefixes = _vendored_prefixes() + self.assertIn("backends/cadence/utils/FACTO", prefixes) + self.assertIn("extension/llm/tokenizers", prefixes) + for prefix in ("backends/cadence/utils/FACTO", "extension/llm/tokenizers"): + self.assertTrue(_is_vendored_path(f"executorch/{prefix}")) + self.assertTrue(_is_vendored_path(f"src/executorch/{prefix}/setup.py")) + self.assertFalse( + _is_vendored_path("executorch/extension/llm/custom_ops/op_sdpa.py") + ) + + def test_root_level_submodules_are_not_listed(self) -> None: + """A submodule at the repository root is not a wheel path. + + Those are build tooling, never copied into the package, and listing one would put a + bare single-word name into the matcher. That would then drop any directory sharing the + name, anywhere in the tree, which is a much wider rule than intended. + """ + for prefix in _vendored_prefixes(): + self.assertIn( + "/", + prefix, + f"{prefix!r} is a root-level submodule and must not be listed", + ) + self.assertFalse(_is_vendored_path("executorch/some/nested/shim")) + + +if __name__ == "__main__": + unittest.main() diff --git a/pyproject.toml b/pyproject.toml index 105565862df..f7cf98679fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,9 +100,6 @@ Changelog = "https://github.com/pytorch/executorch/releases" [project.scripts] flatc = "executorch.data.bin:flatc" -# TODO(dbort): Could use py_modules to restrict the set of modules we -# package, and package_data to restrict the set up non-python files we -# include. See also setuptools/discovery.py for custom finders. [tool.setuptools] license-files = ["LICENSE"] @@ -118,10 +115,9 @@ license-files = ["LICENSE"] "executorch" = "src/executorch" [tool.setuptools.package-data] -# TODO(dbort): Prune /test[s]/ dirs, /third-party/ dirs, yaml files that we -# don't need. -# TODO(RobertKalmar): When test[s] dirs pruned the PROJECT_DIR resolution in backends.nxp.tests_models.config.py can -# avoid exporting and reading env variable. +# TODO(RobertKalmar): the PROJECT_DIR env variable is still needed. Test directories still install, +# since the suites import shared helpers from them, and the artifacts that config.py resolves are +# not in the wheel, so a path derived from __file__ would point at files that are not there. "*" = [ # Some backends like XNNPACK need their .fbs files. "*.fbs", diff --git a/setup.py b/setup.py index 8980d1bd94a..132e9bcd1e3 100644 --- a/setup.py +++ b/setup.py @@ -47,7 +47,9 @@ # other computer software, distribute, and sublicense such enhancements or # derivative works thereof, in binary and source code form. +import ast import contextlib +import functools # Import this before distutils so that setuptools can intercept the distuils # imports. @@ -64,7 +66,7 @@ from distutils import log # type: ignore[import-not-found] from distutils.sysconfig import get_python_lib # type: ignore[import-not-found] from pathlib import Path, PurePosixPath -from typing import List, Optional +from typing import Dict, FrozenSet, List, Optional, Set, Tuple # Clean dynamic import using importlib _install_utils_path = Path(__file__).parent / "install_utils.py" @@ -177,10 +179,117 @@ def _minimal_cmake_flags() -> List[str]: ] +_VENDORED_DIR_NAMES = frozenset({"third-party", "third_party"}) + +# Used only when .gitmodules cannot be read, as in a source distribution. A test keeps it in step. +_VENDORED_SUBMODULE_FALLBACK = ( + "backends/cadence/utils/FACTO", + "extension/llm/tokenizers", +) + + +@functools.lru_cache(maxsize=None) +def _vendored_prefixes() -> Tuple[str, ...]: + """Source-tree prefixes holding code from another repository. + + Two shapes reach the wheel. Most vendored code sits in a directory named third-party, + which the name above covers wherever it appears. The rest are git submodules checked out + under an ordinary name, so they can only be recognized by asking git what they are. + + None of them are importable from where they sit. FACTO is pure Python but its nested copy + cannot satisfy backends/cadence/utils/facto_util.py, which imports the top level facto.specdb, + and the tokenizers ship separately as pytorch-tokenizers in the dependency list. The rest, + XNNPACK and the Vulkan headers among them, are C++ sources that the wheel has no use for once + the libraries are built. + + Read through git rather than by scanning the file, so only real submodule entries count. + A hand-rolled reader accepts a `path` line from any section, and one stray line elsewhere + in the file would drop a first-party package from the wheel with nothing to warn about. + + Submodules at the repository root are skipped. Those are build tooling, never copied into + the package, and carrying a bare single-word name here would make the match below drop any + directory that happened to share it. + """ + root = Path(__file__).parent + if not (root / ".gitmodules").is_file(): + # A source distribution carries no .gitmodules, so nothing can be read there. Fall + # back to the directories the vendored trees occupy, or the exclusion would quietly + # do half its job and those files would ship again. + return _VENDORED_SUBMODULE_FALLBACK + try: + listed = subprocess.run( + [ + "git", + "config", + "-z", + "-f", + ".gitmodules", + "--get-regexp", + r"^submodule\..*\.path$", + ], + cwd=root, + capture_output=True, + text=True, + check=False, + ) + except OSError: + # No git on PATH, so fall back for the same reason as above. + return _VENDORED_SUBMODULE_FALLBACK + + if listed.returncode or not listed.stdout.strip(): + # git ran and told us nothing useful, which happens when the file has a bad section + # header or conflict markers in it. Reading that as "no submodules" would turn the + # exclusion off without a word, so fall back rather than trust an empty answer. + return _VENDORED_SUBMODULE_FALLBACK + + prefixes = [] + # -z separates each record with NUL and its key from its value with a newline, so neither a + # name nor a path containing a space can be misread. Splitting the default space-separated + # output cannot do that: "submodule.a b.path c d/e" is ambiguous either way round. + for record in listed.stdout.split("\0"): + if not record: + continue + _, separator, value = record.partition("\n") + if not separator: + continue + # Normalize, because git accepts a trailing slash, a ./ prefix and doubled + # separators as the same path, and the raw text would stop matching the real + # directory. + parts = Path(value).parts + if len(parts) < 2 or any(part in _VENDORED_DIR_NAMES for part in parts): + continue + prefixes.append("/".join(parts)) + return tuple(sorted(prefixes)) + + +def _is_vendored_path(path: str) -> bool: + """Whether a source-tree path holds code from another repository.""" + parts = Path(path).parts + if any(part in _VENDORED_DIR_NAMES for part in parts): + return True + # A submodule path is relative to the repository root, while a path here may be relative + # to src/executorch or carry a src/executorch prefix, so match on any suffix boundary. + # Whole-component match: the prefix must be the entire path, or sit at its start, end, or + # middle bounded by separators. Substring matching would let a directory whose name merely + # begins with a prefix be dropped. + posix = "/".join(parts) + return any( + posix == prefix + or posix.startswith(f"{prefix}/") + or posix.endswith(f"/{prefix}") + or f"/{prefix}/" in posix + for prefix in _vendored_prefixes() + ) + + def _minimal_packages() -> List[str]: return sorted( find_namespace_packages( - where="src", + # Anchored on this file, not the working directory, so the list does not change + # with where the build was started from. A cwd-relative path returns nothing when + # the build runs from anywhere but the repository root, and a wheel with no packages + # in it ships no Python at all. + where=str(Path(__file__).parent / "src"), include=[ "executorch", "executorch.data", @@ -204,6 +313,334 @@ def _minimal_packages() -> List[str]: ) +_WALK_SKIP_DIRS = frozenset( + {".git", "pip-out", "cmake-out", "third-party", "third_party", "__pycache__"} +) + +_TEST_DIR_NAMES = frozenset({"test", "tests"}) + + +@functools.lru_cache(maxsize=None) +def _top_level_package_dirs() -> FrozenSet[str]: + """The first path component of every package the wheel ships. + + Derived from the tree rather than listed, so a new top-level directory is covered without an + edit here. Used to recognize the unprefixed spelling of a first-party import. + """ + root = Path(__file__).parent / "src" / "executorch" + if not root.is_dir(): + return frozenset() + return frozenset(entry.name for entry in root.iterdir() if entry.is_dir()) + + +# Named only by a workflow or by a documented command, so no import reaches them. Listed here +# rather than scanned from .github, which a source distribution does not carry; a test re-derives +# the list so it cannot drift. +_CI_ENTRY_POINTS = ( + "executorch.backends.mlx.test.run_all_tests", + "executorch.backends.mlx.test.test_sample", + "executorch.backends.mlx.test.test_slot_recycling", + "executorch.backends.samsung.test.utils.run_tests", + "executorch.backends.test.suite.generate_markdown_summary_json", + "executorch.examples.models.muse_glimmer.tests.gen_prompt_golden", + "executorch.examples.models.muse_glimmer.tests.test_mlx_pipeline", + "executorch.examples.models.muse_glimmer.tests.test_prompt_tokens", + "executorch.extension.pybindings.test.test_pybindings", +) + +# Directories whose test modules are reached without any import statement naming them, so no scan +# of the source can find them: mlx.yml runs each file it discovers under custom_kernel_ops, the +# webgpu scripts import one module per operator, runner.py resolves a suite root out of a dict and +# then walks it, and the llava README documents a `python -m` command. Directories rather than file +# names, so a new test is covered when it is added. +_CI_ENTRY_POINT_DIRS = ( + "executorch.backends.mlx.custom_kernel_ops", + "executorch.backends.webgpu.test", + "executorch.backends.test.suite", + "executorch.examples.models.llava.test", +) + + +def _is_test_module(dotted: str) -> bool: + return any(part in _TEST_DIR_NAMES for part in dotted.split(".")) + + +def _module_name(root: Path, path: Path) -> str: + parts = list(path.relative_to(root).parts) + if parts[-1] == "__init__.py": + parts = parts[:-1] + else: + parts[-1] = parts[-1][: -len(".py")] + return ".".join(["executorch"] + parts) + + +def _first_party_module(name: str) -> Optional[str]: + """The `executorch.`-prefixed spelling of an import target, or None if it is not ours. + + This repository imports itself two ways. Most code says `executorch.backends.x`, but some + says `backends.x`, which resolves because pytest puts the repository root on sys.path. Both + name the same file, so both have to count as a reference or a helper reached only by the + second spelling is dropped from the wheel while its importers still expect it. + """ + if name.startswith("executorch."): + return name + if name.split(".", 1)[0] in _top_level_package_dirs(): + return f"executorch.{name}" + return None + + +def _scan_imports(path: Path, package: str, out: Set[str], dynamic: Set[str]) -> None: + """Collect into out the executorch modules one file refers to.""" + try: + tree = ast.parse(path.read_text(encoding="utf-8", errors="replace")) + except SyntaxError: + return + for node in ast.walk(tree): + out.update(_import_targets(node, package, dynamic)) + + +_GENERATED_DIR_NAMES = frozenset( + { + ".venv", + "venv", + "build", + "dist", + "buck-out", + ".cache", + ".hypothesis", + ".mypy_cache", + ".pytest_cache", + ".tox", + "test-build", + "arm_test", + "riscv_test", + } +) + + +def _unshipped_directories(root: Path) -> List[Path]: + """Checkout directories the wheel does not carry, whose imports still have to be followed. + + src/executorch is a subset of the repository, so a file under test/ or tools/ is never + packaged, yet a module it imports still has to ship. + + Generated directories are left out, because a build tree or an in-tree virtualenv holds an + INSTALLED copy of this package, and reading it would let the last wheel vote on what the next + one ships. Listed by name rather than asked of git, because `git check-ignore` needs a working + repository and answers differently for a pattern with a trailing slash depending on whether the + directory exists yet, which made the same build behave differently on two platforms. + """ + repository = Path(__file__).parent + if not repository.is_dir() or not root.is_dir(): + return [] + shipped = {entry.name for entry in root.iterdir()} + return [ + entry + for entry in sorted(repository.iterdir()) + if entry.is_dir() + and entry.name not in shipped + and entry.name not in _WALK_SKIP_DIRS + and entry.name not in _GENERATED_DIR_NAMES + and entry.name not in ("src", ".github") + ] + + +def _import_graph(root: Path) -> Tuple[Set[str], Dict[str, Set[str]], Set[str]]: + """Every module under root, what each imports, and literal importlib targets. + + The walk covers root, but the SEED covers more: a file elsewhere in the checkout can import + a module that ships, so its imports are collected too and attributed to a synthetic name. + Without that, a helper whose only importer lives outside the shipped tree looks unreachable. + """ + modules: Set[str] = set() + edges: Dict[str, Set[str]] = {} + dynamic: Set[str] = set() + + # followlinks, because src/executorch is a tree of symlinks into the repository root. + for dirpath, dirnames, filenames in os.walk(root, followlinks=True): + # Vendored trees are skipped by the same test that excludes them from the package list, + # not only by directory name. A submodule checked out under an ordinary name, FACTO and + # the tokenizers among them, is otherwise read as first-party, and its imports would keep + # test modules the wheel has no reason to carry. + dirnames[:] = [ + d + for d in dirnames + if d not in _WALK_SKIP_DIRS + and not _is_vendored_path(os.path.relpath(os.path.join(dirpath, d), root)) + ] + for filename in filenames: + if not filename.endswith(".py"): + continue + path = Path(dirpath) / filename + me = _module_name(root, path) + modules.add(me) + package = me if filename == "__init__.py" else me.rsplit(".", 1)[0] + _scan_imports(path, package, edges.setdefault(me, set()), dynamic) + + # Directories of the checkout that the wheel does not ship, test/ among them. Their files are + # never packaged, so they are not modules, but what they import must still ship: for example + # test/end2end/test_end2end.py imports two model helpers out of exir/tests. + for entry in _unshipped_directories(root): + for dirpath, dirnames, filenames in os.walk(entry, followlinks=False): + dirnames[:] = [d for d in dirnames if d not in _WALK_SKIP_DIRS] + for filename in filenames: + if not filename.endswith(".py"): + continue + outside = f"{dirpath}/{filename}" + _scan_imports( + Path(dirpath) / filename, + "", + edges.setdefault(outside, set()), + dynamic, + ) + + return modules, edges, dynamic + + +def _import_targets(node: ast.AST, package: str, dynamic: Set[str]) -> Set[str]: + """The executorch modules one AST node refers to.""" + found: Set[str] = set() + if isinstance(node, ast.Import): + found.update( + name for name in (_first_party_module(a.name) for a in node.names) if name + ) + elif isinstance(node, ast.ImportFrom): + if node.level: + if not package: + # A file outside the shipped tree, so a relative import stays inside that tree + # and cannot name anything the wheel carries. + return found + # A relative import names a real module too, and inside a kept package its target + # has to ship: stages/__init__.py does `from .export import Export`, so dropping + # stages.export would break every importer of that package. + parts = package.split(".") + if node.level > 1: + parts = parts[: len(parts) - (node.level - 1)] + base = ".".join(parts + (node.module.split(".") if node.module else [])) + elif node.module and (prefixed := _first_party_module(node.module)): + base = prefixed + else: + return found + if base.startswith("executorch"): + found.add(base) + # `from pkg import name` may name a submodule rather than an attribute, and there + # is no way to tell without importing, so both readings are kept. + found.update(f"{base}.{a.name}" for a in node.names) + elif ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "import_module" + and node.args + and isinstance(node.args[0], ast.Constant) + and isinstance(node.args[0].value, str) + ): + target = node.args[0].value + if target.startswith("."): + if not package: + return found + target = package + target + resolved = _first_party_module(target) + if resolved: + dynamic.add(resolved) + return found + + +@functools.lru_cache(maxsize=None) +def _reachable_test_modules() -> FrozenSet[str]: + """Test modules something can still reach once the wheel is installed. + + A test case that nothing imports is dead weight in the wheel: pytest loads it from a path in + the checkout, never through the installed name. A shared helper is the opposite, because the + suites import each other by installed name, so it has to ship or collection breaks. + + Reachable means named by something, anywhere in the checkout, including by a test module + itself. That looks circular and is not: a test collected from the checkout still resolves + `from executorch.x.test import helper` through the INSTALLED package, so the helper must be + in the wheel even though the file importing it is not. + """ + root = Path(__file__).parent / "src" / "executorch" + modules, edges, dynamic = _import_graph(root) + + # Every name anything refers to. No transitive walk is needed: this is already the union of + # every edge target, so following an edge could only rediscover a name that is in here. + referenced = set(dynamic) + referenced.update(_CI_ENTRY_POINTS) + for targets in edges.values(): + referenced.update(targets) + + keep = {name for name in referenced if _is_test_module(name)} & modules + # Everything under a directory whose tests are run one file at a time by a discovery loop. + keep |= { + name + for name in modules + if _is_test_module(name) + and any( + name == prefix or name.startswith(f"{prefix}.") + for prefix in _CI_ENTRY_POINT_DIRS + ) + } + # Parent packages of anything kept, or the dotted path cannot resolve. + for name in list(keep): + parts = name.split(".") + for end in range(2, len(parts)): + parent = ".".join(parts[:end]) + if _is_test_module(parent): + keep.add(parent) + return frozenset(keep) + + +_SHADER_TEMPLATE_MARKERS = ( + "parameter_names_with_default_values", + "shader_variants", + "generate_variant_forall", +) + + +@functools.lru_cache(maxsize=None) +def _is_shader_template(path: str) -> bool: + """Whether a yaml file is a shader codegen input rather than data the wheel needs. + + gen_vulkan_spv.py and gen_wgsl_headers.py expand these into SPIR-V and WGSL headers during + the cmake build, so the wheel already carries the compiled result. Matched on content rather + than on a directory list, because the same shape appears under vulkan and webgpu and a path + list goes stale as soon as a backend adds one. The op and kernel definitions that ARE read + at run time, edge.yaml among them, carry none of these keys. + """ + if not path.endswith(".yaml"): + return False + full = Path(__file__).parent / path + try: + head = full.read_text(encoding="utf-8", errors="replace") + except OSError: + return False + return any(marker in head for marker in _SHADER_TEMPLATE_MARKERS) + + +def _full_packages() -> List[str]: + """Every package the full wheel ships. + + Without an explicit list setuptools discovers all of src/executorch, which pulls in the + Python files and codegen scripts of the vendored third-party checkouts. Those exist to + build the C++ targets, so once the libraries are built no shipped module imports them. + + Test packages deliberately stay. The suites in this repository import each other through + the installed name, for example `from executorch.backends.arm.test import common`, so + dropping them from the wheel stops the suites collecting under a non-editable install. + """ + return sorted( + package + # Anchored on this file rather than the working directory, so the list does not + # change with where the build or a test was started from. + for package in find_namespace_packages( + where=str(Path(__file__).parent / "src"), + include=["executorch", "executorch.*"], + ) + # The include patterns above DO match these, since they are ordinary dotted names, + # which is exactly why they have to be removed here instead. + if not _is_vendored_path(package.replace(".", "/")) + ) + + # The published project names for the CUDA runtime components a CUDA wheel links but # does not bundle, keyed by CUDA major version. Not derivable from a suffix rule: the # CUDA 12 wheels carry a "-cu12" suffix while the CUDA 13 ones are published under @@ -1628,6 +2065,68 @@ class CustomBuildPy(build_py): a file to a different relative location under the output package directory. """ + def _prune_unstaged_files(self) -> None: + """Delete .py and .yaml this command staged in an earlier build and no longer wants. + + Restricted to files that exist in the source tree, because build_py is the first of + build's sub commands and everything a later one stages is still sitting in the build + directory when this runs. build_ext generates executorch/data/bin/__init__.py, the + target of the flatc console script, from a template that lives elsewhere, so a walk + that removed anything absent from build_py's own file list would delete it. + """ + if self.editable_mode: + return + if not self.packages: + # Nothing to compare against, so every staged file would look unwanted. Refuse + # rather than empty the build directory. + return + + wanted = set() + for package, module, _ in self.find_all_modules(): + parts = package.split(".") if package else [] + wanted.add(os.path.join(self.build_lib, *parts, f"{module}.py")) + for package in self.packages or (): + src_dir = self.get_package_dir(package) + build_dir = os.path.join(*([self.build_lib] + package.split("."))) + for filename in self.find_data_files(package, src_dir): + wanted.add(os.path.join(build_dir, os.path.relpath(filename, src_dir))) + + source_root = Path(__file__).parent / "src" + for dirpath, _, filenames in os.walk(self.build_lib): + for filename in filenames: + if not filename.endswith((".py", ".yaml")): + continue + staged = os.path.join(dirpath, filename) + if staged in wanted: + continue + relative = os.path.relpath(staged, self.build_lib) + if not (source_root / relative).is_file(): + # Generated by another command, so build_py must not remove it. + continue + os.remove(staged) + + def find_package_modules(self, package, package_dir): + modules = super().find_package_modules(package, package_dir) + if self.editable_mode or not _is_test_module(package): + # An editable install exposes the whole source tree whatever is listed here, and a + # package outside a test directory has nothing to drop. + return modules + keep = _reachable_test_modules() + return [ + entry + for entry in modules + if entry[1] == "__init__" or f"{package}.{entry[1]}" in keep + ] + + def find_data_files(self, package, src_dir): + files = super().find_data_files(package, src_dir) + if self.editable_mode: + return files + root = os.path.dirname(os.path.abspath(__file__)) + return [ + _f for _f in files if not _is_shader_template(os.path.relpath(_f, root)) + ] + def analyze_manifest(self): super().analyze_manifest() # Recent versions of setuptools may include bare directory symlinks from version @@ -1642,6 +2141,13 @@ def analyze_manifest(self): _f for _f in self.manifest_files[_pkg] if os.path.isfile(os.path.join(_root, _f)) + # A directory left out of `packages` is not simply skipped. setuptools + # walks up to the nearest listed package and records the file as that + # package's data, so a vendored *.yaml still arrives under its parent. + # Filter with the same list so the two agree. + and not _is_vendored_path(_f) + # Shader templates are consumed by the cmake build, not at run time. + and not _is_shader_template(_f) ] def _copy_extra_files(self, src_to_dst, dst_root: str) -> None: @@ -1684,6 +2190,12 @@ def run(self): # defined by the py_module list and package_data patterns. build_py.run(self) + # A rebuild over a staging directory left by an earlier build keeps whatever that build + # put there, because build_py only ever copies and never deletes. So a file this build + # deliberately leaves out is still present from last time, and the wheel packages it. + # Remove what is no longer wanted rather than only skipping the copy. + self._prune_unstaged_files() + # dst_root is the root of the `executorch` module in the output package # directory. build_lib is the platform-independent root of the output # package, and will look like `pip-out/lib`. It can contain multiple @@ -2325,6 +2837,7 @@ def iter_distribution_names(self): setup_kwargs["packages"] = _minimal_packages() setup_kwargs["install_requires"] = _minimal_dependencies() else: + setup_kwargs["packages"] = _full_packages() # A CUDA wheel links the CUDA runtime but does not bundle it, so the wheels that # carry it are declared here. A CPU wheel adds nothing. setup_kwargs["install_requires"] = _base_dependencies() + _cuda_dependencies()