From 9cf82eafd91e0643411f2c4f4ae57d975d3d46ae Mon Sep 17 00:00:00 2001 From: Gantaphon Chalumporn Date: Thu, 26 Jun 2025 10:11:20 -0700 Subject: [PATCH] Add Manifold wrapper (attemp 2) Part2 Frontend Summary: Summary This diff is a second attemp to add monifold wrapper. This diff merged two reverted diffs D76186007, D76603891. The reversion is due to S532997. The two reveretd diffs are merged. Then, this diff attemps to include frontend part of the manifold wrapper. *** Only land this diff when the backend has already landed *** Differential Revision: D77241255 --- fbgemm_gpu/fbgemm_gpu/utils/filestore.py | 51 ++++++++++++++++++++- fbgemm_gpu/test/utils/filestore_test.py | 58 ++++++++++++++++++++++-- 2 files changed, 105 insertions(+), 4 deletions(-) diff --git a/fbgemm_gpu/fbgemm_gpu/utils/filestore.py b/fbgemm_gpu/fbgemm_gpu/utils/filestore.py index 9261f85922..d5639c5639 100644 --- a/fbgemm_gpu/fbgemm_gpu/utils/filestore.py +++ b/fbgemm_gpu/fbgemm_gpu/utils/filestore.py @@ -155,4 +155,53 @@ def exists(self, path: str) -> bool: True if file exists, False otherwise. """ filepath = f"{self.bucket}/{path}" - return os.path.isfile(filepath) + return os.path.exists(filepath) + + def create_directory(self, path: str) -> "FileStore": + """ + Creates a directory in the file store. + + Args: + path (str): The path of the node or symlink to a directory (relative + to `self.bucket`) to be created. + + Returns: + self. This allows for method-chaining. + """ + filepath = f"{self.bucket}/{path}" + event = f"creating directory {filepath}" + logger.info(f"FileStore: {event}") + + try: + if not os.path.exists(filepath): + os.makedirs(filepath, exist_ok=True) + except Exception as e: + logger.error(f"FileStore: exception occurred when {event}: {e}") + raise e + + return self + + def remove_directory(self, path: str) -> "FileStore": + """ + Removes a directory from the file store. + + Args: + path (str): The path of the node or symlink to a directory (relative + to `self.bucket`) to be removed. + + Returns: + self. This allows for method-chaining. + """ + filepath = f"{self.bucket}/{path}" + event = f"deleting {filepath}" + logger.info(f"FileStore: {event}") + + try: + if os.path.isdir(filepath): + os.rmdir(filepath) + + except Exception as e: + logger.error(f"Manifold: exception occurred when {event}: {e}") + raise e + + return self diff --git a/fbgemm_gpu/test/utils/filestore_test.py b/fbgemm_gpu/test/utils/filestore_test.py index 38877e6785..329f0ddea5 100644 --- a/fbgemm_gpu/test/utils/filestore_test.py +++ b/fbgemm_gpu/test/utils/filestore_test.py @@ -60,6 +60,41 @@ def _test_filestore_readwrite( store.remove(path) assert not store.exists(path), f"{path} is not removed" + def _test_filestore_directory( + self, + # pyre-fixme[2] + store, # FileStore + root_dir: Optional[str] = None, + ) -> None: + """ + Generic FileStore routines to test creating and removing directories + + Args: + store (FileStore): The FileStore to test + root_dir (str): The root directory to create + """ + if root_dir is not None: + root_dir += "/" + else: + root_dir = "" + + dir1 = f"{''.join(random.choices(string.ascii_letters, k=15))}" + dir2 = f"{''.join(random.choices(string.ascii_letters, k=15))}" + + store.create_directory(f"{root_dir}{dir1}/{dir2}") + assert store.exists( + f"{root_dir}{dir1}/{dir2}" + ), f"Failed creating directories /{dir1}/{dir2}, directoies does not exist" + + store.remove_directory(f"{root_dir}{dir1}/{dir2}") + assert not store.exists( + f"{root_dir}{dir1}/{dir2}" + ), f"Failed removing directories /{dir1}/{dir2}, directory still exists" + store.remove_directory(f"{root_dir}{dir1}") + assert not store.exists( + f"{root_dir}{dir1}" + ), f"Failed removing directories /{dir1}, directory still exists" + def test_filestore_oss_bad_bucket(self) -> None: """ Test that OSS FileStore raises ValueError when an invalid bucket is provided @@ -104,6 +139,14 @@ def test_filestore_oss_file(self) -> None: self._test_filestore_readwrite(FileStore("/tmp"), Path(infile.name)) + def test_filestore_oss_directory(self) -> None: + """ + Test that OSS FileStore can create and remove directories + """ + from fbgemm_gpu.utils import FileStore + + self._test_filestore_directory(FileStore("/tmp")) + @unittest.skipIf(open_source, "Test does not apply to OSS") def test_filestore_fb_bad_bucket(self) -> None: """ @@ -125,7 +168,7 @@ def test_filestore_fb_binaryio(self) -> None: self._test_filestore_readwrite( FileStore("tlparse_reports"), io.BytesIO("".join(random.choices(string.ascii_letters, k=128)).encode()), - f"tree/{''.join(random.choices(string.ascii_letters, k=15))}.unittest", + f"tree/unit_tests/{''.join(random.choices(string.ascii_letters, k=15))}.unittest", ) @unittest.skipIf(open_source, "Test does not apply to OSS") @@ -138,7 +181,7 @@ def test_filestore_fb_tensor(self) -> None: self._test_filestore_readwrite( FileStore("tlparse_reports"), torch.rand((random.randint(100, 1000), random.randint(100, 1000))), - f"tree/{''.join(random.choices(string.ascii_letters, k=15))}.unittest", + f"tree/unit_tests/{''.join(random.choices(string.ascii_letters, k=15))}.unittest", ) @unittest.skipIf(open_source, "Test does not apply to OSS") @@ -155,5 +198,14 @@ def test_filestore_fb_file(self) -> None: self._test_filestore_readwrite( FileStore("tlparse_reports"), Path(infile.name), - f"tree/{''.join(random.choices(string.ascii_letters, k=15))}.unittest", + f"tree/unit_tests/{''.join(random.choices(string.ascii_letters, k=15))}.unittest", ) + + @unittest.skipIf(open_source, "Test does not apply to OSS") + def test_filestore_fb_directory(self) -> None: + """ + Test that FB FileStore can create and remove directories + """ + from fbgemm_gpu.fb.utils import FileStore + + self._test_filestore_directory(FileStore("tlparse_reports"), "tree/unit_tests")