diff --git a/backends/mlx/pte_inspector.py b/backends/mlx/pte_inspector.py index 20ec9e9d9df..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,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 diff --git a/exir/_serialize/_program.py b/exir/_serialize/_program.py index 0f6adce08d6..8477cdfda7b 100644 --- a/exir/_serialize/_program.py +++ b/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: diff --git a/exir/_serialize/test/test_program.py b/exir/_serialize/test/test_program.py index dab3175f861..a8c73320da5 100644 --- a/exir/_serialize/test/test_program.py +++ b/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 diff --git a/extension/flat_tensor/serialize/serialize.py b/extension/flat_tensor/serialize/serialize.py index f6f6c7cedad..01c1b9a463b 100644 --- a/extension/flat_tensor/serialize/serialize.py +++ b/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." ) 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 35d12ca0552..c8ad6d0c9b9 100644 --- a/extension/flat_tensor/test/test_serialize.py +++ b/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 diff --git a/pytest.ini b/pytest.ini index efefecf1335..d34adc6ed5a 100644 --- a/pytest.ini +++ b/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