Add AppendingByteSerializer class (pytorch#148226)

This PR adds a new util class that enables efficient appending of sequential byte data with custom serialization and deserialization.
Pull Request resolved: pytorch#148226
Approved by: https://github.com/aorenste

Author: oulgen

test/test_appending_byte_serializer.py

@@ -0,0 +1,85 @@
+# Owner(s): ["module: inductor"]
+
+import dataclasses
+
+from torch.testing._internal.common_utils import TestCase
+from torch.utils._appending_byte_serializer import (
+    AppendingByteSerializer,
+    BytesReader,
+    BytesWriter,
+)
+
+
+class TestAppendingByteSerializer(TestCase):
+    def test_write_and_read_int(self) -> None:
+        def int_serializer(writer: BytesWriter, i: int) -> None:
+            writer.write_uint64(i)
+
+        def int_deserializer(reader: BytesReader) -> int:
+            return reader.read_uint64()
+
+        s = AppendingByteSerializer(serialize_fn=int_serializer)
+
+        data = [1, 2, 3, 4]
+        s.extend(data)
+        self.assertListEqual(
+            data,
+            AppendingByteSerializer.to_list(
+                s.to_bytes(), deserialize_fn=int_deserializer
+            ),
+        )
+
+        data2 = [8, 9, 10, 11]
+        s.extend(data2)
+        self.assertListEqual(
+            data + data2,
+            AppendingByteSerializer.to_list(
+                s.to_bytes(), deserialize_fn=int_deserializer
+            ),
+        )
+
+    def test_write_and_read_class(self) -> None:
+        @dataclasses.dataclass(frozen=True, eq=True)
+        class Foo:
+            x: int
+            y: str
+            z: bytes
+
+            @staticmethod
+            def serialize(writer: BytesWriter, cls: "Foo") -> None:
+                writer.write_uint64(cls.x)
+                writer.write_str(cls.y)
+                writer.write_bytes(cls.z)
+
+            @staticmethod
+            def deserialize(reader: BytesReader) -> "Foo":
+                x = reader.read_uint64()
+                y = reader.read_str()
+                z = reader.read_bytes()
+                return Foo(x, y, z)
+
+        a = Foo(5, "ok", bytes([15]))
+        b = Foo(10, "lol", bytes([25]))
+
+        s = AppendingByteSerializer(serialize_fn=Foo.serialize)
+        s.append(a)
+        self.assertListEqual(
+            [a],
+            AppendingByteSerializer.to_list(
+                s.to_bytes(), deserialize_fn=Foo.deserialize
+            ),
+        )
+
+        s.append(b)
+        self.assertListEqual(
+            [a, b],
+            AppendingByteSerializer.to_list(
+                s.to_bytes(), deserialize_fn=Foo.deserialize
+            ),
+        )
+
+
+if __name__ == "__main__":
+    from torch._inductor.test_case import run_tests
+
+    run_tests()

torch/utils/_appending_byte_serializer.py

@@ -0,0 +1,109 @@
+from collections.abc import Iterable
+from typing import Callable, Generic, TypeVar
+
+
+T = TypeVar("T")
+
+_ENCODING_VERSION: int = 1
+
+__all__ = ["AppendingByteSerializer"]
+
+
+#######################################
+# Helper classes
+#######################################
+
+
+class BytesWriter:
+    def __init__(self, preallocate_size: int) -> None:
+        self._data = bytearray(preallocate_size)
+
+    def write_uint64(self, i: int) -> None:
+        self._data.extend(i.to_bytes(8, byteorder="big", signed=False))
+
+    def write_str(self, s: str) -> None:
+        payload = s.encode("utf-8")
+        self.write_bytes(payload)
+
+    def write_bytes(self, b: bytes) -> None:
+        self.write_uint64(len(b))
+        self._data.extend(b)
+
+    def to_bytes(self) -> bytes:
+        return bytes(self._data)
+
+
+class BytesReader:
+    def __init__(self, data: bytes) -> None:
+        self._data = data
+        self._i = 0
+
+    def is_finished(self) -> bool:
+        return len(self._data) == self._i
+
+    def read_uint64(self) -> int:
+        result = int.from_bytes(
+            self._data[self._i : self._i + 8], byteorder="big", signed=False
+        )
+        self._i += 8
+        return result
+
+    def read_str(self) -> str:
+        return self.read_bytes().decode("utf-8")
+
+    def read_bytes(self) -> bytes:
+        size = self.read_uint64()
+        result = self._data[self._i : self._i + size]
+        self._i += size
+        return result
+
+
+#######################################
+# AppendingByteSerializer
+#######################################
+
+
+class AppendingByteSerializer(Generic[T]):
+    """
+    Provides efficient serialization and deserialization of list of bytes
+    Note that this does not provide any guarantees around byte order
+    """
+
+    _serialize_fn: Callable[[BytesWriter, T], None]
+    _writer: BytesWriter
+    _preallocate_size: int
+
+    def __init__(
+        self,
+        *,
+        serialize_fn: Callable[[BytesWriter, T], None],
+        preallocate_size: int = 0,
+    ) -> None:
+        self._serialize_fn = serialize_fn
+        self._preallocate_size = preallocate_size
+        self.clear()
+
+    def clear(self) -> None:
+        self._writer = BytesWriter(preallocate_size=self._preallocate_size)
+        # First 8-bytes are for version
+        self._writer.write_uint64(_ENCODING_VERSION)
+
+    def append(self, data: T) -> None:
+        self._serialize_fn(self._writer, data)
+
+    def extend(self, elems: Iterable[T]) -> None:
+        for elem in elems:
+            self.append(elem)
+
+    def to_bytes(self) -> bytes:
+        return self._writer.to_bytes()
+
+    @staticmethod
+    def to_list(data: bytes, *, deserialize_fn: Callable[[BytesReader], T]) -> list[T]:
+        reader = BytesReader(data)
+        assert reader.read_uint64() == _ENCODING_VERSION
+
+        result: list[T] = []
+        while not reader.is_finished():
+            result.append(deserialize_fn(reader))
+        return result