Skip to content

Commit 3fd434b

Browse files
committed
raise error when trying to save an etrecord missing essential info
Pull Request resolved: #13143 as title ghstack-source-id: 301624419 @exported-using-ghexport Differential Revision: [D79687142](https://our.internmc.facebook.com/intern/diff/D79687142/)
1 parent a8fe653 commit 3fd434b

File tree

2 files changed

+56
-1
lines changed

2 files changed

+56
-1
lines changed

devtools/etrecord/_etrecord.py

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,22 @@ def __init__(
7070
_reference_outputs: Optional[Dict[str, List[ProgramOutput]]] = None,
7171
_representative_inputs: Optional[List[ProgramInput]] = None,
7272
):
73+
"""
74+
Please do not construct an ETRecord object directly.
75+
76+
If you want to create an ETRecord for logging AOT information to further analysis, please mark `generate_etrecord`
77+
as True in your export api, and get the ETRecord object from the `ExecutorchProgramManager`.
78+
For exmaple:
79+
```python
80+
exported_program = torch.export.export(model, inputs)
81+
edge_program = to_edge_transform_and_lower(exported_program, generate_etrecord=True)
82+
executorch_program = edge_program.to_executorch()
83+
etrecord = executorch_program.get_etrecord()
84+
```
85+
86+
If user need to create an ETRecord manually, please use the `create_etrecord` function.
87+
"""
88+
7389
self.exported_program = exported_program
7490
self.export_graph_id = export_graph_id
7591
self.edge_dialect_program = edge_dialect_program
@@ -81,15 +97,25 @@ def __init__(
8197

8298
def save(self, path: Union[str, os.PathLike, BinaryIO, IO[bytes]]) -> None:
8399
"""
84-
Serialize and save the ETRecord to the specified path.
100+
Serialize and save the ETRecord to the specified path for use in Inspector. The ETRecord
101+
should contains at least edge dialect program and executorch program information for further
102+
analysis, otherwise it will raise an exception.
85103
86104
Args:
87105
path: Path where the ETRecord file will be saved to.
106+
107+
Raises:
108+
RuntimeError: If the ETRecord does not contain essential information for Inpector.
88109
"""
89110
if isinstance(path, (str, os.PathLike)):
90111
# pyre-ignore[6]: In call `os.fspath`, for 1st positional argument, expected `str` but got `Union[PathLike[typing.Any], str]`
91112
path = os.fspath(path)
92113

114+
if not (self.edge_dialect_program and self._debug_handle_map):
115+
raise RuntimeError(
116+
"ETRecord must contain edge dialect program and executorch program to be saved"
117+
)
118+
93119
etrecord_zip = ZipFile(path, "w")
94120

95121
try:

devtools/etrecord/tests/etrecord_test.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1462,3 +1462,32 @@ def test_update_apis_and_save_parse(self):
14621462
custom_outputs["forward"], parsed_etrecord._reference_outputs["forward"]
14631463
):
14641464
self.assertTrue(torch.equal(expected[0], actual[0]))
1465+
1466+
def test_save_missing_essential_info(self):
1467+
def expected_runtime_error(etrecord, etrecord_path):
1468+
with self.assertRaises(RuntimeError) as context:
1469+
etrecord.save(etrecord_path)
1470+
1471+
self.assertIn(
1472+
"ETRecord must contain edge dialect program and executorch program to be saved",
1473+
str(context.exception),
1474+
)
1475+
1476+
"""Test that save raises RuntimeError when essential info is missing."""
1477+
_, edge_output, et_output = self.get_test_model()
1478+
1479+
etrecord = ETRecord()
1480+
1481+
with tempfile.TemporaryDirectory() as tmpdirname:
1482+
etrecord_path = tmpdirname + "/etrecord_no_edge.bin"
1483+
1484+
expected_runtime_error(etrecord, etrecord_path)
1485+
etrecord.add_edge_dialect_program(edge_output)
1486+
1487+
# Should raise runtime error due to missing executorch program related info
1488+
expected_runtime_error(etrecord, etrecord_path)
1489+
1490+
etrecord.add_executorch_program(et_output)
1491+
1492+
# All essential components are now present, so save should succeed
1493+
etrecord.save(etrecord_path)

0 commit comments

Comments
 (0)