From 4088ba16f61140ac815d866bcf50684685880746 Mon Sep 17 00:00:00 2001 From: Utkarsh Bahuguna Date: Mon, 12 Oct 2026 00:19:43 +0530 Subject: [PATCH] Normalize ETRecord map keys in Inspector --- devtools/inspector/_inspector.py | 27 +++++++++++---- devtools/inspector/tests/inspector_test.py | 39 ++++++++++++++++++++++ 2 files changed, 59 insertions(+), 7 deletions(-) diff --git a/devtools/inspector/_inspector.py b/devtools/inspector/_inspector.py index 1d5a52715c9..a15dcb0d86c 100644 --- a/devtools/inspector/_inspector.py +++ b/devtools/inspector/_inspector.py @@ -1132,15 +1132,28 @@ def _consume_etrecord(self, reference_graph_name) -> None: return # (1) Debug Handle Symbolification + # These maps are keyed by instruction id: ints on an ETRecord object + # from get_etrecord(), strs once it went through save() and + # parse_etrecord() (JSON). _gen_resolve_debug_handles looks up strs. + def with_str_keys(mapping: Optional[Dict]) -> Optional[Dict]: + if mapping is None: + return None + return {str(key): value for key, value in mapping.items()} + + debug_handle_map = with_str_keys(self._etrecord._debug_handle_map[FORWARD]) + delegate_map = ( + with_str_keys(self._etrecord._delegate_map[FORWARD]) + if self._etrecord._delegate_map is not None + else None + ) + instruction_id_to_num_outs_map = with_str_keys( + self._etrecord._instruction_id_to_num_outs_map[FORWARD] + ) for event_block in self.event_blocks: event_block._gen_resolve_debug_handles( - self._etrecord._debug_handle_map[FORWARD], - ( - self._etrecord._delegate_map[FORWARD] - if self._etrecord._delegate_map is not None - else None - ), - self._etrecord._instruction_id_to_num_outs_map[FORWARD], + debug_handle_map, + delegate_map, + instruction_id_to_num_outs_map, ) # (2) Event Metadata Association diff --git a/devtools/inspector/tests/inspector_test.py b/devtools/inspector/tests/inspector_test.py index 4c59190650c..320ec2f9e21 100644 --- a/devtools/inspector/tests/inspector_test.py +++ b/devtools/inspector/tests/inspector_test.py @@ -390,6 +390,45 @@ def test_inspector_get_exported_program(self): ) ) + def test_inspector_with_etrecord_object_resolves_debug_handles(self): + _, _, et_output = TestETRecord().get_test_model(generate_etrecord=True) + etrecord = et_output.get_etrecord() + # An ETRecord object from get_etrecord() keys its maps by int instruction id + instruction_id = next(iter(etrecord._debug_handle_map["forward"])) + self.assertIsInstance(instruction_id, int) + + def gen_event_blocks(*args, **kwargs): + return [ + EventBlock( + name=EVENT_BLOCK_NAME, + events=[Event(name="op", _instruction_id=instruction_id)], + ) + ] + + with tempfile.TemporaryDirectory() as tmp_dir: + etrecord_path = os.path.join(tmp_dir, "etrecord.bin") + etrecord.save(etrecord_path) + + with patch.object( + _inspector, "gen_etdump_object", return_value=None + ), patch.object( + EventBlock, "_gen_from_etdump", side_effect=gen_event_blocks + ): + inspector_from_object = Inspector( + etdump_path=ETDUMP_PATH, etrecord=etrecord + ) + inspector_from_file = Inspector( + etdump_path=ETDUMP_PATH, etrecord=etrecord_path + ) + + event_from_object = inspector_from_object.event_blocks[0].events[0] + event_from_file = inspector_from_file.event_blocks[0].events[0] + self.assertIsNotNone(event_from_file.debug_handles) + # Passing the ETRecord object directly must give the same result as + # passing the path of the saved ETRecord + self.assertEqual(event_from_object.debug_handles, event_from_file.debug_handles) + self.assertEqual(event_from_object.op_types, event_from_file.op_types) + def test_populate_debugging_related_fields_raises_for_inconsistent_events(self): ret_event: Event = Event( name="event",