Skip to content

Commit 0d904b6

Browse files
Add minimal wheel build mode (pytorch#19899)
Adds an opt-in `EXECUTORCH_BUILD_MINIMAL=1` wheel build mode that packages only the Python EXIR export path (plus `flatc`), for distributors that need ExecuTorch's ahead-of-time `.pte` export but not its runtime. For example, Torch-TensorRT's `output_format="executorch"` uses ExecuTorch only to export and runs the result with its own runtime. The minimal wheel: - omits runtime pybindings, kernels, backend packages, headers, examples, and devtools; - declares only the dependencies the export path needs (`flatbuffers`, `numpy`, `packaging`, `pyyaml`, `ruamel.yaml`, `sympy`, `tabulate`, `typing-extensions`) instead of the full set (`coremltools`, `scikit-learn`, `pandas`, `hydra-core`, `omegaconf`, and so on), so a normal install stays small; - produces byte-identical `.pte` output to the full wheel. The default (non-minimal) wheel is unchanged: its dependencies move from static `pyproject.toml` to a dynamic `install_requires` in `setup.py`, but the declared set is identical. Build from source and bundle it: ``` EXECUTORCH_BUILD_MINIMAL=1 pip wheel . --no-deps # or: EXECUTORCH_BUILD_MINIMAL=1 pip install . ``` A redistributor (e.g. NVIDIA, for a Torch-TensorRT container) can build the slim wheel at a pinned ExecuTorch version and ship it. `torch` is consumer-provided in both modes. CI (`test-minimal-wheel-linux`) builds the minimal wheel, asserts the excluded runtime/backend content and heavy deps are absent, installs it in a clean venv with full dependency resolution (no `--no-deps`), runs the bundled `flatc`, and exports MobileNetV2 to a `.pte`. Local result: minimal Linux x86_64 wheel ~2.1 MiB compressed vs ~15 MiB for the full wheel; MobileNetV2 `.pte` is 13,995,880 bytes, byte-identical to the published 1.3.1 wheel. --------- Co-authored-by: shoumikhin <anthony@shoumikh.in>
1 parent f976e63 commit 0d904b6

6 files changed

Lines changed: 571 additions & 201 deletions

File tree

.ci/scripts/test_minimal_wheel.sh

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
#!/bin/bash
2+
# Copyright (c) Meta Platforms, Inc. and affiliates.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
8+
set -euxo pipefail
9+
10+
PYTHON_EXECUTABLE="${PYTHON_EXECUTABLE:-python}"
11+
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
12+
BUILD_VENV="${REPO_ROOT}/.venv-minimal-build"
13+
TEST_VENV="${REPO_ROOT}/.venv-minimal-test"
14+
15+
rm -rf "${BUILD_VENV}" "${TEST_VENV}" "${REPO_ROOT}/dist" "${REPO_ROOT}/pip-out"
16+
17+
"${PYTHON_EXECUTABLE}" -m venv "${BUILD_VENV}"
18+
source "${BUILD_VENV}/bin/activate"
19+
python -m pip install --upgrade pip
20+
python -m pip install \
21+
"cmake>=3.24,<4.0.0" \
22+
"numpy>=2.0.0" \
23+
packaging \
24+
pyyaml \
25+
setuptools \
26+
wheel \
27+
zstd \
28+
certifi \
29+
torch \
30+
torchvision \
31+
--index-url https://download.pytorch.org/whl/cpu \
32+
--extra-index-url https://pypi.org/simple
33+
34+
(
35+
cd "${REPO_ROOT}"
36+
EXECUTORCH_BUILD_MINIMAL=1 python setup.py bdist_wheel
37+
)
38+
39+
WHEEL_FILE="$(find "${REPO_ROOT}/dist" -maxdepth 1 -name 'executorch-*.whl' | head -1)"
40+
test -n "${WHEEL_FILE}"
41+
42+
python - "${WHEEL_FILE}" <<'PY'
43+
import re
44+
import sys
45+
import zipfile
46+
47+
wheel_file = sys.argv[1]
48+
with zipfile.ZipFile(wheel_file) as wheel:
49+
names = wheel.namelist()
50+
metadata_name = next(
51+
(name for name in names if name.endswith(".dist-info/METADATA")), None
52+
)
53+
if metadata_name is None:
54+
raise AssertionError(f"{wheel_file} has no METADATA")
55+
metadata_text = wheel.read(metadata_name).decode("utf-8")
56+
57+
for forbidden in (
58+
"executorch/backends/",
59+
"executorch/examples/",
60+
"executorch/kernels/",
61+
"executorch/runtime/",
62+
"executorch/devtools/",
63+
"executorch/extension/pybindings/",
64+
):
65+
matches = [name for name in names if name.startswith(forbidden)]
66+
if matches:
67+
raise AssertionError(f"{wheel_file} unexpectedly contains {matches[:5]}")
68+
69+
extensions = [
70+
name
71+
for name in names
72+
if name.endswith((".so", ".dylib", ".dll", ".pyd")) and "flatc" not in name
73+
]
74+
if extensions:
75+
raise AssertionError(f"{wheel_file} unexpectedly contains extensions: {extensions}")
76+
77+
78+
def _dist_name(requirement):
79+
name = re.split(r"[ ;\[<>=!~(]", requirement.strip(), maxsplit=1)[0]
80+
return re.sub(r"[-_.]+", "-", name).lower()
81+
82+
83+
# Only the core (non-extra) Requires-Dist entries define what a plain
84+
# "pip install" pulls; ignore the optional extras (cortex_m, vgf, ...).
85+
declared = {
86+
_dist_name(line.split(":", 1)[1])
87+
for line in metadata_text.splitlines()
88+
if line.startswith("Requires-Dist:") and "extra==" not in line.replace(" ", "")
89+
}
90+
# The minimal wheel must declare EXACTLY this core set and nothing else -- the
91+
# same names as `keep` in setup.py:_minimal_dependencies(). Exact match catches
92+
# both a heavy full-wheel dep leaking in (coremltools, pandas, or a re-added
93+
# mpmath/torch) and a required dep going missing.
94+
expected = {
95+
"flatbuffers",
96+
"numpy",
97+
"packaging",
98+
"pyyaml",
99+
"ruamel-yaml",
100+
"sympy",
101+
"tabulate",
102+
"typing-extensions",
103+
}
104+
if declared != expected:
105+
raise AssertionError(
106+
f"{wheel_file} minimal core deps mismatch: "
107+
f"unexpected={sorted(declared - expected)} missing={sorted(expected - declared)}"
108+
)
109+
PY
110+
111+
deactivate
112+
113+
"${PYTHON_EXECUTABLE}" -m venv "${TEST_VENV}"
114+
source "${TEST_VENV}/bin/activate"
115+
python -m pip install --upgrade pip
116+
# torch and torchvision are needed to export a model but are intentionally not
117+
# declared as wheel dependencies (consumers are expected to bring their own).
118+
python -m pip install \
119+
"torch" \
120+
"torchvision" \
121+
--index-url https://download.pytorch.org/whl/cpu \
122+
--extra-index-url https://pypi.org/simple
123+
# Install the minimal wheel WITHOUT --no-deps so pip resolves its declared
124+
# dependencies, confirming the slim set is correct and resolvable. (That no heavy
125+
# deps sneak in is guaranteed by the METADATA exact-match check above, which
126+
# covers the wheel's direct Requires-Dist.)
127+
python -m pip install \
128+
"${WHEEL_FILE}" \
129+
--index-url https://download.pytorch.org/whl/cpu \
130+
--extra-index-url https://pypi.org/simple
131+
132+
# flatc is the only compiled artifact in the minimal wheel and the reason it is
133+
# platform specific. Confirm it ships, resolves through _get_flatc_path() (the
134+
# executorch.data.bin lookup added for this build mode), and actually runs.
135+
python - <<'PY'
136+
import subprocess
137+
138+
from executorch.exir._serialize._flatbuffer import _get_flatc_path
139+
140+
flatc_path = _get_flatc_path()
141+
print(f"flatc resolved to: {flatc_path}")
142+
subprocess.run([flatc_path, "--version"], check=True)
143+
PY
144+
145+
python - <<'PY'
146+
from pathlib import Path
147+
148+
import torch
149+
from torch.export import export
150+
from torchvision.models import mobilenet_v2
151+
152+
from executorch.exir import to_edge_transform_and_lower
153+
154+
model = mobilenet_v2(weights=None).eval()
155+
example_inputs = (torch.randn(1, 3, 224, 224),)
156+
157+
edge_program = to_edge_transform_and_lower(export(model, example_inputs))
158+
executorch_program = edge_program.to_executorch()
159+
160+
output_path = Path("mv2_minimal.pte")
161+
with output_path.open("wb") as output_file:
162+
executorch_program.write_to_file(output_file)
163+
164+
assert output_path.stat().st_size > 0
165+
PY

.github/workflows/pull.yml

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,36 @@ jobs:
6161
6262
PYTHON_EXECUTABLE=python bash .ci/scripts/test_wheel_package_qnn.sh "${{ matrix.python-version }}"
6363
64+
test-minimal-wheel-linux:
65+
needs: changed-files
66+
if: |
67+
github.event_name != 'pull_request' ||
68+
contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_minimal_wheel.sh') ||
69+
contains(needs.changed-files.outputs.changed-files, '.github/workflows/pull.yml') ||
70+
contains(needs.changed-files.outputs.changed-files, 'exir/') ||
71+
contains(needs.changed-files.outputs.changed-files, 'extension/flat_tensor') ||
72+
contains(needs.changed-files.outputs.changed-files, 'extension/pytree') ||
73+
contains(needs.changed-files.outputs.changed-files, 'pyproject.toml') ||
74+
contains(needs.changed-files.outputs.changed-files, 'schema/') ||
75+
contains(needs.changed-files.outputs.changed-files, 'setup.py') ||
76+
contains(needs.changed-files.outputs.changed-files, 'tools/cmake/')
77+
name: test-minimal-wheel-linux
78+
uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
79+
permissions:
80+
id-token: write
81+
contents: read
82+
with:
83+
runner: linux.2xlarge
84+
docker-image: ci-image:executorch-ubuntu-22.04-clang12
85+
submodules: 'recursive'
86+
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
87+
timeout: 120
88+
script: |
89+
CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
90+
conda activate "${CONDA_ENV}"
91+
92+
PYTHON_EXECUTABLE=python bash .ci/scripts/test_minimal_wheel.sh
93+
6494
test-setup-linux-gcc:
6595
name: test-setup-linux-gcc
6696
uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main

README-wheel.md

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,16 @@ The `executorch` pip package is in beta.
88
* Supported python versions: 3.10, 3.11, 3.12, 3.13
99
* Compatible systems: Linux x86_64, Linux aarch64, macOS aarch64
1010

11+
To build a minimal wheel from source, set
12+
`EXECUTORCH_BUILD_MINIMAL=1` when running `pip wheel` or `pip install`.
13+
That wheel contains the Python EXIR export path and `flatc` for `.pte`
14+
serialization, but omits runtime pybindings, kernels, backend packages, headers,
15+
examples, and devtools. It also declares only the Python dependencies the export
16+
path needs (no `coremltools`, `pandas`, `scikit-learn`, `hydra-core`, or
17+
`omegaconf`), so a normal install stays small. Like the full wheel it does not
18+
bundle PyTorch, so install a compatible `torch` separately. The wheel is still
19+
platform specific because it ships `flatc`.
20+
1121
The prebuilt `executorch.runtime` module included in this package provides a way
1222
to run ExecuTorch `.pte` files, with some restrictions:
1323
* Only [core ATen operators](docs/source/ir-ops-set-definition.md) are linked into the prebuilt module

exir/_serialize/_flatbuffer.py

Lines changed: 28 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -268,27 +268,35 @@ def _get_flatc_path() -> str:
268268
if _flatc_cached_path is not None:
269269
return _flatc_cached_path
270270

271-
flatc_resource = importlib.resources.files(__package__).joinpath(
272-
_FLATC_RESOURCE_NAME
273-
)
274-
if flatc_resource.is_file():
275-
exit_stack = contextlib.ExitStack()
276-
flatc_path = exit_stack.enter_context(
277-
importlib.resources.as_file(flatc_resource)
278-
)
271+
for package, resource_name in (
272+
(__package__, _FLATC_RESOURCE_NAME),
273+
("executorch.data.bin", "flatc"),
274+
):
279275
try:
280-
current_mode = flatc_path.stat().st_mode
281-
if not (current_mode & stat.S_IXUSR):
282-
flatc_path.chmod(
283-
current_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
284-
)
285-
except OSError:
286-
pass
287-
_flatc_exit_stack = exit_stack
288-
# Clean up the extracted temp file on normal process exit.
289-
atexit.register(exit_stack.close)
290-
_flatc_cached_path = str(flatc_path)
291-
else:
276+
flatc_resource = importlib.resources.files(package).joinpath(
277+
resource_name
278+
)
279+
except ModuleNotFoundError:
280+
continue
281+
if flatc_resource.is_file():
282+
exit_stack = contextlib.ExitStack()
283+
flatc_path = exit_stack.enter_context(
284+
importlib.resources.as_file(flatc_resource)
285+
)
286+
try:
287+
current_mode = flatc_path.stat().st_mode
288+
if not (current_mode & stat.S_IXUSR):
289+
flatc_path.chmod(
290+
current_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
291+
)
292+
except OSError:
293+
pass
294+
_flatc_exit_stack = exit_stack
295+
# Clean up the extracted temp file on normal process exit.
296+
atexit.register(exit_stack.close)
297+
_flatc_cached_path = str(flatc_path)
298+
break
299+
if _flatc_cached_path is None:
292300
_flatc_cached_path = os.getenv("FLATC_EXECUTABLE", "flatc")
293301

294302
return _flatc_cached_path

pyproject.toml

Lines changed: 8 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,9 @@ name = "executorch"
1616
dynamic = [
1717
# setup.py will set the version.
1818
'version',
19+
# setup.py sets dependencies, which vary by build mode (the
20+
# EXECUTORCH_BUILD_MINIMAL wheel declares a slimmer runtime set).
21+
'dependencies',
1922
]
2023
description = "On-device AI across mobile, embedded and edge for PyTorch"
2124
readme = "README-wheel.md"
@@ -51,30 +54,11 @@ classifiers = [
5154
]
5255

5356
requires-python = ">=3.10,<3.14"
54-
dependencies=[
55-
"expecttest",
56-
"flatbuffers",
57-
"hypothesis",
58-
"kgb",
59-
"mpmath==1.3.0",
60-
"numpy>=2.0.0; python_version >= '3.10'",
61-
"packaging",
62-
"pandas>=2.2.2; python_version >= '3.10'",
63-
"parameterized",
64-
"pytorch-tokenizers",
65-
"pyyaml",
66-
"ruamel.yaml",
67-
"sympy",
68-
"tabulate",
69-
# See also third-party/TARGETS for buck's typing-extensions version.
70-
"typing-extensions>=4.10.0",
71-
# Keep this version in sync with: ./backends/apple/coreml/scripts/install_requirements.sh
72-
"coremltools==9.0; platform_system == 'Darwin' or platform_system == 'Linux'",
73-
# scikit-learn is used to support palettization in the coreml backend
74-
"scikit-learn==1.7.1",
75-
"hydra-core>=1.3.0",
76-
"omegaconf>=2.3.0",
77-
]
57+
58+
# Runtime dependencies are declared dynamically (see `dynamic` above) and
59+
# computed in setup.py, so the EXECUTORCH_BUILD_MINIMAL wheel can ship a slimmer
60+
# set than the full wheel. See `_base_dependencies()` / `_minimal_dependencies()`
61+
# in setup.py.
7862

7963
[project.optional-dependencies]
8064
cortex_m = [

0 commit comments

Comments
 (0)