[build] Properly implement editable mode (#8722)

* [build] Properly implement editable mode

* Add CI

* Fix typo

* Fix macos job

Author: larryliu0820

.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
 

.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

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])

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",

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.

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,