From 1a61cc198ddf528b176b501653d629da05676066 Mon Sep 17 00:00:00 2001 From: HiroForYou Date: Wed, 7 Oct 2026 17:52:18 -0500 Subject: [PATCH 1/5] Avoid copying the whole file when reading .pte and .ptd headers Reading the extended header with data[8:] copies the entire file, even though the header is at most 40 bytes. Slice only the bytes needed for the header instead. --- exir/_serialize/_program.py | 5 ++++- extension/flat_tensor/serialize/serialize.py | 9 +++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/exir/_serialize/_program.py b/exir/_serialize/_program.py index cf253475589..0a28385294e 100644 --- a/exir/_serialize/_program.py +++ b/exir/_serialize/_program.py @@ -262,7 +262,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: diff --git a/extension/flat_tensor/serialize/serialize.py b/extension/flat_tensor/serialize/serialize.py index c050b6dbeb9..5a1987dbc16 100644 --- a/extension/flat_tensor/serialize/serialize.py +++ b/extension/flat_tensor/serialize/serialize.py @@ -233,7 +233,10 @@ def __init__(self, data: Cord, alignment: Optional[int] = None) -> None: 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: @@ -394,7 +397,9 @@ def deserialize(self, blob: Cord) -> DataPayload: data = bytes(blob) # Read header. Verify that it's valid. - header = FlatTensorHeader.from_bytes(data[8:]) + header = FlatTensorHeader.from_bytes( + data[8 : 8 + FlatTensorHeader.EXPECTED_LENGTH] + ) if not header.is_valid(): raise RuntimeError( "Flat tensor header is invalid. File is likely incorrect format or corrupt." From 38a11261e90cebd99a665f3a1ac3c7180aafe4f3 Mon Sep 17 00:00:00 2001 From: HiroForYou Date: Wed, 7 Oct 2026 18:23:05 -0500 Subject: [PATCH 2/5] Add regression tests for reading only the header bytes --- exir/_serialize/test/test_program.py | 15 +++++++++++++++ extension/flat_tensor/test/test_serialize.py | 16 ++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/exir/_serialize/test/test_program.py b/exir/_serialize/test/test_program.py index c53b52ae418..96e69c76e6b 100644 --- a/exir/_serialize/test/test_program.py +++ b/exir/_serialize/test/test_program.py @@ -1300,6 +1300,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: + # Slicing to the end of the data would copy the whole file. + class NoSliceToEnd(bytes): + def __getitem__(self, key): + if isinstance(key, slice) and key.stop is None: + raise AssertionError("sliced to the end of the data") + return super().__getitem__(key) + + pte_data = NoSliceToEnd(b"\x00" * 8 + EXAMPLE_HEADER_DATA + b"\x55" * 1024) + eh = _get_extended_header(pte_data) + + self.assertIsNotNone(eh) + self.assertEqual(eh.program_size, EXAMPLE_PROGRAM_SIZE) + self.assertEqual(eh.segment_base_offset, EXAMPLE_SEGMENT_BASE_OFFSET) + 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 diff --git a/extension/flat_tensor/test/test_serialize.py b/extension/flat_tensor/test/test_serialize.py index 0f447d4bc78..09226f024cc 100644 --- a/extension/flat_tensor/test/test_serialize.py +++ b/extension/flat_tensor/test/test_serialize.py @@ -34,6 +34,7 @@ _deserialize_to_flat_tensor, _FLAT_TENSOR_VERSION, _FLATBUFFER_ALIGNMENT, + _get_extended_header, FlatTensorConfig, FlatTensorHeader, FlatTensorSerializer, @@ -284,6 +285,21 @@ def test_round_trip(self) -> None: TEST_DATA_PAYLOAD.named_data, deserialized_payload.named_data ) + def test_get_extended_header_reads_only_header_bytes(self) -> None: + # Slicing to the end of the data would copy the whole file. + class NoSliceToEnd(bytes): + def __getitem__(self, key): + if isinstance(key, slice) and key.stop is None: + raise AssertionError("sliced to the end of the data") + return super().__getitem__(key) + + serializer: DataSerializer = FlatTensorSerializer(FlatTensorConfig()) + serialized_data = NoSliceToEnd(serializer.serialize(TEST_DATA_PAYLOAD)) + header = _get_extended_header(serialized_data) + + self.assertIsNotNone(header) + self.assertTrue(header.is_valid()) + 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 From 75cec3290db20829c9dc18a32d11517fdc919826 Mon Sep 17 00:00:00 2001 From: HiroForYou Date: Fri, 9 Oct 2026 09:53:20 -0500 Subject: [PATCH 3/5] Cover deserialize() in the .ptd test and fix the MLX inspector --- backends/mlx/pte_inspector.py | 7 +++++- extension/flat_tensor/test/test_serialize.py | 24 ++++++++++---------- 2 files changed, 18 insertions(+), 13 deletions(-) diff --git a/backends/mlx/pte_inspector.py b/backends/mlx/pte_inspector.py index 20ec9e9d9df..cb2b9182ba8 100644 --- a/backends/mlx/pte_inspector.py +++ b/backends/mlx/pte_inspector.py @@ -470,7 +470,12 @@ def parse_executorch_program(pte_data: bytes) -> Dict[str, Any]: # noqa: C901 extended_header_offset = 8 if len(pte_data) > extended_header_offset + 32: try: - header = _ExtendedHeader.from_bytes(pte_data[extended_header_offset:]) + header = _ExtendedHeader.from_bytes( + pte_data[ + extended_header_offset : extended_header_offset + + _ExtendedHeader.EXPECTED_LENGTH + ] + ) if header.is_valid(): result["extended_header"] = { "magic": header.magic.decode("utf-8", errors="replace"), diff --git a/extension/flat_tensor/test/test_serialize.py b/extension/flat_tensor/test/test_serialize.py index 09226f024cc..9922ffa8920 100644 --- a/extension/flat_tensor/test/test_serialize.py +++ b/extension/flat_tensor/test/test_serialize.py @@ -285,20 +285,20 @@ def test_round_trip(self) -> None: TEST_DATA_PAYLOAD.named_data, deserialized_payload.named_data ) - def test_get_extended_header_reads_only_header_bytes(self) -> None: - # Slicing to the end of the data would copy the whole file. - class NoSliceToEnd(bytes): - def __getitem__(self, key): - if isinstance(key, slice) and key.stop is None: - raise AssertionError("sliced to the end of the data") - return super().__getitem__(key) - + 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 = NoSliceToEnd(serializer.serialize(TEST_DATA_PAYLOAD)) - header = _get_extended_header(serialized_data) + serialized_data = bytes(serializer.serialize(TEST_DATA_PAYLOAD)) - self.assertIsNotNone(header) - self.assertTrue(header.is_valid()) + with mock.patch.object( + FlatTensorHeader, "from_bytes", wraps=FlatTensorHeader.from_bytes + ) as from_bytes: + serializer.deserialize(Cord(serialized_data)) + self.assertIsNotNone(_get_extended_header(serialized_data)) + + self.assertEqual(from_bytes.call_count, 2) + for call in from_bytes.call_args_list: + self.assertEqual(len(call.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 From 5efeb57e017f6e33a3a2bbbb15d297628695ac29 Mon Sep 17 00:00:00 2001 From: HiroForYou Date: Fri, 9 Oct 2026 23:43:21 -0500 Subject: [PATCH 4/5] Address review: run the .ptd test in CI, tighten the .pte test, reuse the helper in the MLX inspector --- backends/mlx/pte_inspector.py | 39 ++++++++------------ exir/_serialize/test/test_program.py | 20 +++++----- extension/flat_tensor/serialize/serialize.py | 3 +- pytest.ini | 1 - 4 files changed, 26 insertions(+), 37 deletions(-) diff --git a/backends/mlx/pte_inspector.py b/backends/mlx/pte_inspector.py index cb2b9182ba8..1d96904e193 100644 --- a/backends/mlx/pte_inspector.py +++ b/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,29 +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 : extended_header_offset - + _ExtendedHeader.EXPECTED_LENGTH - ] - ) - 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 diff --git a/exir/_serialize/test/test_program.py b/exir/_serialize/test/test_program.py index 96e69c76e6b..09e591ee5c6 100644 --- a/exir/_serialize/test/test_program.py +++ b/exir/_serialize/test/test_program.py @@ -1301,19 +1301,19 @@ def test_from_bytes_with_more_data_than_necessary(self) -> None: self.assertEqual(eh.segment_data_size, EXAMPLE_SEGMENT_DATA_SIZE) def test_get_extended_header_reads_only_header_bytes(self) -> None: - # Slicing to the end of the data would copy the whole file. - class NoSliceToEnd(bytes): - def __getitem__(self, key): - if isinstance(key, slice) and key.stop is None: - raise AssertionError("sliced to the end of the data") - return super().__getitem__(key) - - pte_data = NoSliceToEnd(b"\x00" * 8 + EXAMPLE_HEADER_DATA + b"\x55" * 1024) - eh = _get_extended_header(pte_data) + # 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(eh.segment_base_offset, EXAMPLE_SEGMENT_BASE_OFFSET) + 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 diff --git a/extension/flat_tensor/serialize/serialize.py b/extension/flat_tensor/serialize/serialize.py index 5a1987dbc16..f6d696880e1 100644 --- a/extension/flat_tensor/serialize/serialize.py +++ b/extension/flat_tensor/serialize/serialize.py @@ -233,7 +233,6 @@ def __init__(self, data: Cord, alignment: Optional[int] = None) -> None: def _get_extended_header(flat_tensor_data: bytes) -> Optional[FlatTensorHeader]: """Returns the extended header of the flat_tensor data, if present and valid.""" try: - # flat_tensor_data[8:] would copy the whole file. eh = FlatTensorHeader.from_bytes( flat_tensor_data[8 : 8 + FlatTensorHeader.EXPECTED_LENGTH] ) @@ -396,7 +395,7 @@ def deserialize(self, blob: Cord) -> DataPayload: data = bytes(blob) - # Read header. Verify that it's valid. + # Read header. Verify that it's valid. data[8:] would copy the whole file. header = FlatTensorHeader.from_bytes( data[8 : 8 + FlatTensorHeader.EXPECTED_LENGTH] ) diff --git a/pytest.ini b/pytest.ini index 576ed2b78ce..a8b7aa1385b 100644 --- a/pytest.ini +++ b/pytest.ini @@ -55,7 +55,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 From d1059f0da0838a0566b181310813cfa95c54187d Mon Sep 17 00:00:00 2001 From: HiroForYou Date: Wed, 14 Oct 2026 15:03:37 -0500 Subject: [PATCH 5/5] Reuse _get_extended_header() in deserialize() and add __init__.py to the flat_tensor tests --- extension/flat_tensor/serialize/serialize.py | 9 ++++----- extension/flat_tensor/test/__init__.py | 0 extension/flat_tensor/test/test_serialize.py | 8 +++----- 3 files changed, 7 insertions(+), 10 deletions(-) create mode 100644 extension/flat_tensor/test/__init__.py diff --git a/extension/flat_tensor/serialize/serialize.py b/extension/flat_tensor/serialize/serialize.py index 1be635e4952..01c1b9a463b 100644 --- a/extension/flat_tensor/serialize/serialize.py +++ b/extension/flat_tensor/serialize/serialize.py @@ -217,6 +217,7 @@ 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: + # flat_tensor_data[8:] would copy the whole file. eh = FlatTensorHeader.from_bytes( flat_tensor_data[8 : 8 + FlatTensorHeader.EXPECTED_LENGTH] ) @@ -373,11 +374,9 @@ def deserialize(self, blob: Cord) -> DataPayload: data = bytes(blob) - # Read header. Verify that it's valid. data[8:] would copy the whole file. - header = FlatTensorHeader.from_bytes( - data[8 : 8 + FlatTensorHeader.EXPECTED_LENGTH] - ) - if not header.is_valid(): + # Read header. Verify that it's valid. + header = _get_extended_header(data) + if header is None: raise RuntimeError( "Flat tensor header is invalid. File is likely incorrect format or corrupt." ) diff --git a/extension/flat_tensor/test/__init__.py b/extension/flat_tensor/test/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/extension/flat_tensor/test/test_serialize.py b/extension/flat_tensor/test/test_serialize.py index fd8b0d5c0ee..c8ad6d0c9b9 100644 --- a/extension/flat_tensor/test/test_serialize.py +++ b/extension/flat_tensor/test/test_serialize.py @@ -34,7 +34,6 @@ _deserialize_to_flat_tensor, _FLAT_TENSOR_VERSION, _FLATBUFFER_ALIGNMENT, - _get_extended_header, FlatTensorConfig, FlatTensorHeader, FlatTensorSerializer, @@ -319,11 +318,10 @@ def test_reads_only_header_bytes(self) -> None: FlatTensorHeader, "from_bytes", wraps=FlatTensorHeader.from_bytes ) as from_bytes: serializer.deserialize(Cord(serialized_data)) - self.assertIsNotNone(_get_extended_header(serialized_data)) - self.assertEqual(from_bytes.call_count, 2) - for call in from_bytes.call_args_list: - self.assertEqual(len(call.args[0]), FlatTensorHeader.EXPECTED_LENGTH) + 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