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