Add abstractions for writing out data (flesh out D71147675, pt 1) (pytorch#947)

q10 · facebook-github-bot · commit 6ccad58c37d3 · 2025-03-20T17:47:05.000-07:00
Summary: Pull Request resolved: facebookresearch/FBGEMM#947 X-link: pytorch#3856 - Add abstractions for writing out data to FileStorage. This folds in common code from D71147675, D69883868, and D64932387 Reviewed By: spcyppt Differential Revision: D71350076 fbshipit-source-id: 6f0e01d15e5afa79e62db9c4af46c7fab2980975
diff --git a/fbgemm_gpu/fbgemm_gpu/utils/__init__.py b/fbgemm_gpu/fbgemm_gpu/utils/__init__.py
@@ -7,4 +7,5 @@
 
 # pyre-unsafe
 
-from fbgemm_gpu.utils.torch_library import TorchLibraryFragment  # noqa F401
+from .filestore import FileStore  # noqa F401
+from .torch_library import TorchLibraryFragment  # noqa F401
diff --git a/fbgemm_gpu/fbgemm_gpu/utils/filestore.py b/fbgemm_gpu/fbgemm_gpu/utils/filestore.py
@@ -0,0 +1,160 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+# pyre-ignore-all-errors[56]
+
+import io
+import logging
+import os
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import BinaryIO, Union
+
+import torch
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class FileStore:
+    """
+    A basic file store implementation for easy data reads / writes / deletes.
+
+    This class is intended to be used as a utility inside the FBGEMM_GPU codebase
+    for consistent writing of tensors and other objects to the filesystem.
+
+    Attribute:
+        bucket (str): A directory in the filesystem.
+    """
+
+    bucket: str
+
+    def __post_init__(self) -> None:
+        # self.bucket = bucket
+
+        if not os.path.isdir(self.bucket):
+            raise ValueError(f"Directory {self.bucket} does not exist")
+
+    def write(
+        self,
+        path: str,
+        raw_input: Union[BinaryIO, torch.Tensor, Path],
+        ttls: int = 864000,
+    ) -> "FileStore":
+        """
+        Writes a binary stream, or a torch.Tensor to the file located at `path`
+        (relative to `self.bucket`).
+
+        Args:
+            path (str): The path of the node or symlink to a directory.
+            raw_input (BinaryIO | torch.Tensor | Path): The data to write.
+
+            ttls (int): The time to live for the data in seconds. Defaults to
+            10 days.
+
+        Returns:
+            self.  This allows for method-chaining.
+        """
+
+        filepath = f"{self.bucket}/{path}"
+        event = f"writing to {filepath}"
+        logger.info(f"FileStore: {event}")
+
+        try:
+            if os.path.isfile(filepath):
+                raise FileExistsError(
+                    f"File {filepath} already exists in the filesystem"
+                )
+
+            if isinstance(raw_input, torch.Tensor):
+                torch.save(raw_input, filepath)
+
+            elif isinstance(raw_input, Path):
+                if not os.path.exists(raw_input):
+                    raise FileNotFoundError(f"File {raw_input} does not exist")
+                shutil.copyfile(raw_input, filepath)
+
+            elif isinstance(raw_input, io.BytesIO) or isinstance(raw_input, BinaryIO):
+                with open(filepath, "wb") as file:
+                    raw_input.seek(0)
+                    while chunk := raw_input.read(4096):  # Read 4 KB at a time
+                        file.write(chunk)
+            else:
+                raise TypeError(f"Unsupported input type: {type(raw_input)}")
+
+        except Exception as e:
+            logger.error(f"FileStore: exception occurred when {event}: {e}")
+            raise e
+
+        return self
+
+    def read(self, path: str) -> io.BytesIO:
+        """
+        Reads a file into a BytesIO object.
+
+        Args:
+            path (str): The path of the node or symlink to a directory (relative
+            to `self.bucket`) to be read.
+
+        Returns:
+            Data from the file in BytesIO object format.
+        """
+        filepath = f"{self.bucket}/{path}"
+        event = f"reading from {filepath}"
+        logger.info(f"FileStore: {event}")
+
+        try:
+            if not os.path.isfile(filepath):
+                raise FileNotFoundError(
+                    f"File {filepath} does not exist in the FileStore"
+                )
+
+            return io.BytesIO(open(filepath, "rb").read())
+
+        except Exception as e:
+            logger.error(f"FileStore: exception occurred when {event}: {e}")
+            raise e
+
+    def remove(self, path: str) -> "FileStore":
+        """
+        Removes a file or 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.isfile(filepath):
+                os.remove(filepath)
+
+        except Exception as e:
+            logger.error(f"Manifold: exception occurred when {event}: {e}")
+            raise e
+
+        return self
+
+    def exists(self, path: str) -> bool:
+        """
+        Checks for existence of file in the file store.
+
+        Args:
+            path (str): The Manifold target path (relative to `self.bucket`).
+
+        Returns:
+            True if file exists, False otherwise.
+        """
+        filepath = f"{self.bucket}/{path}"
+        return os.path.isfile(filepath)
diff --git a/fbgemm_gpu/test/utils/filestore_test.py b/fbgemm_gpu/test/utils/filestore_test.py
@@ -0,0 +1,159 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import io
+import random
+import string
+import tempfile
+import unittest
+from pathlib import Path
+from typing import BinaryIO, Optional, Union
+
+import fbgemm_gpu
+import torch
+
+# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`.
+open_source: bool = getattr(fbgemm_gpu, "open_source", False)
+
+
+class FileStoreTest(unittest.TestCase):
+    def _test_filestore_readwrite(
+        self,
+        # pyre-fixme[2]
+        store,  # FileStore
+        input: Union[BinaryIO, torch.Tensor, Path],
+        path: Optional[str] = None,
+    ) -> None:
+        """
+        Generic FileStore routines to test reading and writing data
+
+        Args:
+            store (FileStore): The FileStore to test
+            input (torch.Tensor | BinaryIO | Path): The data to write to the FileStore
+            path (str, optional): The path to write the data to. If not provided, a random path will be generated.
+        """
+        if path is None:
+            path = "".join(random.choices(string.ascii_letters, k=15))
+
+        assert not store.exists(path), f"{path} already exists"
+        store.write(path, input)
+        assert store.exists(path), f"{path} does not exist"
+
+        if isinstance(input, torch.Tensor):
+            assert torch.load(store.read(path)).equal(input), "tensors do not match"
+
+        elif isinstance(input, io.BytesIO) or isinstance(input, BinaryIO):
+            input.seek(0)
+            assert store.read(path).read() == input.read(), "byte streams do not match"
+
+        elif isinstance(input, Path):
+            assert (
+                store.read(path).read() == input.read_bytes()
+            ), "file contents do not match"
+
+        store.remove(path)
+        assert not store.exists(path), f"{path} is not removed"
+
+    def test_filestore_oss_bad_bucket(self) -> None:
+        """
+        Test that OSS FileStore raises ValueError when an invalid bucket is provided
+        """
+        from fbgemm_gpu.utils import FileStore
+
+        self.assertRaises(
+            ValueError, FileStore, "".join(random.choices(string.ascii_letters, k=15))
+        )
+
+    def test_filestore_oss_binaryio(self) -> None:
+        """
+        Test that OSS FileStore can read and write binary data
+        """
+        from fbgemm_gpu.utils import FileStore
+
+        self._test_filestore_readwrite(
+            FileStore("/tmp"),
+            io.BytesIO("".join(random.choices(string.ascii_letters, k=128)).encode()),
+        )
+
+    def test_filestore_oss_tensor(self) -> None:
+        """
+        Test that OSS FileStore can read and write tensors
+        """
+        from fbgemm_gpu.utils import FileStore
+
+        self._test_filestore_readwrite(
+            FileStore("/tmp"),
+            torch.rand((random.randint(100, 1000), random.randint(100, 1000))),
+        )
+
+    def test_filestore_oss_file(self) -> None:
+        """
+        Test that OSS FileStore can read and write files
+        """
+        from fbgemm_gpu.utils import FileStore
+
+        input = torch.rand((random.randint(100, 1000), random.randint(100, 1000)))
+        infile = tempfile.NamedTemporaryFile()
+        torch.save(input, infile)
+
+        self._test_filestore_readwrite(FileStore("/tmp"), Path(infile.name))
+
+    @unittest.skipIf(open_source, "Test does not apply to OSS")
+    def test_filestore_fb_bad_bucket(self) -> None:
+        """
+        Test that FB FileStore raises ValueError when an invalid bucket is provided
+        """
+        from fbgemm_gpu.fb.utils import FileStore
+
+        self.assertRaises(
+            ValueError, FileStore, "".join(random.choices(string.ascii_letters, k=15))
+        )
+
+    @unittest.skipIf(open_source, "Test does not apply to OSS")
+    def test_filestore_fb_binaryio(self) -> None:
+        """
+        Test that FB FileStore can read and write binary data
+        """
+        from fbgemm_gpu.fb.utils import FileStore
+
+        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",
+        )
+
+    @unittest.skipIf(open_source, "Test does not apply to OSS")
+    def test_filestore_fb_tensor(self) -> None:
+        """
+        Test that FB FileStore can read and write tensors
+        """
+        from fbgemm_gpu.fb.utils import FileStore
+
+        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",
+        )
+
+    @unittest.skipIf(open_source, "Test does not apply to OSS")
+    def test_filestore_fb_file(self) -> None:
+        """
+        Test that FB FileStore can read and write files
+        """
+        from fbgemm_gpu.fb.utils import FileStore
+
+        input = torch.rand((random.randint(100, 1000), random.randint(100, 1000)))
+        infile = tempfile.NamedTemporaryFile()
+        torch.save(input, infile)
+
+        self._test_filestore_readwrite(
+            FileStore("tlparse_reports"),
+            Path(infile.name),
+            f"tree/{''.join(random.choices(string.ascii_letters, k=15))}.unittest",
+        )
