Add file-backed serialization support (#21626)

This PR adds file-backed serialization support to ExecuTorch, allowing
large named-data blobs to be hashed, deduplicated, and written in chunks
without loading them entirely into memory. It also introduces backend
hooks for controlling program copying and AOTI weight materialization
while preserving existing behavior by default.

Author: Gasoonjia

backends/aoti/aoti_backend.py

@@ -112,6 +112,23 @@ def codesign_so(cls, so_path: str, compile_specs: List[CompileSpec]) -> None:
         """
         return
 
+    @classmethod
+    def load_weights_blob(
+        cls, blob_path: str, compile_specs: List[CompileSpec]
+    ) -> tuple[Any, str]:
+        """Load an AOTI weights blob and return its data and SHA-256 digest."""
+        with open(blob_path, "rb") as f:
+            blob_data = f.read()
+        os.remove(blob_path)
+        return blob_data, hashlib.sha256(blob_data).hexdigest()
+
+    @classmethod
+    def materialize_weights_blob(
+        cls, paths: Any, compile_specs: List[CompileSpec]
+    ) -> Any:
+        """Materialize backend-specific weight outputs into an AOTI blob."""
+        return paths
+
     @classmethod
     def move_program_to_device(
         cls,
@@ -257,6 +274,8 @@ def preprocess(
                 edge_program_module, tuple(user_input_placeholders), options=options
             )
 
+            paths = cls.materialize_weights_blob(paths, compile_specs)
+
             if len(missing_fallback_kernels) > 0:
                 formatted_kernels = "\n  - ".join(sorted(missing_fallback_kernels))
                 method_name = cls.method_name_from_compile_specs(compile_specs)
@@ -290,9 +309,7 @@ def preprocess(
         with open(so_path, "rb") as f:
             so_data = f.read()
 
-        # Read weights blob
-        with open(blob_path, "rb") as f:
-            blob_data = f.read()
+        blob_data, weights_blob_hash = cls.load_weights_blob(blob_path, compile_specs)
 
         # Create named data store
         named_data_store = NamedDataStore()
@@ -301,7 +318,7 @@ def preprocess(
         # keys (a method-name-only key collides). Runtime recovers them from
         # processed_bytes below.
         so_blob_key = hashlib.sha256(so_data).hexdigest() + "_so_blob"
-        weights_blob_key = hashlib.sha256(blob_data).hexdigest() + "_weights_blob"
+        weights_blob_key = weights_blob_hash + "_weights_blob"
 
         named_data_store.add_named_data(so_blob_key, so_data, 1, None)
         # Determine whether to save named data externally based on backend setting
@@ -314,7 +331,6 @@ def preprocess(
 
         # Clean up the generated files
         os.remove(so_path)
-        os.remove(blob_path)
 
         # Release device memory held by tensors that ``move_to_device_pass``
         # placed on the target device. Default impl is a no-op; concrete

exir/_serialize/_cord.py

@@ -4,10 +4,88 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import hashlib
 import io
+import os
+import shutil
+import tempfile
+import weakref
 from typing import List, Optional, Union
 
 
+class FileBackedData:
+    """A byte buffer that stays on disk until explicitly closed."""
+
+    _COPY_CHUNK_SIZE = 8 * 1024 * 1024
+
+    def __init__(self, path: str, cleanup: bool = False) -> None:
+        self._path = path
+        self._size = os.path.getsize(path)
+        self._sha256: Optional[bytes] = None
+        self._finalizer = (
+            weakref.finalize(self, self._remove, path) if cleanup else None
+        )
+
+    @staticmethod
+    def _remove(path: str) -> None:
+        try:
+            os.remove(path)
+        except OSError:
+            pass
+
+    @classmethod
+    def move_from(cls, path: str) -> "FileBackedData":
+        """Take ownership of ``path`` without loading its contents."""
+        directory = os.path.dirname(path) or "."
+        fd, owned_path = tempfile.mkstemp(
+            prefix=".executorch_", suffix=".data", dir=directory
+        )
+        os.close(fd)
+        try:
+            os.replace(path, owned_path)
+        except Exception:
+            os.remove(owned_path)
+            raise
+        return cls(owned_path, cleanup=True)
+
+    def __len__(self) -> int:
+        return self._size
+
+    def prefix(self, size: int) -> bytes:
+        with open(self._path, "rb") as f:
+            return f.read(size)
+
+    def sha256(self) -> bytes:
+        if self._sha256 is None:
+            digest = hashlib.sha256()
+            with open(self._path, "rb") as f:
+                while chunk := f.read(self._COPY_CHUNK_SIZE):
+                    digest.update(chunk)
+            self._sha256 = digest.digest()
+        return self._sha256
+
+    def to_bytes(self) -> bytes:
+        with open(self._path, "rb") as f:
+            return f.read()
+
+    def write_to_file(self, outfile: io.BufferedIOBase) -> None:
+        with open(self._path, "rb") as f:
+            shutil.copyfileobj(f, outfile, length=self._COPY_CHUNK_SIZE)
+
+    def close(self) -> None:
+        if self._finalizer is not None:
+            self._finalizer()
+
+    def __enter__(self) -> "FileBackedData":
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        self.close()
+
+
+CordBuffer = Union[bytes, FileBackedData]
+
+
 class Cord:
     """A `bytes`-like sequence of bytes, stored non-contiguously.
 
@@ -16,9 +94,9 @@ class Cord:
     `bytes` or `bytearray` object.
     """
 
-    def __init__(self, data: Optional[Union[bytes, "Cord"]] = None) -> None:
+    def __init__(self, data: Optional[Union[CordBuffer, "Cord"]] = None) -> None:
         """Initialize Cord data structure."""
-        self._buffers: List[bytes] = []
+        self._buffers: List[CordBuffer] = []
         self._byte_size: int = 0
 
         if data is not None:
@@ -30,20 +108,28 @@ def __len__(self):
 
     def __bytes__(self) -> bytes:
         """Return the contents of the Cord as a single `bytes` object."""
-        return b"".join(self._buffers)
+        return b"".join(
+            item if isinstance(item, bytes) else item.to_bytes()
+            for item in self._buffers
+        )
 
-    def append(self, data: Union[bytes, "Cord"]) -> None:
+    def append(self, data: Union[CordBuffer, "Cord"]) -> None:
         """Append a bytes or Cord to the current Cord."""
-        if isinstance(data, bytes):
+        if isinstance(data, (bytes, FileBackedData)):
             self._buffers.append(data)
             self._byte_size += len(data)
         elif isinstance(data, Cord):
             self._buffers.extend(data._buffers)
             self._byte_size += len(data)
         else:
-            raise TypeError(f"Can only append bytes or Cords, received {type(data)}")
+            raise TypeError(
+                f"Can only append bytes, FileBackedData, or Cords, received {type(data)}"
+            )
 
     def write_to_file(self, outfile: io.BufferedIOBase) -> None:
         """Write the Cord to a file."""
         for item in self._buffers:
-            outfile.write(item)
+            if isinstance(item, bytes):
+                outfile.write(item)
+            else:
+                item.write_to_file(outfile)

exir/_serialize/_named_data_store.py

@@ -11,6 +11,7 @@
 from typing import Dict, List, Optional, Tuple, Union
 
 import torch
+from executorch.exir._serialize._cord import CordBuffer, FileBackedData
 from executorch.exir._serialize.data_serializer import DataEntry
 from executorch.exir.tensor_layout import TensorLayout
 
@@ -47,7 +48,7 @@ class NamedDataStoreOutput:
             from {filename: {key: DataEntry}}.
     """
 
-    buffers: List[bytes]
+    buffers: List[CordBuffer]
     pte_data: Dict[str, DataEntry]
     external_data: Dict[str, Dict[str, DataEntry]]
 
@@ -68,7 +69,7 @@ class NamedDataStore:
     """
 
     # List of unique blobs.
-    buffers: List[bytes]
+    buffers: List[CordBuffer]
     # Named data stored inside the PTE file. Map of {key: DataEntry}.
     pte_data: Dict[str, DataEntry]
     # Named data stored outside of the PTE file.
@@ -93,17 +94,29 @@ def __init__(self) -> None:
         self.buffer_sha256 = {}
         self.key_to_buffer_idx = {}
 
+    @staticmethod
+    def _sha256(data: CordBuffer) -> bytes:
+        if isinstance(data, FileBackedData):
+            return data.sha256()
+        return hashlib.sha256(data).digest()
+
+    @staticmethod
+    def _prefix(data: CordBuffer, size: int) -> bytes:
+        if isinstance(data, FileBackedData):
+            return data.prefix(size)
+        return data[:size]
+
     def _get_buffer_sha256(self, buffer_idx: int) -> bytes:
         sha = self.buffer_sha256.get(buffer_idx)
         if sha is None:
-            sha = hashlib.sha256(self.buffers[buffer_idx]).digest()
+            sha = self._sha256(self.buffers[buffer_idx])
             self.buffer_sha256[buffer_idx] = sha
         return sha
 
     def _add_named_data_to_map(
         self,
         key: str,
-        data: bytes,
+        data: CordBuffer,
         alignment: int,
         local_key_to_buffer_idx: Dict[str, DataEntry],
         tensor_layout: Optional[TensorLayout] = None,
@@ -127,7 +140,9 @@ def _add_named_data_to_map(
         # Check if the key exists.
         buffer_idx = self.key_to_buffer_idx.get(key, -1)
         if buffer_idx != -1:
-            if data != self.buffers[buffer_idx]:
+            if len(data) != len(self.buffers[buffer_idx]) or self._sha256(
+                data
+            ) != self._get_buffer_sha256(buffer_idx):
                 raise ValueError(
                     f"Duplicate key {key} with different data. "
                     f"Existing data size: {len(self.buffers[buffer_idx])} bytes. "
@@ -136,10 +151,10 @@ def _add_named_data_to_map(
         else:
             # Two-level dedup: cheap fingerprint rejects non-matches fast,
             # SHA-256 confirms matches without full byte comparison.
-            fingerprint = (len(data), data[:32])
+            fingerprint = (len(data), self._prefix(data, 32))
             candidates = self.fingerprint_to_buffer_idx.get(fingerprint)
             if candidates is not None:
-                new_sha = hashlib.sha256(data).digest()
+                new_sha = self._sha256(data)
                 for candidate in candidates:
                     if new_sha == self._get_buffer_sha256(candidate):
                         buffer_idx = candidate
@@ -162,7 +177,7 @@ def _add_named_data_to_map(
     def add_named_data(
         self,
         key: str,
-        data: Union[bytes, torch.Tensor],
+        data: Union[bytes, FileBackedData, torch.Tensor],
         alignment: Optional[int] = 1,
         external_tag: Optional[str] = None,
         tensor_layout: Optional[TensorLayout] = None,
@@ -171,7 +186,8 @@ def add_named_data(
         Adds a named blob to the NamedDataStore.
         Args:
             key (str): key associated with the data.
-            data (Union[bytes, torch.Tensor]): Union of bytes, or torch.Tensor to serialize. Note: if a tensor is passed, it must have contiguous memory layout. The tensor_layout will be inferred from the tensor and should not be passed in.
+            data: Bytes, file-backed data, or a torch.Tensor to serialize. If a
+                tensor is passed, its layout is inferred.
             alignment (int): alignment for bytes to be serialized with.
             external (Optional[str]): the external filename that this data is saved to.
             tensor_layout (Optional[TensorLayout]): layout of the tensor, if applicable.
@@ -194,8 +210,10 @@ def add_named_data(
                 )
             tensor_layout = real_tensor_layout
             byte_data = _tensor_to_bytes(data)
-        else:
+        elif isinstance(data, (bytes, FileBackedData)):
             byte_data = data
+        else:
+            raise TypeError(f"Unsupported named data type: {type(data)}")
 
         if external_tag is None:
             self._add_named_data_to_map(

exir/_serialize/data_serializer.py

@@ -2,7 +2,7 @@
 from dataclasses import dataclass
 from typing import Dict, Optional, Sequence
 
-from executorch.exir._serialize._cord import Cord
+from executorch.exir._serialize._cord import Cord, CordBuffer
 from executorch.exir.tensor_layout import TensorLayout
 
 
@@ -36,7 +36,7 @@ class DataPayload:
         key_to_data: a map from unique keys to serializable data.
     """
 
-    buffers: Sequence[bytes]
+    buffers: Sequence[CordBuffer]
     named_data: Dict[str, DataEntry]
 
 

exir/_serialize/test/test_cord.py

@@ -6,9 +6,11 @@
 
 
 import io
+import os
+import tempfile
 import unittest
 
-from executorch.exir._serialize._cord import Cord
+from executorch.exir._serialize._cord import Cord, FileBackedData
 
 
 class TestCord(unittest.TestCase):
@@ -61,3 +63,22 @@ def test_cord_write_to_file(self) -> None:
         outfile = io.BytesIO()
         cord.write_to_file(outfile)
         self.assertEqual(b"HelloWorld", outfile.getvalue())
+
+    def test_file_backed_data(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            source_path = os.path.join(directory, "source.bin")
+            with open(source_path, "wb") as f:
+                f.write(b"FileBacked")
+
+            with FileBackedData.move_from(source_path) as data:
+                self.assertFalse(os.path.exists(source_path))
+                self.assertEqual(10, len(data))
+
+                cord = Cord(b"Prefix")
+                cord.append(data)
+                outfile = io.BytesIO()
+                cord.write_to_file(outfile)
+                self.assertEqual(b"PrefixFileBacked", outfile.getvalue())
+                self.assertEqual(b"PrefixFileBacked", bytes(cord))
+
+            self.assertEqual([], os.listdir(directory))