fix deviceProperties ingestion when combining multiple torch traces

src/aiu_trace_analyzer/ingest/ingestion.py

@@ -63,6 +63,20 @@ def set_ts_offset(self, offset):
     def is_torch_profile(self, data: dict) -> bool:
         return ("deviceProperties" in data)
 
+    def set_rank_pid(self, rank: int) -> None:
+        self.rank_pid = rank
+        if "deviceProperties" not in self.other_metadata:
+            return
+
+        if not isinstance(self.other_metadata["deviceProperties"], list) and \
+                "id" in self.other_metadata["deviceProperties"]:
+            self.other_metadata["deviceProperties"]["id"] = rank
+            return
+
+        for entry in self.other_metadata["deviceProperties"]:
+            if "id" in entry:
+                entry["id"] = rank
+
     def __iter__(self):
         raise NotImplementedError("Class %s doesn't implement __iter__" % (self.__class__.__name__))
 
@@ -99,7 +113,21 @@ def ftype_to_str(self, ftype):
         return prl[ftype]
 
     def combine_metadata(self, metadata):
+        # special treatment of deviceProperties list
+        current_dev_properties = self.other_metadata.pop("deviceProperties", [])
+        additional_dev_properties = metadata.pop("deviceProperties", [{}])
+        assert len(additional_dev_properties) == 1, \
+            f"Combining incoming metadata with multiple deviceProperties is not supported {additional_dev_properties}"
+
+        if "id" in additional_dev_properties[0]:
+            new_id = additional_dev_properties[0]["id"]
+            for entry in current_dev_properties:
+                if "id" in entry and entry["id"] == new_id:
+                    additional_dev_properties = []
+                    break
+
         self.other_metadata.update(metadata)
+        self.other_metadata["deviceProperties"] = current_dev_properties + additional_dev_properties
 
     def get_passthrough_meta(self) -> dict:
         return self.other_metadata
@@ -114,15 +142,22 @@ def updated_event(self, event: TraceEvent) -> TraceEvent:
         if "dur" in event:
             event["dur"] = float(event["dur"] * self.scale)
         dialect = GlobalIngestData.get_dialect(self.jobhash)
-        if self.rank_pid >= 0 and not isinstance(dialect, InputDialectTORCH):
-            event["pid"] = self.rank_pid
 
         the_args = "args"
         if "attr" in event:
             the_args = "attr"
         if the_args not in event:
             event[the_args] = {}
 
+        if isinstance(dialect, InputDialectTORCH):
+            event[the_args]["rank"] = self.rank_pid
+            if event["pid"] == 0:
+                event["pid"] = self.rank_pid
+        elif isinstance(dialect, InputDialectFLEX):
+            event[the_args]["rank"] = self.rank_pid
+            if self.rank_pid >= 0:
+                event["pid"] = self.rank_pid
+
         # make sure the pid/tid entries are numbers
         try:
             event["pid"] = int(event["pid"])
@@ -193,13 +228,10 @@ def _initialize_data(self, data_stream) -> None:
         metadata_keys = list(self.data.keys())
         for k in metadata_keys:
             if k in processing_keys:
+                # only copy entries that are not processed otherwise
                 continue
             self.other_metadata[k] = deepcopy(self.data.pop(k))
 
-        if "distributedInfo" in self.data and "rank" in self.data["distributedInfo"]:
-            self.rank_pid = self.data["distributedInfo"]["rank"]
-            aiulog.log(aiulog.DEBUG, "INGEST: Detected distributedInfo Rank", self.rank_pid)
-
         if "otherData" in self.data and \
             "Application" in self.data["otherData"] and \
                 "Acelyzer" in self.data["otherData"]["Application"]:
@@ -211,6 +243,11 @@ def _initialize_data(self, data_stream) -> None:
                     self.source_uri)
                 self.data["traceEvents"] = []
 
+        if "distributedInfo" in self.data and "rank" in self.data["distributedInfo"]:
+            self.set_rank_pid(self.data["distributedInfo"]["rank"])
+#            self.check_update_device_properties(self.rank_pid)
+            aiulog.log(aiulog.DEBUG, "INGEST: Detected distributedInfo Rank", self.rank_pid)
+
         if "traceEvents" in self.data:
             self.data = self.data["traceEvents"]
 

src/aiu_trace_analyzer/pipeline/tb_refinement.py

@@ -108,79 +108,87 @@ def update_event_data_heavy(self, event: TraceEvent) -> TraceEvent:
         event = self._update_for_collective(event)
         return event
 
-    def _queue_add_device(self, pid, ts, is_acc: bool = True):
-        if pid not in self.queues:
+    def _queue_add_device(self, dev_id, ts, is_acc: bool = True):
+        if dev_id not in self.queues:
             if is_acc:
-                self.queues[pid] = ("AIU Device"+str(pid), pid*2+1, "AIU", ts)
+                self.queues[dev_id] = ("AIU Device"+str(dev_id), dev_id*2+1, "AIU", ts)
                 self.exporter.add_device(
-                    pid,
+                    dev_id,
                     {"type": "AIU",
                      "name": "AIU",
                      "core": "PT Array"})
             else:
