[build] Fix flatc (#8858)

[build] Fix flatc (#8816)

* [build] Fix flatc

We need to install `build/pip_data_bin_init.py.in` into `<executorch
root>/data/bin/__init__.py`. This PR rewrite the logic into a
`BuiltFile` so that it works well in editable mode.

Test:

```
python -c "from executorch.data.bin import flatc"
```
Will add unit test in next PR.

[ghstack-poisoned]

* Update on "[build] Fix flatc"

Fixes #8784

We need to install `build/pip_data_bin_init.py.in` into `<executorch
root>/data/bin/__init__.py`. This PR rewrite the logic into a
`BuiltFile` so that it works well in editable mode.

Since `BuiltFile` by default looks into cmake cache directory, this PR adds a placeholder `%CMAKE_CACHE_DIR%` for those are actually built by CMake and for `build/pip_data_bin_init.py.in` we don't add this placeholder.

Test:

```
python -c "from executorch.data.bin import flatc"
```
Will add unit test in next PR.

[ghstack-poisoned]

* Update on "[build] Fix flatc"

Fixes #8784

We need to install `build/pip_data_bin_init.py.in` into `<executorch
root>/data/bin/__init__.py`. This PR rewrite the logic into a
`BuiltFile` so that it works well in editable mode.

Since `BuiltFile` by default looks into cmake cache directory, this PR adds a placeholder `%CMAKE_CACHE_DIR%` for those are actually built by CMake and for `build/pip_data_bin_init.py.in` we don't add this placeholder.

Test:

```
python -c "from executorch.data.bin import flatc"
```
Will add unit test in next PR.

[ghstack-poisoned]

* Update on "[build] Fix flatc"

Fixes #8784

We need to install `build/pip_data_bin_init.py.in` into `<executorch
root>/data/bin/__init__.py`. This PR rewrite the logic into a
`BuiltFile` so that it works well in editable mode.

Since `BuiltFile` by default looks into cmake cache directory, this PR adds a placeholder `%CMAKE_CACHE_DIR%` for those are actually built by CMake and for `build/pip_data_bin_init.py.in` we don't add this placeholder.

Test:

```
python -c "from executorch.data.bin import flatc"
```
Will add unit test in next PR.

[ghstack-poisoned]

(cherry picked from commit d0b27b5)

Co-authored-by: Mengwei Liu <[email protected]>

Author: pytorchbot

.github/workflows/pull.yml

@@ -371,6 +371,17 @@ jobs:
       build-tool: cmake
       docker-image: executorch-ubuntu-22.04-clang12
 
+  unittest-editable:
+    uses: ./.github/workflows/_unittest.yml
+    permissions:
+      id-token: write
+      contents: read
+    with:
+      build-mode: Debug
+      build-tool: cmake
+      editable: true
+      docker-image: executorch-ubuntu-22.04-clang12
+
   unittest-buck:
     uses: ./.github/workflows/_unittest.yml
     permissions:

data/bin/README.md

@@ -0,0 +1,31 @@
+## PLEASE DO NOT REMOVE THIS DIRECTORY!
+
+This directory is used to host binaries installed during pip wheel build time.
+
+## How to add a binary into pip wheel
+
+1. Update `[project.scripts]` section of `pyproject.toml` file. Add the new binary name and it's corresponding module name similar to:
+
+```
+flatc = "executorch.data.bin:flatc"
+```
+
+For example, `flatc` is built during wheel packaging, we first build `flatc` through CMake and copy the file to `<executorch root>/data/bin/flatc` and ask `setuptools` to generate a commandline wrapper for `flatc`, then route it to `<executorch root>/data/bin/flatc`.
+
+This way after installing `executorch`, a user will be able to call `flatc` directly in commandline and it points to `<executorch root>/data/bin/flatc`
+
+2. Update `setup.py` to include the logic of building the new binary and copying the binary to this directory.
+
+```python
+BuiltFile(
+    src_dir="%CMAKE_CACHE_DIR%/third-party/flatbuffers/%BUILD_TYPE%/",
+    src_name="flatc",
+    dst="executorch/data/bin/",
+    is_executable=True,
+),
+```
+This means find `flatc` in `CMAKE_CACHE_DIR` and copy it to `<executorch root>/data/bin`. Notice that this works for both pip wheel packaging as well as editable mode install.
+
+## Why we can't create this directory at wheel build time?
+
+The reason is without `data/bin/` present in source file, we can't tell `setuptools` to generate a module `executorch.data.bin` in editable mode, partially because we don't have a good top level module `executorch` and have to enumerate all the second level modules, including `executorch.data.bin`.

pyproject.toml

@@ -88,6 +88,7 @@ flatc = "executorch.data.bin:flatc"
 [tool.setuptools.package-dir]
 "executorch.backends" = "backends"
 "executorch.codegen" = "codegen"
+"executorch.data.bin" = "data/bin"
 # 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.

setup.py

@@ -220,37 +220,50 @@ def src_path(self, installer: "InstallerBuildExt") -> Path:
         """
         # Share the cmake-out location with CustomBuild.
         build_cmd = installer.get_finalized_command("build")
-        if hasattr(build_cmd, "cmake_cache_dir"):
-            cmake_cache_dir = Path(build_cmd.cmake_cache_dir)
+        if "%CMAKE_CACHE_DIR%" in self.src:
+            if not hasattr(build_cmd, "cmake_cache_dir"):
+                raise RuntimeError(
+                    f"Extension {self.name} has a src {self.src} that contains"
+                    " %CMAKE_CACHE_DIR% but CMake does not run in the `build` "
+                    "command. Please double check if the command is correct."
+                )
+            else:
+                build_dir = Path(build_cmd.cmake_cache_dir)
         else:
-            # If we're in editable mode, use a default or fallback value for cmake_cache_dir
-            # This could be a hardcoded path, or a path derived from the current working directory
-            cmake_cache_dir = Path(".")
+            # If the src path doesn't contain %CMAKE_CACHE_DIR% placeholder,
+            # try to find it under the current directory.
+            build_dir = Path(".")
+
+        src_path = self.src.replace("%CMAKE_CACHE_DIR%/", "")
+
         cfg = get_build_type(installer.debug)
 
         if os.name == "nt":
             # Replace %BUILD_TYPE% with the current build type.
-            self.src = self.src.replace("%BUILD_TYPE%", cfg)
+            src_path = src_path.replace("%BUILD_TYPE%", cfg)
         else:
             # Remove %BUILD_TYPE% from the path.
-            self.src = self.src.replace("/%BUILD_TYPE%", "")
+            src_path = src_path.replace("/%BUILD_TYPE%", "")
 
         # Construct the full source path, resolving globs. If there are no glob
         # pattern characters, this will just ensure that the source file exists.
-        srcs = tuple(cmake_cache_dir.glob(self.src))
+        srcs = tuple(build_dir.glob(src_path))
         if len(srcs) != 1:
             raise ValueError(
-                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.
-
-Try:
-
-EXECUTORCH_BUILD_FLATC=OFF EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT=OFF pip install -e .
-"""
+                f"Expecting exactly 1 file matching {self.src} in {build_dir}, "
+                f"found {repr(srcs)}. Resolved src pattern: {src_path}."
             )
         return srcs[0]
 
+    def inplace_dir(self, installer: "InstallerBuildExt") -> Path:
+        """Returns the path of this file to be installed to, under inplace mode.
+
+        It will be a relative path to the project root directory. For more info
+        related to inplace/editable mode, please checkout this doc:
+        https://setuptools.pypa.io/en/latest/userguide/development_mode.html
+        """
+        raise NotImplementedError()
+
 
 class BuiltFile(_BaseExtension):
     """An extension that installs a single file that was built by cmake.
@@ -316,6 +329,18 @@ def dst_path(self, installer: "InstallerBuildExt") -> Path:
             # Destination looks like a file.
             return dst_root / Path(self.dst)
 
+    def inplace_dir(self, installer: "InstallerBuildExt") -> Path:
+        """For a `BuiltFile`, we use self.dst as its inplace directory path.
+        Need to handle directory vs file.
+        """
+        # HACK: get rid of the leading "executorch" in ext.dst.
+        # This is because we don't have a root level "executorch" module.
+        package_dir = self.dst.removeprefix("executorch/")
+        # If dst is a file, use it's directory
+        if not package_dir.endswith("/"):
+            package_dir = os.path.dirname(package_dir)
+        return Path(package_dir)
+
 
 class BuiltExtension(_BaseExtension):
     """An extension that installs a python extension that was built by cmake."""
@@ -335,7 +360,7 @@ def __init__(self, src: str, modpath: str):
             "/" not in modpath
         ), f"modpath must be a dotted python module path: saw '{modpath}'"
         # This is a real extension, so use the modpath as the name.
-        super().__init__(src=src, dst=modpath, name=modpath)
+        super().__init__(src=f"%CMAKE_CACHE_DIR%/{src}", dst=modpath, name=modpath)
 
     def src_path(self, installer: "InstallerBuildExt") -> Path:
         """Returns the path to the source file, resolving globs.
@@ -369,6 +394,15 @@ def dst_path(self, installer: "InstallerBuildExt") -> Path:
         # path: that's the file we're creating.
         return Path(installer.get_ext_fullpath(self.dst))
 
+    def inplace_dir(self, installer: "InstallerBuildExt") -> Path:
+        """For BuiltExtension, deduce inplace dir path from extension name."""
+        build_py = installer.get_finalized_command("build_py")
+        modpath = self.name.split(".")
+        package = ".".join(modpath[:-1])
+        package_dir = os.path.abspath(build_py.get_package_dir(package))
+
+        return Path(package_dir)
+
 
 class InstallerBuildExt(build_ext):
     """Installs files that were built by cmake."""
@@ -399,23 +433,15 @@ def copy_extensions_to_source(self) -> None:
 
         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/")
+            package_dir = ext.inplace_dir(self)
 
             # 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))
+                package_dir, os.path.basename(ext.dst_path(self))
             )
 
             # Always copy, even if source is older than destination, to ensure
