Add CI

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-mode --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-mode --pybind xnnpack --use-pt-pinned-commit
+        # 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,16 +22,14 @@
     op2 = torch.ops.llama.fast_hadamard_transform.default
     assert op2 is not None
 except:
-    import glob
-    import os
-
     # 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
     from pathlib import Path
+
     package_path = Path(__file__).parent.resolve()
     logging.info(f"Looking for libcustom_ops_aot_lib.so in {package_path}")
 

install_executorch.py

@@ -138,6 +138,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 +234,7 @@ def main(args):
             "-m",
             "pip",
             "install",
+            "-e" if args.editable_mode else "",
             ".",
             "--no-build-isolation",
             "-v",

setup.py

@@ -50,8 +50,8 @@
 import os
 import platform
 import re
-import sys
 import site
+import sys
 
 # Import this before distutils so that setuptools can intercept the distuils
 # imports.
@@ -372,12 +372,13 @@ 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 
+        # 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
@@ -387,10 +388,9 @@ def run(self):
             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 
+        """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.
@@ -399,11 +399,11 @@ def copy_extensions_to_source(self) -> None:
 
         Returns:
         """
-        build_py = self.get_finalized_command('build_py')
+        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])
+                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.
@@ -412,9 +412,11 @@ def copy_extensions_to_source(self) -> None:
 
             # 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)))
+            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
@@ -428,7 +430,6 @@ def copy_extensions_to_source(self) -> None:
                 # 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:
@@ -688,7 +689,9 @@ def run(self):
             # 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["PYTHONPATH"]
+        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