Add minimal wheel build mode

.ci/scripts/test_minimal_wheel.sh

@@ -0,0 +1,111 @@
+#!/bin/bash
+# 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.
+
+set -euxo pipefail
+
+PYTHON_EXECUTABLE="${PYTHON_EXECUTABLE:-python}"
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+BUILD_VENV="${REPO_ROOT}/.venv-minimal-build"
+TEST_VENV="${REPO_ROOT}/.venv-minimal-test"
+
+rm -rf "${BUILD_VENV}" "${TEST_VENV}" "${REPO_ROOT}/dist" "${REPO_ROOT}/pip-out"
+
+"${PYTHON_EXECUTABLE}" -m venv "${BUILD_VENV}"
+source "${BUILD_VENV}/bin/activate"
+python -m pip install --upgrade pip
+python -m pip install \
+  "cmake>=3.24,<4.0.0" \
+  "numpy>=2.0.0" \
+  packaging \
+  pyyaml \
+  setuptools \
+  wheel \
+  zstd \
+  certifi \
+  torch \
+  torchvision \
+  --index-url https://download.pytorch.org/whl/cpu \
+  --extra-index-url https://pypi.org/simple
+
+(
+  cd "${REPO_ROOT}"
+  EXECUTORCH_BUILD_MINIMAL=1 python setup.py bdist_wheel
+)
+
+WHEEL_FILE="$(find "${REPO_ROOT}/dist" -maxdepth 1 -name 'executorch-*.whl' | head -1)"
+test -n "${WHEEL_FILE}"
+
+python - "${WHEEL_FILE}" <<'PY'
+import sys
+import zipfile
+
+wheel_file = sys.argv[1]
+with zipfile.ZipFile(wheel_file) as wheel:
+    names = wheel.namelist()
+
+for forbidden in (
+    "executorch/backends/",
+    "executorch/examples/",
+    "executorch/kernels/",
+    "executorch/runtime/",
+    "executorch/devtools/",
+    "executorch/extension/pybindings/",
+):
+    matches = [name for name in names if name.startswith(forbidden)]
+    if matches:
+        raise AssertionError(f"{wheel_file} unexpectedly contains {matches[:5]}")
+
+extensions = [
+    name
+    for name in names
+    if name.endswith((".so", ".dylib", ".dll", ".pyd")) and "flatc" not in name
+]
+if extensions:
+    raise AssertionError(f"{wheel_file} unexpectedly contains extensions: {extensions}")
+PY
+
+deactivate
+
+"${PYTHON_EXECUTABLE}" -m venv "${TEST_VENV}"
+source "${TEST_VENV}/bin/activate"
+python -m pip install --upgrade pip
+python -m pip install \
+  "flatbuffers" \
+  "numpy>=2.0.0" \
+  "packaging" \
+  "pyyaml" \
+  "ruamel.yaml" \
+  "sympy" \
+  "tabulate" \
+  "torch" \
+  "torchvision" \
+  "typing-extensions>=4.10.0" \
+  --index-url https://download.pytorch.org/whl/cpu \
+  --extra-index-url https://pypi.org/simple
+python -m pip install --no-deps "${WHEEL_FILE}"
+
+python - <<'PY'
+from pathlib import Path
+
+import torch
+from torch.export import export
+from torchvision.models import mobilenet_v2
+
+from executorch.exir import to_edge_transform_and_lower
+
+model = mobilenet_v2(weights=None).eval()
+example_inputs = (torch.randn(1, 3, 224, 224),)
+
+edge_program = to_edge_transform_and_lower(export(model, example_inputs))
+executorch_program = edge_program.to_executorch()
+
+output_path = Path("mv2_minimal.pte")
+with output_path.open("wb") as output_file:
+    executorch_program.write_to_file(output_file)
+
+assert output_path.stat().st_size > 0
+PY

.github/workflows/pull.yml

@@ -48,6 +48,36 @@ jobs:
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/test_wheel_package_qnn.sh "${{ matrix.python-version }}"
 