@@ -724,20 +750,6 @@ def run(self):
         # Build the system.
         self.spawn(["cmake", "--build", cmake_cache_dir, *build_args])
 
-        # Non-python files should live under this data directory.
-        data_root = os.path.join(self.build_lib, "executorch", "data")
-
-        # Directories like bin/ and lib/ live under data/.
-        bin_dir = os.path.join(data_root, "bin")
-
-        # Copy the bin wrapper so that users can run any executables under
-        # data/bin, as long as they are listed in the [project.scripts] section
-        # of pyproject.toml.
-        self.mkpath(bin_dir)
-        self.copy_file(
-            "build/pip_data_bin_init.py.in",
-            os.path.join(bin_dir, "__init__.py"),
-        )
         # Share the cmake-out location with _BaseExtension.
         self.cmake_cache_dir = cmake_cache_dir
 
@@ -749,13 +761,20 @@ def get_ext_modules() -> List[Extension]:
     """Returns the set of extension modules to build."""
     ext_modules = []
     if ShouldBuild.flatc():
-        ext_modules.append(
-            BuiltFile(
-                src_dir="third-party/flatbuffers/%BUILD_TYPE%/",
-                src_name="flatc",
-                dst="executorch/data/bin/",
-                is_executable=True,
-            )
+        ext_modules.extend(
+            [
+                BuiltFile(
+                    src_dir="%CMAKE_CACHE_DIR%/third-party/flatbuffers/%BUILD_TYPE%/",
+                    src_name="flatc",
+                    dst="executorch/data/bin/",
+                    is_executable=True,
+                ),
+                BuiltFile(
+                    src_dir="build/",
+                    src_name="pip_data_bin_init.py.in",
+                    dst="executorch/data/bin/__init__.py",
+                ),
+            ]
         )
 
     if ShouldBuild.pybindings():
@@ -778,16 +797,16 @@ def get_ext_modules() -> List[Extension]:
     if ShouldBuild.llama_custom_ops():
         ext_modules.append(
             BuiltFile(
-                src_dir="extension/llm/custom_ops/%BUILD_TYPE%/",
+                src_dir="%CMAKE_CACHE_DIR%/extension/llm/custom_ops/%BUILD_TYPE%/",
                 src_name="custom_ops_aot_lib",
-                dst="executorch/extension/llm/custom_ops",
+                dst="executorch/extension/llm/custom_ops/",
                 is_dynamic_lib=True,
             )
         )
         ext_modules.append(
             # Install the prebuilt library for quantized ops required by custom ops.
             BuiltFile(
-                src_dir="kernels/quantized/%BUILD_TYPE%/",
+                src_dir="%CMAKE_CACHE_DIR%/kernels/quantized/%BUILD_TYPE%/",
                 src_name="quantized_ops_aot_lib",
                 dst="executorch/kernels/quantized/",
                 is_dynamic_lib=True,