-                self.queues[pid] = ("Host"+str(pid), pid*2, "cpu", ts)
-
-    def update_event_data_light(self, event) -> TraceEvent:
-
-        def _cat_for_regular_event(event: TraceEvent) -> str:
-            # Use DmaI and DmaO with Sen to check memcpy event
-            if DmaI in event["name"] or DmaO in event["name"]:
-
-                # if no RDMA, a memory copy event
-                # else RDMA send and recv event
-                if RDMA not in event["name"]:
-                    return "gpu_memcpy"
-                else:
-                    return "user_annotation"
-
+                self.queues[dev_id] = ("Host"+str(dev_id), dev_id*2, "cpu", ts)
+
+    def _cat_for_regular_event(self, event: TraceEvent) -> str:
+        # Use DmaI and DmaO with Sen to check memcpy event
+        if DmaI in event["name"] or DmaO in event["name"]:
+            # if no RDMA, a memory copy event
+            # else RDMA send and recv event
+            if RDMA not in event["name"]:
+                return "gpu_memcpy"
             else:
-                if AllReduce in event["name"] and 'Cmpt Exec' in event["name"]:
-                    return "user_annotation"
-                else:
-                    return self.dialect.get("acc_category_kernel")
-
-        def _update_collective_event(event: TraceEvent) -> TraceEvent:
-            # if collective call block, change 'cat'
-            # and 'name' for communication tb calculation
-            if "args" in event and Coll_data_size in event["args"] and AllReduce in event["name"]:
-                event["cat"] = event["cat"] if "cat" in event else "user_annotation"
-                event["args"]["orig_name"] = event["name"]
-                event["name"] = "gloo:all_reduce"
-                event["external id"] = re.search(r"_(\d+)", event["args"]["orig_name"]).group(1)
+                return "user_annotation"
+
+        else:
+            if AllReduce in event["name"] and 'Cmpt Exec' in event["name"]:
+                return "user_annotation"
             else:
-                event["cat"] = event["cat"] if "cat" in event else "cpu_op"
-            return event
+                return self.dialect.get("acc_category_kernel")
+
+    @staticmethod
+    def _update_collective_event(event: TraceEvent) -> TraceEvent:
+        # if collective call block, change 'cat'
+        # and 'name' for communication tb calculation
+        if "args" in event and Coll_data_size in event["args"] and AllReduce in event["name"]:
+            event["cat"] = event["cat"] if "cat" in event else "user_annotation"
+            event["args"]["orig_name"] = event["name"]
+            event["name"] = "gloo:all_reduce"
+            event["external id"] = re.search(r"_(\d+)", event["args"]["orig_name"]).group(1)
+        else:
+            event["cat"] = event["cat"] if "cat" in event else "cpu_op"
+        return event
 
-        def _resolve_string_pids(pid) -> int:
-            if isinstance(pid, str):
-                aiulog.log(aiulog.WARN, f'TBR: input pid is string: {pid}')
-                try:
-                    return int(pid)
-                except (ValueError, TypeError):
-                    return hash(pid) % 10000 + 10000
-            return pid
+    @staticmethod
+    def _resolve_string_pids(pid) -> int:
+        if isinstance(pid, str):
+            aiulog.log(aiulog.WARN, f'TBR: input pid is string: {pid}')
+            try:
+                return int(pid)
+            except (ValueError, TypeError):
+                return hash(pid) % 10000 + 10000
+        return pid
+
+    @staticmethod
+    def _restore_pid_tid(event: TraceEvent) -> TraceEvent:
+        if "opid" in event["args"]:
+            event["pid"] = event["args"]["opid"]
+            event["args"].pop("opid")
+        if "otid" in event["args"]:
+            event["tid"] = event["args"]["otid"]
+            event["args"].pop("otid")
+        return event
+
+    def update_event_data_light(self, event) -> TraceEvent:
 
         self.dialect = GlobalIngestData.get_dialect(event["args"]["jobhash"])
 
         if self.dialect.get("NAME") == "TORCH":
-            if "opid" in event["args"]:
-                event["pid"] = event["args"]["opid"]
-                event["args"].pop("opid")
-            if "otid" in event["args"]:
-                event["tid"] = event["args"]["otid"]
-                event["args"].pop("otid")
+            event = RefinementContext._restore_pid_tid(event)
+            if PipelineContextTool.is_acc_event(event):
+                self._queue_add_device(event["args"]["rank"], event["ts"], is_acc=True)
             return event
 
-        pid = _resolve_string_pids(event["pid"])
+        pid = RefinementContext._resolve_string_pids(event["pid"])
 
         if PipelineContextTool.is_acc_event(event):
-            event["cat"] = _cat_for_regular_event(event)
+            event["cat"] = self._cat_for_regular_event(event)
 
-            event["args"]["device"] = pid
-            self._queue_add_device(pid, event["ts"], is_acc=True)
+            event["args"]["device"] = event["args"]["rank"]
+            self._queue_add_device(event["args"]["rank"], event["ts"], is_acc=True)
 
         else:
-            event = _update_collective_event(event)
+            event = self._update_collective_event(event)
             event["pid"] = pid + 1000
 
             self._queue_add_device(pid, event["ts"], is_acc=False)