forked from pytorch/executorch
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathsetup.py
More file actions
3180 lines (2827 loc) · 152 KB
/
Copy pathsetup.py
File metadata and controls
3180 lines (2827 loc) · 152 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# Copyright 2024, 2026 Arm Limited and/or its affiliates.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
# Part of this code is from pybind11 cmake_example, so attach the license below.
# That project has since dropped setup.py, so this points at the last revision
# that still had it instead of at a branch.
# https://github.com/pybind/cmake_example/blob/7a94877f581a14de4de1a096fb053a55fc2a66bf/setup.py
# Copyright (c) 2016 The Pybind Development Team, All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its contributors
# may be used to endorse or promote products derived from this software
# without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# You are under no obligation whatsoever to provide any bug fixes, patches, or
# upgrades to the features, functionality or performance of the source code
# ("Enhancements") to anyone; however, if you choose to make your Enhancements
# available either publicly, or directly to the author of this software, without
# imposing a separate written license agreement for such Enhancements, then you
# hereby grant the following license: a non-exclusive, royalty-free perpetual
# license to install, use, modify, prepare derivative works, incorporate into
# 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.
import importlib.util
import logging
import os
import re
import shlex
import shutil
import site
import stat
import subprocess
import sys
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 Dict, FrozenSet, List, Optional, Set, Tuple
# Clean dynamic import using importlib
_install_utils_path = Path(__file__).parent / "install_utils.py"
_spec = importlib.util.spec_from_file_location("install_utils", _install_utils_path)
if _spec is None:
raise ImportError(f"Could not create module spec for {_install_utils_path}")
install_utils = importlib.util.module_from_spec(_spec)
if _spec.loader is None:
raise ImportError(f"Module spec has no loader for {_install_utils_path}")
_spec.loader.exec_module(install_utils)
from setuptools import Distribution, Extension, find_namespace_packages, setup
from setuptools.command.build import build
from setuptools.command.build_ext import build_ext
from setuptools.command.build_py import build_py
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(message)s",
)
# Headers swept in by a directory copy that a consumer of the wheel cannot use, because each needs
# something the wheel does not carry. Publishing one is worse than leaving it out: the failure arrives in
# someone else's project rather than here.
#
# Matched on the path ending, not the bare file name, so an entry names one specific header rather than
# every file that happens to share its name. Anything under a directory beginning with "test" is already
# dropped separately, so those need no entry here.
#
# Only headers that nothing else the wheel installs includes belong here. A header other shipped headers
# pull in must keep shipping even when it cannot be compiled on its own.
_UNSHIPPABLE_HEADERS = frozenset(
{
# Needs a header generated when the schema is compiled, which in turn needs the FlatBuffers C++
# headers. Those are a third-party library this wheel does not vendor.
"runtime/executor/tensor_parser.h",
# Reads processor details through cpuinfo, whose headers the wheel does not publish.
"extension/threadpool/cpuinfo_utils.h",
# Holds a pthreadpool member by value, so it needs that library's header, which the wheel does not
# publish either. The component it belongs to is a link dependency the runtime carries, not
# something a consumer includes.
"extension/threadpool/threadpool.h",
# Declares CPUCachingAllocator, whose implementation is in a component no shipped library links,
# so including it compiles and then fails at link time with an undefined reference.
"extension/memory_allocator/cpu_caching_malloc_allocator.h",
# Declares BundledModule, which is built only for the Python bindings, so its implementation is in
# the Python extension. A C++ application cannot link that, and building the source instead needs
# bundled-program headers the wheel does not publish.
"extension/module/bundled_module.h",
# Declares FileDescriptorDataLoader, whose implementation is in no CMake target at all, so no
# shipped library defines it. Including it compiles and then fails at link time.
"extension/data_loader/file_descriptor_data_loader.h",
}
)
try:
from tools.cmake.cmake_cache import CMakeCache
except ImportError:
sys.path.insert(
0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "tools", "cmake")
)
from cmake_cache import CMakeCache # type: ignore[no-redef, import-not-found]
def _is_macos() -> bool:
return sys.platform == "darwin"
def _is_windows() -> bool:
return sys.platform == "win32"
def _is_env_flag_enabled(name: str) -> bool:
return os.environ.get(name, "").strip().upper() in {"1", "ON", "TRUE", "YES"}
def _is_minimal_build() -> bool:
return _is_env_flag_enabled("EXECUTORCH_BUILD_MINIMAL")
def _minimal_cmake_flags() -> List[str]:
return [
"-DEXECUTORCH_BUILD_COREML=OFF",
"-DEXECUTORCH_BUILD_CUDA=OFF",
"-DEXECUTORCH_BUILD_DEVTOOLS=OFF",
"-DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=OFF",
"-DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=OFF",
"-DEXECUTORCH_BUILD_EXTENSION_LLM=OFF",
"-DEXECUTORCH_BUILD_EXTENSION_LLM_RUNNER=OFF",
"-DEXECUTORCH_BUILD_EXTENSION_MODULE=OFF",
"-DEXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP=OFF",
"-DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=OFF",
"-DEXECUTORCH_BUILD_EXTENSION_TENSOR=OFF",
"-DEXECUTORCH_BUILD_EXTENSION_TRAINING=OFF",
"-DEXECUTORCH_BUILD_KERNELS_CUSTOM_AOT=OFF",
"-DEXECUTORCH_BUILD_KERNELS_LLM=OFF",
"-DEXECUTORCH_BUILD_KERNELS_LLM_AOT=OFF",
"-DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=OFF",
"-DEXECUTORCH_BUILD_KERNELS_QUANTIZED=OFF",
"-DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=OFF",
"-DEXECUTORCH_BUILD_KERNELS_TORCHAO=OFF",
"-DEXECUTORCH_BUILD_MLX=OFF",
"-DEXECUTORCH_BUILD_OPENVINO=OFF",
"-DEXECUTORCH_BUILD_PORTABLE_OPS=OFF",
"-DEXECUTORCH_BUILD_PYBIND=OFF",
"-DEXECUTORCH_BUILD_QNN=OFF",
"-DEXECUTORCH_BUILD_TESTS=OFF",
"-DEXECUTORCH_BUILD_VULKAN=OFF",
"-DEXECUTORCH_BUILD_XNNPACK=OFF",
"-DEXECUTORCH_BUILD_CMSIS_NN_PYBINDS=OFF",
]
_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(
# 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",
"executorch.data.bin",
"executorch.exir",
"executorch.exir.*",
"executorch.extension",
"executorch.extension.flat_tensor",
"executorch.extension.flat_tensor.*",
"executorch.extension.pytree",
],
exclude=[
"*.test",
"*.test.*",
"*.tests",
"*.tests.*",
"*.__pycache__",
"*.__pycache__.*",
],
)
)
_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"<outside>{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
# unsuffixed names. A train with no entry here declares nothing rather than guessing a
# name that may not exist.
#
# Only what a shipped library actually loads. Measured on a built wheel, the CUDA
# libraries need the CUDA runtime and nothing else, because cuRAND is used only through
# its device-side header API, which compiles into the object rather than linking a
# library, and the generated model library embeds its kernels rather than compiling them
# at run time, so there is no runtime compiler to satisfy either.
# Bounded to the train the binaries were built against. The shipped libraries record
# NEEDED libcudart.so.<major> and a runtime path into that train's own directory, so a
# resolution to a different major installs a different layout and a different soname and the
# wheel is unimportable. Nothing catches that at install time; it surfaces as an unresolved
# libcudart on first import. The 12 package is already major-specific by name, but the 13 one
# is not, so it needs the specifier to say what the name does not.
_CUDA_RUNTIME_PACKAGES = {
"12": ("nvidia-cuda-runtime-cu12>=12,<13",),
"13": ("nvidia-cuda-runtime>=13,<14",),
}
# Where each train installs its libraries under site-packages. CUDA 13 collects them in
# one directory while CUDA 12 gives each component its own, so the search path differs by
# train and cannot be a single literal.
#
# Every declared package needs its directory here, and nothing else belongs. The loader only
# searches what is recorded here, so a missing directory leaves a shipped library unable to find
# a package that is installed, and an extra one implies a dependency the wheel does not have.
_CUDA_LIBRARY_DIRECTORIES = {
"12": ("nvidia/cuda_runtime/lib",),
"13": ("nvidia/cu13/lib",),
}
# Arch subdirectory names MKL's exported link interface appends to its prefix.
_MKL_ARCH_DIRECTORIES = ("intel64", "intel64_win", "win-x64")
def _cmake_args() -> List[str]:
"""CMAKE_ARGS split into arguments, tolerating an unbalanced quote.
shlex is the correct parser for a value that names a shell argument list, but it raises on an
unbalanced quote, and a path containing an apostrophe is enough to trigger it. Both callers run at
module scope, so the exception surfaced as a traceback during the build rather than a diagnosable
error. Falling back to whitespace splitting keeps the build working for the case that caused it.
"""
raw = os.environ.get("CMAKE_ARGS", "")
try:
return shlex.split(raw)
except ValueError:
return raw.split()
# Release rows that ship no CUDA. Only the metadata side reads this. The shell classifier is
# deliberately broader, treating every name it does not recognise as CPU, so the two do not
# agree on rows like rocm or xpu and are not meant to. What matters is that a row this side
# calls CPU declares no CUDA runtime, which is what the list is for.
_CPU_ROW_NAMES = ("cpu", "cpu-aarch64")
def _row_is_cpu_only() -> bool:
"""Whether the release row this build belongs to names itself a CPU row.
The metadata side already reads the row to decide which NVIDIA packages to declare, so the build
has to read the same input or the two disagree and the wheel ships a delegate it cannot load.
Absent means unknown rather than CPU, which keeps a plain local build behaving as before.
"""
raw = (
(os.environ.get("CU_VERSION") or os.environ.get("DESIRED_CUDA") or "")
.strip()
.lower()
)
return raw in _CPU_ROW_NAMES
def _cuda_train() -> str:
"""The CUDA major version this wheel is being built for, or "" for a CPU wheel.
The release row's own field wins when it is set, because a row states the train it
targets and that is more authoritative than whichever toolkit happens to sit on the
builder. The wheel build exports CU_VERSION; DESIRED_CUDA is the matrix field name.
Falling back to the installed toolkit matters for every build that is not a release
job. The build turns CUDA on by detecting a toolkit, so keying only off the release
field produced a wheel that carried the CUDA libraries with no dependency declarations
and no way to find the CUDA runtime.
Returns "" when the build did not enable CUDA, so a CPU wheel declares nothing even on
a machine that has a toolkit installed.
Raises when a release row names a train the installed toolkit does not provide. The
declared packages and the loader paths both come from this value, so disagreeing with
the toolkit that compiled the libraries produces a wheel that installs cleanly and then
cannot load: a cu126 row built against a 13.0 toolkit declares the CUDA 12 runtime for
binaries that need libcudart.so.13.
"""
# An explicit OFF first, ahead of the release field. A CPU row on a builder that has a
# toolkit installed sets both, so reading the row field first would declare a runtime
# the wheel never loads.
if not install_utils.is_cmake_option_on(
_cmake_args(),
"EXECUTORCH_BUILD_CUDA",
default=True,
):
return ""
raw = os.environ.get("CU_VERSION") or os.environ.get("DESIRED_CUDA") or ""
# A row spelled "cpu" is a CPU row, unless the caller asked for CUDA explicitly. The
# shortcut exists so a local build named that way with nothing requested does not fall
# through and raise on the unsupported-train branch below. It must not swallow an explicit
# request, because that request still reaches CMake, which builds the CUDA libraries: the
# wheel would then carry them with no dependency declared and no path to the runtime.
if raw.lower() in _CPU_ROW_NAMES and not install_utils.is_cmake_option_on(
_cmake_args(), "EXECUTORCH_BUILD_CUDA", default=False
):
return ""
if raw.lower() in _CPU_ROW_NAMES:
# Reached only when the caller turned CUDA on for a CPU row. Refused rather than
# guessed, because the option already reached CMake and built the libraries, so
# declaring nothing would ship them with no runtime dependency and no way to find
# it. Named here so the reader is not sent to add "cpu" to a list of CUDA trains.
raise RuntimeError(
f"the build names a CPU row while also asking for CUDA "
f"(EXECUTORCH_BUILD_CUDA=ON in CMAKE_ARGS). Those contradict: the CUDA "
f"libraries would be built and shipped with no runtime dependency declared. "
f"Pick one, either drop the option or name a CUDA row instead of {raw!r}."
)
# Reduce to digits and match against the same (major, minor) trains the shell classifier
# uses. Previously this took the first two digits and matched against major only, so a
# row spelled with an unsupported minor (say cu125) was classified CPU by the shell and
# CUDA 12 here, and the wheel then declared CUDA runtime packages for a CPU build.
digits = re.sub(r"[^0-9]", "", raw)
# The same rows as `trains` below, keeping both numbers rather than reducing to the
# major, so the guard can tell two rows of one major apart.
_CUDA_ROW_MINORS = {
f"{major}{minor}": (major, minor)
for major, minor in install_utils.SUPPORTED_CUDA_VERSIONS
}
trains = {
f"{major}{minor}": str(major)
for major, minor in install_utils.SUPPORTED_CUDA_VERSIONS
}
requested = trains.get(digits, "")
# Read the toolkit version directly, without the (major, minor) validator, so the guard
# below fires on any mismatch rather than only on the three listed pairs.
detected_version = install_utils._detected_cuda_version()
detected_major = detected_version[0] if detected_version is not None else None
detected = (
str(detected_major)
if detected_major is not None and str(detected_major) in _CUDA_RUNTIME_PACKAGES
else ""
)
if requested:
# A row that names a train has to be buildable for that train. Reported here
# rather than left to produce a mismatched wheel, because nothing downstream
# compares the two: the metadata comes from the row and the binaries come from
# the toolkit.
if detected and detected != requested:
raise RuntimeError(
f"this build targets CUDA {requested} (from "
f"{'CU_VERSION' if os.environ.get('CU_VERSION') else 'DESIRED_CUDA'}="
f"{raw!r}) but the installed toolkit is CUDA {detected}. The declared "
"runtime packages and the loader search paths come from the requested "
"train while the libraries are compiled by the installed one, so the "
"wheel would install and then fail to load. Install a matching toolkit "
"or build the row that matches this one."
)
# The check above compares majors, which two rows of the same major share. A row
# names its train down to the minor, so cu130 and cu132 both reduce to 13 and a
# cu132 row built against a 13.0 toolkit passed. The device code and the version
# in the wheel's local label both come from the row, so that wheel claims a
# toolkit it was not compiled by. Compared here rather than folded above so the
# message can name both numbers.
if detected_version is not None and digits in _CUDA_ROW_MINORS:
requested_pair = _CUDA_ROW_MINORS[digits]
if detected_version != requested_pair:
requested_text = f"{requested_pair[0]}.{requested_pair[1]}"
detected_text = f"{detected_version[0]}.{detected_version[1]}"
raise RuntimeError(
f"this build targets CUDA {requested_text} (from "
f"{'CU_VERSION' if os.environ.get('CU_VERSION') else 'DESIRED_CUDA'}="
f"{raw!r}) but the installed toolkit is CUDA {detected_text}. Those "
"share a major version, so the runtime dependency this wheel declares "
"is correct while the device code and the version recorded in its "
"local label are not. Install a matching toolkit or build the row "
"that matches this one."
)
return requested
if raw and not requested:
# A row named something this packaging does not recognise. Silently reporting the
# builder's toolkit instead contradicts "the row's field wins" and produced a
# wheel tagged for one train carrying another.
supported = ", ".join(
f"cu{major}{minor}"
for major, minor in install_utils.SUPPORTED_CUDA_VERSIONS
)
raise RuntimeError(
f"the release row requests CUDA {raw!r}, which is not a train this project "
f"supports ({supported}). Add it to SUPPORTED_CUDA_VERSIONS in install_utils "
"and to _CUDA_RUNTIME_PACKAGES and _CUDA_LIBRARY_DIRECTORIES here, or build "
"a supported row. Falling back to whatever toolkit this builder has would tag "
"the wheel for one train and fill it with another."
)
# Fall back to the installed toolkit, because keying this off a release variable alone produced a wheel
# that carried the CUDA libraries while declaring no CUDA runtime and recording no way to reach one.
#
# Two ways CUDA gets built, and both have to agree with what is declared here. The build gate turns it
# on when a SUPPORTED train is installed, so a toolkit whose minor is unlisted builds CPU-only and
# declaring runtime packages for it would make a CPU wheel demand four CUDA wheels. An explicit ON
# bypasses that gate and reaches CMake directly, where find_package(CUDAToolkit) accepts a toolkit
# this packaging does not list, so the libraries ship and the runtime has to be declared for them.
# Asking only whether the train is supported got the first case right and the second wrong.
explicit_on = install_utils.is_cmake_option_on(
_cmake_args(),
"EXECUTORCH_BUILD_CUDA",
default=False,
)
if not install_utils.is_cuda_available() and not explicit_on:
return ""
if explicit_on and not detected:
# The explicit request reaches CMake either way, so returning "" here shipped the
# CUDA libraries with no runtime declared and no path to one, which is the wheel
# this whole function exists to prevent. An unlisted minor still resolves to a
# major and is fine; an unlisted major has nothing to declare.
installed = (
f"CUDA {detected_major}"
if detected_major is not None
else "no CUDA toolkit"
)
raise RuntimeError(
f"the build asks for CUDA (EXECUTORCH_BUILD_CUDA=ON in CMAKE_ARGS) but found "
f"{installed}, and this packaging declares a runtime only for CUDA "
f"{', '.join(sorted(_CUDA_RUNTIME_PACKAGES))}. The libraries would still be "
"built and shipped with nothing to load them against. Install a toolkit on one "
"of those majors, or add this one to _CUDA_RUNTIME_PACKAGES and "
"_CUDA_LIBRARY_DIRECTORIES."
)
return detected
def _cuda_libraries_built(cmake_cache_dir: Optional[str]) -> bool:
"""Whether this build produced the CUDA libraries, read from the CMake cache.
The build turns CUDA on from the cache, so the cache is the fact that decides what ships. The
release row's CUDA version is a different question: a build on a toolkit whose train this packaging
does not recognise still produces the libraries while declaring no train, and gating anything else on
the train left that wheel carrying libraries with no matching header.
Falls back to the train when no cache is readable, which is the case for a source distribution where
nothing was built here anyway.
"""
cache_path = os.path.join(cmake_cache_dir or "", "CMakeCache.txt")
if os.path.exists(cache_path):
return CMakeCache(cache_path=cache_path).is_enabled("EXECUTORCH_BUILD_CUDA")
return bool(_cuda_train())
def _verify_cuda_runtime_matches_train(cmake_cache_dir: Optional[str]) -> None:
"""Fail the build when the linked CUDA runtime is not the train being declared.
The declared packages come from the compiler version, while the library that actually gets
linked comes from find_package(CUDAToolkit). Those are normally the same toolkit, but
CUDAToolkit_ROOT steers the second and not the first, so they can split inside a single
find_package call: measured with the compiler at 13.0 and that variable at 12.8,
CUDAToolkit_VERSION reported 13.0.88 while the binary needed libcudart.so.12. Packaging
would then declare the CUDA 13 runtime for a wheel that cannot load without CUDA 12.
Read from the CMake cache rather than from the environment, because the cache records what
the build resolved rather than what was requested.
"""
train = _cuda_train()
if not train:
return
cache_path = os.path.join(cmake_cache_dir or "", "CMakeCache.txt")
if not os.path.exists(cache_path):
return
cache = CMakeCache(cache_path=cache_path)
if not cache.is_enabled("EXECUTORCH_BUILD_CUDA"):
return
linked = cache.get("CUDA_cudart_LIBRARY")
if linked is None or not linked.value:
return
# Read the major from the resolved file name rather than from the recorded path, because the
# conventional way to name a toolkit is the versionless /usr/local/cuda symlink, which carries
# no version at all. Matching the directory accepted that spelling silently, which is the one
# the guard's own message tells the user to set. The resolved name ends in the soname the
# loader will ask for, which is the thing the declared package has to agree with.
found = re.search(r"libcudart\.so\.(\d+)", os.path.realpath(linked.value))
if found is None or found.group(1) == train:
return
raise RuntimeError(
f"this build declares the CUDA {train} runtime but linked the CUDA "
f"{found.group(1)} one from {linked.value!r}, so the wheel would install and then "
"fail to load. The declared train follows the CUDA compiler while the linked "
"libraries follow find_package(CUDAToolkit), so point CUDACXX and CUDAToolkit_ROOT "
"at the same toolkit."
)
def _cuda_dependencies() -> List[str]:
"""Runtime libraries a CUDA wheel needs but does not bundle.
Declared rather than vendored, the way the PyTorch CUDA wheels do it, so one copy is
shared with torch instead of shipping a second one.
"""
train = _cuda_train()
# Marked for Linux, because a CUDA wheel is only built there and these nvidia wheels publish no
# distribution for the other platforms, so an unmarked requirement would make a source install
# elsewhere fail on a dependency it cannot satisfy and does not need.
return [
f"{name}; platform_system == 'Linux'"
for name in _CUDA_RUNTIME_PACKAGES.get(train, ())
]
# Directories inside the wheel that hold libraries a shipped library links, relative to the package
# root rather than to the linking library, because the wheel ships libraries at more than one depth.
#
# The CUDA libraries are split across two directories and reference each other in both directions:
# the delegate in lib/ links the shims library in backends/cuda/, and the shims library links the
# stream helper back in lib/. So both hops are needed.
#
# Applied to every shipped library rather than mapping each library to the directories it happens to
# need. An unused hop costs nothing at load time, while a missing one produces a wheel that installs
# and then fails to load, and a per-library mapping would have to be revisited every time a library
# moves.
_SIBLING_LIBRARY_DIRECTORIES = ("backends/cuda", "lib", "src/executorch/lib")
def _sibling_library_search_paths(depth: int = 1) -> List[str]:
"""Loader paths that reach another directory inside this same package.
`depth` is how many directories separate the linking library from the package root, and it has to
be honoured for the same reason the CUDA hops honour it: the wheel ships libraries at depth one
(lib/) and depth two (backends/cuda/, extension/pybindings/ and others). Measured with a fixed
pair sized for one depth, six of twelve hops landed somewhere that does not exist, and the hop
from lib/ escaped the package entirely into a sibling of it, where an unrelated library with a
matching SONAME could satisfy the dependency first.
"""
up = "/".join([".."] * depth)
token = _loader_relative_token()
return [f"{token}/{up}/{directory}" for directory in _SIBLING_LIBRARY_DIRECTORIES]
def _loader_relative_token() -> str:
"""The token a runtime search path uses to mean "the directory this file is in".
ELF spells it $ORIGIN and Mach-O spells it @loader_path. Both are literal text in the
recorded path, so the wrong one becomes a directory of that name and resolves to
nothing.
"""
return "@loader_path" if sys.platform == "darwin" else "$ORIGIN"