+  test-minimal-wheel-linux:
+    needs: changed-files
+    if: |
+      github.event_name != 'pull_request' ||
+      contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_minimal_wheel.sh') ||
+      contains(needs.changed-files.outputs.changed-files, '.github/workflows/pull.yml') ||
+      contains(needs.changed-files.outputs.changed-files, 'exir/') ||
+      contains(needs.changed-files.outputs.changed-files, 'extension/flat_tensor') ||
+      contains(needs.changed-files.outputs.changed-files, 'extension/pytree') ||
+      contains(needs.changed-files.outputs.changed-files, 'pyproject.toml') ||
+      contains(needs.changed-files.outputs.changed-files, 'schema/') ||
+      contains(needs.changed-files.outputs.changed-files, 'setup.py') ||
+      contains(needs.changed-files.outputs.changed-files, 'tools/cmake/')
+    name: test-minimal-wheel-linux
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    permissions:
+      id-token: write
+      contents: read
+    with:
+      runner: linux.2xlarge
+      docker-image: ci-image:executorch-ubuntu-22.04-clang12
+      submodules: 'recursive'
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      timeout: 120
+      script: |
+        CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
+        conda activate "${CONDA_ENV}"
+
+        PYTHON_EXECUTABLE=python bash .ci/scripts/test_minimal_wheel.sh
+
   test-setup-linux-gcc:
     name: test-setup-linux-gcc
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main

README-wheel.md

@@ -8,6 +8,12 @@ The `executorch` pip package is in beta.
 * Supported python versions: 3.10, 3.11, 3.12, 3.13
 * Compatible systems: Linux x86_64, Linux aarch64, macOS aarch64
 
+To build a minimal wheel from source, set
+`EXECUTORCH_BUILD_MINIMAL=1` when running `pip wheel` or `pip install`.
+That wheel contains the Python EXIR export path and `flatc` for `.pte`
+serialization, but omits runtime pybindings, kernels, backend packages, headers,
+examples, and devtools.
+
 The prebuilt `executorch.runtime` module included in this package provides a way
 to run ExecuTorch `.pte` files, with some restrictions:
 * Only [core ATen operators](docs/source/ir-ops-set-definition.md) are linked into the prebuilt module

exir/_serialize/_flatbuffer.py

@@ -269,27 +269,35 @@ def _get_flatc_path() -> str:
         if _flatc_cached_path is not None:
             return _flatc_cached_path
 
-        flatc_resource = importlib.resources.files(__package__).joinpath(
-            _FLATC_RESOURCE_NAME
-        )
-        if flatc_resource.is_file():
-            exit_stack = contextlib.ExitStack()
-            flatc_path = exit_stack.enter_context(
-                importlib.resources.as_file(flatc_resource)
-            )
+        for package, resource_name in (
+            (__package__, _FLATC_RESOURCE_NAME),
+            ("executorch.data.bin", "flatc"),
+        ):
             try:
-                current_mode = flatc_path.stat().st_mode
-                if not (current_mode & stat.S_IXUSR):
-                    flatc_path.chmod(
-                        current_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
-                    )
-            except OSError:
-                pass
-            _flatc_exit_stack = exit_stack
-            # Clean up the extracted temp file on normal process exit.
-            atexit.register(exit_stack.close)
-            _flatc_cached_path = str(flatc_path)
-        else:
+                flatc_resource = importlib.resources.files(package).joinpath(
+                    resource_name
+                )
+            except ModuleNotFoundError:
+                continue
+            if flatc_resource.is_file():
+                exit_stack = contextlib.ExitStack()
+                flatc_path = exit_stack.enter_context(
+                    importlib.resources.as_file(flatc_resource)
+                )
+                try:
+                    current_mode = flatc_path.stat().st_mode
+                    if not (current_mode & stat.S_IXUSR):
+                        flatc_path.chmod(
+                            current_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
+                        )
+                except OSError:
+                    pass
+                _flatc_exit_stack = exit_stack
+                # Clean up the extracted temp file on normal process exit.
+                atexit.register(exit_stack.close)
+                _flatc_cached_path = str(flatc_path)
+                break
+        if _flatc_cached_path is None:
             _flatc_cached_path = os.getenv("FLATC_EXECUTABLE", "flatc")
 
         return _flatc_cached_path