Fix Inspector losing debug handles when given an ETRecord object

devtools/inspector/_inspector.py

@@ -8,6 +8,7 @@
 # pyre-unsafe
 
 import dataclasses
+import json
 import logging
 import sys
 import warnings
@@ -932,7 +933,7 @@ def _gen_resolve_debug_handles(
         self,
         handle_map: Dict[str, List[int]],
         delegate_map: Optional[Dict[str, DelegateMetadata]] = None,
-        instruction_id_to_num_outs_map: Dict[int, int] = None,
+        instruction_id_to_num_outs_map: Dict[str, int] = None,
     ):
         """
         Given mappings from instruction id to debug handles, populate the
@@ -1132,16 +1133,25 @@ def _consume_etrecord(self, reference_graph_name) -> None:
             return
 
         # (1) Debug Handle Symbolification
-        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],
+        if self.event_blocks:
+            # An ETRecord object from get_etrecord() has int keys (at every level)
+            # and tuples; save()/parse_etrecord() gives the JSON shape (str keys,
+            # lists) that _gen_resolve_debug_handles expects, so use that for both.
+            debug_handle_map = json.loads(json.dumps(self._etrecord._debug_handle_map))
+            delegate_map = (
+                json.loads(json.dumps(self._etrecord._delegate_map))[FORWARD]
+                if self._etrecord._delegate_map is not None
+                else None
             )
+            num_outs_map = json.loads(
+                json.dumps(self._etrecord._instruction_id_to_num_outs_map)
+            )
+            for event_block in self.event_blocks:
+                event_block._gen_resolve_debug_handles(
+                    debug_handle_map[FORWARD],
+                    delegate_map,
+                    num_outs_map[FORWARD],
+                )
 
         # (2) Event Metadata Association
         self.op_graph_dict = gen_graphs_from_etrecord(

devtools/inspector/tests/inspector_test.py

@@ -390,6 +390,106 @@ 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_inspector_with_delegated_etrecord_object_matches_saved_file(self):
+        _, _, et_output = TestETRecord().get_test_model(generate_etrecord=True)
+        etrecord = et_output.get_etrecord()
+        instruction_id = next(iter(etrecord._debug_handle_map["forward"]))
+        # Shape produced by DelegateMappingBuilder: int keys, tuple values
+        etrecord._delegate_map = {
+            "forward": {
+                instruction_id: {
+                    "name": "TestBackend",
+                    "delegate_map": {0: (1, 2), 5: (3,)},
+                }
+            }
+        }
+
+        def gen_event_blocks(*args, **kwargs):
+            return [
+                EventBlock(
+                    name=EVENT_BLOCK_NAME,
+                    events=[
+                        Event(name="op", _instruction_id=instruction_id),
+                        Event(
+                            name="delegated",
+                            _instruction_id=instruction_id,
+                            delegate_debug_identifier=0,
+                        ),
+                        # not in the delegate map
+                        Event(
+                            name="unmapped",
+                            _instruction_id=instruction_id,
+                            delegate_debug_identifier=7,
+                        ),
+                    ],
+                )
+            ]
+
+        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
+                )
+
+        events_from_object = inspector_from_object.event_blocks[0].events
+        events_from_file = inspector_from_file.event_blocks[0].events
+        self.assertEqual(events_from_file[1].debug_handles, [1, 2])
+        for from_object, from_file in zip(events_from_object, events_from_file):
+            self.assertEqual(from_object.debug_handles, from_file.debug_handles)
+            self.assertEqual(from_object.num_outputs, from_file.num_outputs)
+            self.assertEqual(
+                from_object.delegate_backend_name, from_file.delegate_backend_name
+            )
+
     def test_populate_debugging_related_fields_raises_for_inconsistent_events(self):
         ret_event: Event = Event(
             name="event",