Fix 'cpp_lib.json' file creation and installation

ghstack-source-id: 88885d4d1d96ea7e452c7cbe46d6b24383be047b
Pull Request resolved: https://github.com/fairinternal/xformers/pull/426

__original_commit__ = fairinternal/xformers@c5776a283a60df6507f763fa01134cd799b80df0

.circleci/config.yml

@@ -126,6 +126,7 @@ install_repo: &install_repo
 
         # Test import.
         $CONDA_PYTHON -c 'import sys; sys.path = sys.path[1:]; import xformers'
+        ls xformers
         $CONDA_PYTHON -m xformers.info
 
 install_experimental_repo: &install_experimental_repo

setup.py

@@ -300,14 +300,29 @@ def run(self):
 class BuildExtensionWithMetadata(BuildExtension):
     def __init__(self, *args, **kwargs) -> None:
         self.xformers_build_metadata = kwargs.pop("xformers_build_metadata")
+        self.pkg_name = "xformers"
+        self.metadata_json = "cpp_lib.json"
         super().__init__(*args, **kwargs)
 
     def build_extensions(self) -> None:
         super().build_extensions()
-        metadata_json = os.path.join(this_dir, "xformers", "cpp_lib.json")
-        with open(metadata_json, "w+") as fp:
+        with open(
+            os.path.join(self.build_lib, self.pkg_name, self.metadata_json), "w+"
+        ) as fp:
             json.dump(self.xformers_build_metadata, fp)
 
+    def copy_extensions_to_source(self):
+        """
+        Used for `pip install -e .`
+        Copies everything we built back into the source repo
+        """
+        build_py = self.get_finalized_command("build_py")
+        package_dir = build_py.get_package_dir(self.pkg_name)
+        inplace_file = os.path.join(package_dir, self.metadata_json)
+        regular_file = os.path.join(self.build_lib, self.pkg_name, self.metadata_json)
+        self.copy_file(regular_file, inplace_file, level=self.verbose)
+        super().copy_extensions_to_source()
+
 
 if __name__ == "__main__":
 
@@ -360,8 +375,6 @@ def build_extensions(self) -> None:
             ),
             "clean": clean,
         },
-        data_files=[("xformers", ["xformers/cpp_lib.json"])],
-        include_package_data=True,
         url="https://facebookresearch.github.io/xformers/",
         python_requires=">=3.7",
         author="Facebook AI Research",