Avoid copying the whole file when reading .pte and .ptd headers

backends/mlx/pte_inspector.py

@@ -34,8 +34,8 @@
     MLX_OP_TYPE_NAMES,
 )
 from executorch.exir._serialize._program import (
-    _ExtendedHeader,
     _extract_delegate_payload as extract_delegate_payload,
+    _get_extended_header,
 )
 
 MLX_MAGIC = b"MLX0"
@@ -468,24 +468,20 @@ def parse_executorch_program(pte_data: bytes) -> Dict[str, Any]:  # noqa: C901
     result["flatbuffer_magic"] = fb_magic.decode("utf-8", errors="replace")
 
     extended_header_offset = 8
-    if len(pte_data) > extended_header_offset + 32:
-        try:
-            header = _ExtendedHeader.from_bytes(pte_data[extended_header_offset:])
-            if header.is_valid():
-                result["extended_header"] = {
-                    "magic": header.magic.decode("utf-8", errors="replace"),
-                    "length": header.length,
-                    "program_size": header.program_size,
-                    "segment_base_offset": header.segment_base_offset,
-                    "segment_data_size": header.segment_data_size,
-                }
-                fb_start = extended_header_offset + header.length
-                result["flatbuffer_offset"] = fb_start
-                result["flatbuffer_size"] = header.program_size
-                result["segment_offset"] = header.segment_base_offset
-                result["segment_size"] = header.segment_data_size
-        except Exception as e:
-            result["header_parse_error"] = str(e)
+    header = _get_extended_header(pte_data)
+    if header is not None:
+        result["extended_header"] = {
+            "magic": header.magic.decode("utf-8", errors="replace"),
+            "length": header.length,
+            "program_size": header.program_size,
+            "segment_base_offset": header.segment_base_offset,
+            "segment_data_size": header.segment_data_size,
+        }
+        fb_start = extended_header_offset + header.length
+        result["flatbuffer_offset"] = fb_start
+        result["flatbuffer_size"] = header.program_size
+        result["segment_offset"] = header.segment_base_offset
+        result["segment_size"] = header.segment_data_size
 
     try:
         from executorch.exir._serialize._flatbuffer import _program_flatbuffer_to_json

exir/_serialize/_program.py

@@ -248,7 +248,10 @@ def to_bytes(self) -> bytes:
 def _get_extended_header(program_data: bytes) -> Optional[_ExtendedHeader]:
     """Returns the extended header of the program data, if present and valid."""
     try:
-        eh = _ExtendedHeader.from_bytes(program_data[8:])
+        # program_data[8:] would copy the whole file.
+        eh = _ExtendedHeader.from_bytes(
+            program_data[8 : 8 + _ExtendedHeader.EXPECTED_LENGTH]
+        )
         if eh.is_valid():
             return eh
     except ValueError:

exir/_serialize/test/test_program.py

@@ -1332,6 +1332,21 @@ def test_from_bytes_with_more_data_than_necessary(self) -> None:
         self.assertEqual(eh.segment_base_offset, EXAMPLE_SEGMENT_BASE_OFFSET)
         self.assertEqual(eh.segment_data_size, EXAMPLE_SEGMENT_DATA_SIZE)
 
+    def test_get_extended_header_reads_only_header_bytes(self) -> None:
+        # Passing program_data[8:] to from_bytes() would copy the whole file.
+        pte_data = b"\x00" * 8 + EXAMPLE_HEADER_DATA + b"\x55" * 1024
+
+        with patch.object(
+            _ExtendedHeader, "from_bytes", wraps=_ExtendedHeader.from_bytes
+        ) as from_bytes:
+            eh = _get_extended_header(pte_data)
+
+        self.assertIsNotNone(eh)
+        self.assertEqual(eh.program_size, EXAMPLE_PROGRAM_SIZE)
+        self.assertEqual(
+            len(from_bytes.call_args.args[0]), _ExtendedHeader.EXPECTED_LENGTH
+        )
+
     def test_from_bytes_larger_than_needed_header_size_field(self) -> None:
         # Simulate a backwards-compatibility situation. Parse a header
         # with a larger-than expected size. This would typically mean that

extension/flat_tensor/serialize/serialize.py

@@ -217,7 +217,10 @@ def to_bytes(self) -> bytes:
 def _get_extended_header(flat_tensor_data: bytes) -> Optional[FlatTensorHeader]:
     """Returns the extended header of the flat_tensor data, if present and valid."""
     try:
-        eh = FlatTensorHeader.from_bytes(flat_tensor_data[8:])
+        # flat_tensor_data[8:] would copy the whole file.
+        eh = FlatTensorHeader.from_bytes(
+            flat_tensor_data[8 : 8 + FlatTensorHeader.EXPECTED_LENGTH]
+        )
         if eh.is_valid():
             return eh
     except ValueError:
@@ -372,8 +375,8 @@ def deserialize(self, blob: Cord) -> DataPayload:
         data = bytes(blob)
 
         # Read header. Verify that it's valid.
-        header = FlatTensorHeader.from_bytes(data[8:])
-        if not header.is_valid():
+        header = _get_extended_header(data)
+        if header is None:
             raise RuntimeError(
                 "Flat tensor header is invalid. File is likely incorrect format or corrupt."
             )

extension/flat_tensor/test/test_serialize.py

@@ -309,6 +309,20 @@ def test_round_trip(self) -> None:
             TEST_DATA_PAYLOAD.named_data, deserialized_payload.named_data
         )
 
+    def test_reads_only_header_bytes(self) -> None:
+        # Passing data[8:] to from_bytes() would copy the whole file.
+        serializer: DataSerializer = FlatTensorSerializer(FlatTensorConfig())
+        serialized_data = bytes(serializer.serialize(TEST_DATA_PAYLOAD))
+
+        with mock.patch.object(
+            FlatTensorHeader, "from_bytes", wraps=FlatTensorHeader.from_bytes
+        ) as from_bytes:
+            serializer.deserialize(Cord(serialized_data))
+
+        self.assertEqual(
+            len(from_bytes.call_args.args[0]), FlatTensorHeader.EXPECTED_LENGTH
+        )
+
     def test_deserialize_refuses_newer_version(self) -> None:
         # A file whose version is newer than this reader understands must be
         # refused, rather than silently misread. Serialize with a bumped

pytest.ini

@@ -61,7 +61,6 @@ addopts =
 
 
     # Ignore tests with missing dependencies or build issues
-    --ignore=extension/flat_tensor/test/test_serialize.py
     --ignore=extension/llm/custom_ops/test_preprocess_custom_ops.py
     --ignore=extension/llm/tokenizers/third-party/re2/python/re2_test.py
     --ignore=extension/llm/tokenizers/third-party/sentencepiece/python/test/sentencepiece_test.py