diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index b599f2fdc67..d9a7aa6bd5b 100644 --- a/.github/workflows/pull.yml +++ b/.github/workflows/pull.yml @@ -56,6 +56,30 @@ jobs: # Build and test ExecuTorch with the add model on portable backend. PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh "add" "${BUILD_TOOL}" "portable" + test-pip-install-editable-mode-linux: + name: test-pip-install-editable-mode-linux + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + permissions: + id-token: write + contents: read + strategy: + fail-fast: false + with: + runner: linux.2xlarge + docker-image: executorch-ubuntu-22.04-clang12 + submodules: 'true' + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 90 + script: | + # The generic Linux job chooses to use base env, not the one setup by the image + CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]") + conda activate "${CONDA_ENV}" + # Debug + which pip + PYTHON_EXECUTABLE=python bash ./install_executorch.sh --editable --pybind xnnpack --use-pt-pinned-commit + # Try to import extension library + python -c "from executorch.extension.llm.custom_ops import custom_ops" + test-models-linux: name: test-models-linux uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main @@ -480,7 +504,7 @@ jobs: # Setup install_requirements for llama PYTHON_EXECUTABLE=python bash examples/models/llama/install_requirements.sh - + # Test static llama weight sharing and accuracy PYTHON_EXECUTABLE=python bash .ci/scripts/test_qnn_static_llama.sh diff --git a/.github/workflows/trunk.yml b/.github/workflows/trunk.yml index 7f66474bdba..d8ec745b75c 100644 --- a/.github/workflows/trunk.yml +++ b/.github/workflows/trunk.yml @@ -36,6 +36,31 @@ jobs: PYTHONPATH="${PWD}" python .ci/scripts/gather_test_models.py --target-os macos --event "${GITHUB_EVENT_NAME}" + test-pip-install-editable-mode-macos: + name: test-pip-install-editable-mode-macos + uses: pytorch/test-infra/.github/workflows/macos_job.yml@main + permissions: + id-token: write + contents: read + strategy: + fail-fast: false + with: + runner: macos-m1-stable + python-version: '3.11' + submodules: 'true' + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 90 + script: | + # The generic Linux job chooses to use base env, not the one setup by the image + CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]") + conda activate "${CONDA_ENV}" + # Debug + which pip + bash .ci/scripts/setup-conda.sh + PYTHON_EXECUTABLE=python ${CONDA_RUN} bash ./install_executorch.sh --editable --pybind xnnpack + # Try to import extension library + python -c "from executorch.extension.llm.custom_ops import custom_ops" + test-models-macos: name: test-models-macos uses: pytorch/test-infra/.github/workflows/macos_job.yml@main diff --git a/extension/llm/custom_ops/custom_ops.py b/extension/llm/custom_ops/custom_ops.py index b3b05db68fb..d299b314816 100644 --- a/extension/llm/custom_ops/custom_ops.py +++ b/extension/llm/custom_ops/custom_ops.py @@ -22,23 +22,19 @@ op2 = torch.ops.llama.fast_hadamard_transform.default assert op2 is not None except: - import glob - - import executorch - # This is needed to ensure that custom ops are registered from executorch.extension.pybindings import portable_lib # noqa # usort: skip # Ideally package is installed in only one location but usage of # PYATHONPATH can result in multiple locations. # ATM this is mainly used in CI for qnn runner. Will need to revisit this - executorch_package_path = executorch.__path__[-1] - logging.info(f"Looking for libcustom_ops_aot_lib.so in {executorch_package_path}") - libs = list( - glob.glob( - f"{executorch_package_path}/**/libcustom_ops_aot_lib.*", recursive=True - ) - ) + from pathlib import Path + + package_path = Path(__file__).parent.resolve() + logging.info(f"Looking for libcustom_ops_aot_lib.so in {package_path}") + + libs = list(package_path.glob("**/libcustom_ops_aot_lib.*")) + assert len(libs) == 1, f"Expected 1 library but got {len(libs)}" logging.info(f"Loading custom ops library: {libs[0]}") torch.ops.load_library(libs[0]) diff --git a/install_executorch.py b/install_executorch.py index b35f5668eb2..0d82f0a05ca 100644 --- a/install_executorch.py +++ b/install_executorch.py @@ -65,6 +65,7 @@ def clean(): "prelude": "BUCK", "pthreadpool": "CMakeLists.txt", "pybind11": "CMakeLists.txt", + "shim": "BUCK", "XNNPACK": "CMakeLists.txt", } @@ -138,6 +139,14 @@ def build_args_parser() -> argparse.ArgumentParser: action="store_true", help="build from the pinned PyTorch commit instead of nightly", ) + parser.add_argument( + "--editable", + "-e", + action="store_true", + help="build an editable pip wheel, changes to python code will be " + "picked up without rebuilding the wheel. Extension libraries will be " + "installed inside the source tree.", + ) return parser @@ -226,6 +235,9 @@ def main(args): "-m", "pip", "install", + ] + + (["--editable"] if args.editable else []) + + [ ".", "--no-build-isolation", "-v", diff --git a/pyproject.toml b/pyproject.toml index fb4196d99bc..43b0a8c4daf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,6 +82,25 @@ 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.package-dir] +"executorch.backends" = "backends" +"executorch.codegen" = "codegen" +# TODO(mnachin T180504136): Do not put examples/models +# into core pip packages. Refactor out the necessary utils +# or core models files into a separate package. +"executorch.examples.models" = "examples/models" +"executorch.exir" = "exir" +"executorch.extension" = "extension" +"executorch.kernels.quantized" = "kernels/quantized" +"executorch.schema" = "schema" +"executorch.devtools" = "devtools" +"executorch.devtools.bundled_program" = "devtools/bundled_program" +"executorch.runtime" = "runtime" +"executorch.util" = "util" + [tool.setuptools.package-data] # TODO(dbort): Prune /test[s]/ dirs, /third-party/ dirs, yaml files that we # don't need. diff --git a/setup.py b/setup.py index 59b468f8c93..ad5c561a7b3 100644 --- a/setup.py +++ b/setup.py @@ -50,6 +50,7 @@ import os import platform import re +import site import sys # Import this before distutils so that setuptools can intercept the distuils @@ -239,7 +240,7 @@ def src_path(self, installer: "InstallerBuildExt") -> Path: srcs = tuple(cmake_cache_dir.glob(self.src)) if len(srcs) != 1: raise ValueError( - f"""Expected exactly one file matching '{self.src}'; found {repr(srcs)}. + f"""Expected exactly one file matching '{self.src}' in {cmake_cache_dir}; found {repr(srcs)}. If that file is a CMake-built extension module file, and we are installing in editable mode, please disable the corresponding build option since it's not supported yet. @@ -372,6 +373,63 @@ def dst_path(self, installer: "InstallerBuildExt") -> Path: class InstallerBuildExt(build_ext): """Installs files that were built by cmake.""" + def __init__(self, *args, **kwargs): + self._ran_build = False + super().__init__(*args, **kwargs) + + def run(self): + # Run the build command first in editable mode. Since `build` command + # will also trigger `build_ext` command, only run this once. + if self._ran_build: + return + + if self.editable_mode: + self._ran_build = True + self.run_command("build") + super().run() + + def copy_extensions_to_source(self) -> None: + """For each extension in `ext_modules`, we need to copy the extension + file from the build directory to the correct location in the local + directory. + + This should only be triggered when inplace mode (editable mode) is enabled. + + Args: + + Returns: + """ + build_py = self.get_finalized_command("build_py") + for ext in self.extensions: + if isinstance(ext, BuiltExtension): + modpath = ext.name.split(".") + package = ".".join(modpath[:-1]) + package_dir = os.path.abspath(build_py.get_package_dir(package)) + else: + # HACK: get rid of the leading "executorch" in ext.dst. + # This is because we don't have a root level "executorch" module. + package_dir = ext.dst.removeprefix("executorch/") + + # Ensure that the destination directory exists. + self.mkpath(os.fspath(package_dir)) + + regular_file = ext.src_path(self) + inplace_file = os.path.join( + package_dir, os.path.basename(ext.src_path(self)) + ) + + # Always copy, even if source is older than destination, to ensure + # that the right extensions for the current Python/platform are + # used. + if os.path.exists(regular_file) or not ext.optional: + self.copy_file(regular_file, inplace_file, level=self.verbose) + + if ext._needs_stub: + inplace_stub = self._get_equivalent_stub(ext, inplace_file) + self._write_stub_file(inplace_stub, ext, compile=True) + # Always compile stub and remove the original (leave the cache behind) + # (this behaviour was observed in previous iterations of the code) + # TODO(dbort): Depend on the "build" command to ensure it runs first def build_extension(self, ext: _BaseExtension) -> None: @@ -630,6 +688,10 @@ def run(self): if not self.dry_run: # Dry run should log the command but not actually run it. (Path(cmake_cache_dir) / "CMakeCache.txt").unlink(missing_ok=True) + # Set PYTHONPATH to the location of the pip package. + os.environ["PYTHONPATH"] = ( + site.getsitepackages()[0] + ";" + os.environ.get("PYTHONPATH", "") + ) with Buck2EnvironmentFixer(): # The context manager may patch the environment while running this # cmake command, which happens to run buck2 to get some source @@ -741,25 +803,6 @@ def get_ext_modules() -> List[Extension]: setup( version=Version.string(), - # 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. - package_dir={ - "executorch/backends": "backends", - "executorch/codegen": "codegen", - # TODO(mnachin T180504136): Do not put examples/models - # into core pip packages. Refactor out the necessary utils - # or core models files into a separate package. - "executorch/examples/models": "examples/models", - "executorch/exir": "exir", - "executorch/extension": "extension", - "executorch/kernels/quantized": "kernels/quantized", - "executorch/schema": "schema", - "executorch/devtools": "devtools", - "executorch/devtools/bundled_program": "devtools/bundled_program", - "executorch/runtime": "runtime", - "executorch/util": "util", - }, cmdclass={ "build": CustomBuild, "build_ext": InstallerBuildExt,