From e5131ef210fa7ef6d09a18efb7697d9a8cb4e870 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Tue, 5 Aug 2025 16:32:43 -0700 Subject: [PATCH] raise error when trying to save an etrecord missing essential info as title Differential Revision: [D79687142](https://our.internmc.facebook.com/intern/diff/D79687142/) [ghstack-poisoned] --- devtools/etrecord/_etrecord.py | 8 +++++++ devtools/etrecord/tests/etrecord_test.py | 29 ++++++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/devtools/etrecord/_etrecord.py b/devtools/etrecord/_etrecord.py index 6c8a55d6220..ac4de37ec85 100644 --- a/devtools/etrecord/_etrecord.py +++ b/devtools/etrecord/_etrecord.py @@ -85,11 +85,19 @@ def save(self, path: Union[str, os.PathLike, BinaryIO, IO[bytes]]) -> None: Args: path: Path where the ETRecord file will be saved to. + + Raises: + RuntimeError: If the ETRecord does not contain essential information for Inpector. """ if isinstance(path, (str, os.PathLike)): # pyre-ignore[6]: In call `os.fspath`, for 1st positional argument, expected `str` but got `Union[PathLike[typing.Any], str]` path = os.fspath(path) + if not (self.edge_dialect_program and self.graph_map and self._debug_handle_map): + raise RuntimeError( + "ETRecord must contain edge dialect program, graph map, and debug handle map to be saved." + ) + etrecord_zip = ZipFile(path, "w") try: diff --git a/devtools/etrecord/tests/etrecord_test.py b/devtools/etrecord/tests/etrecord_test.py index dbd7fdfb776..88f079f0e62 100644 --- a/devtools/etrecord/tests/etrecord_test.py +++ b/devtools/etrecord/tests/etrecord_test.py @@ -1545,3 +1545,32 @@ def test_update_apis_and_save_parse(self): custom_outputs["forward"], parsed_etrecord._reference_outputs["forward"] ): self.assertTrue(torch.equal(expected[0], actual[0])) + + def test_save_missing_essential_info(self): + def expected_runtime_error(etrecord, etrecord_path): + with self.assertRaises(RuntimeError) as context: + etrecord.save(etrecord_path) + + self.assertIn( + "ETRecord must contain edge dialect program, graph map, and debug handle map to be saved", + str(context.exception), + ) + + """Test that save raises RuntimeError when essential info is missing.""" + _, edge_output, et_output = self.get_test_model() + + etrecord = ETRecord() + + with tempfile.TemporaryDirectory() as tmpdirname: + etrecord_path = tmpdirname + "/etrecord_no_edge.bin" + + expected_runtime_error(etrecord, etrecord_path) + etrecord.add_edge_dialect_program(edge_output) + + # Should raise runtime error due to missing executorch program related info + expected_runtime_error(etrecord, etrecord_path) + + etrecord.add_executorch_program(et_output) + + # All essential components are now present, so save should succeed + etrecord.save(etrecord_path)