Adding script to read the events from trace file (aws#246)

* adding script to read the events from tensorboard trace.

* Refactored code
Introduced base event parser and the child classes TFEventProfiler
and SMTFEventProfiler

Co-authored-by: Vikas Kumar <[email protected]>

Author: leleamol

config/tests.sh

@@ -51,6 +51,7 @@ export REPORT_DIR=$OUT_DIR/pytest_reports
 python -m pytest ${code_coverage_smdebug:+--cov=./ --cov-append} -v -W=ignore --durations=50 --html=$REPORT_DIR/report_analysis.html --self-contained-html tests/analysis
 
 run_for_framework core
+run_for_framework profiler
 
 if [ "$run_pytest_xgboost" = "enable" ] ; then
     run_for_framework xgboost

smdebug/profiler/__init__.py

@@ -0,0 +1,3 @@
+# Local
+from .tf_profiler_parser import SMTFProfilerEvents, TFProfilerEvents
+from .trace_event_file_parser import TraceEvent, TraceEventParser

smdebug/profiler/tf_profiler_parser.py

@@ -0,0 +1,78 @@
+# Standard Library
+import json
+
+# First Party
+from smdebug.profiler.trace_event_file_parser import ProcessInfo, TraceEventParser
+
+
+class SMTFProfilerEvents(TraceEventParser):
+    def __init__(self, trace_file):
+        self._trace_json_file = trace_file
+        super().__init__()
+        self.read_trace_file()
+
+    def _populate_start_time(self, event):
+        event_args = event["args"] if "args" in event else None
+        if self._start_time_known is False:
+            if event_args is None:
+                return
+            if "start_time_since_epoch_in_micros" in event_args:
+                self._start_timestamp = event_args["start_time_since_epoch_in_micros"]
+                self._start_time_known = True
+                self.logger.info(f"Start time for events in uSeconds = {self._start_timestamp}")
+
+    # TODO implementation of below would be changed to support streaming file and incomplete json file
+    def read_trace_file(self):
+        try:
+            with open(self._trace_json_file) as json_data:
+                trace_json_data = json.load(json_data)
+        except Exception as e:
+            self.logger.error(
+                f"Can't open TF trace file {self._trace_json_file}: Exception {str(e)}"
+            )
+            return
+
+        for event in trace_json_data:
+            self._read_event(event)
+
+
+class TFProfilerEvents(TraceEventParser):
+    def __init__(self, trace_file):
+        self._trace_json_file = trace_file
+        super().__init__()
+        self.read_trace_file()
+
+    def _populate_thread_info_for_metaevent(self, event):
+        if event["name"] == "thread_name":
+            name = event["args"]["name"]
+            t_id = event["tid"]
+            pid = event["pid"]
+            if pid not in self._processes:
+                self.logger.warn(
+                    f"Did not find matching process for pid {pid}. Creating a process with name 'Unknown'"
+                )
+                self._processes[pid] = ProcessInfo(pid, "Unknown")
+            self._processes[pid].add_thread(t_id, name)
+
+    def _populate_start_time(self, event):
+        # TODO, not sure if we can implement this right now
+        return
+
+    def read_trace_file(self):
+        try:
+            with open(self._trace_json_file) as json_data:
+                trace_json_data = json.load(json_data)
+        except Exception as e:
+            self.logger.error(
+                f"Can't open TF trace file {self._trace_json_file}: Exception {str(e)} "
+            )
+            return
+        if "traceEvents" not in trace_json_data:
+            self.logger.error(
+                f"The TF trace file {self._trace_json_file} does not contain traceEvents"
+            )
+            return
+        trace_events_json = trace_json_data["traceEvents"]
+
+        for event in trace_events_json:
+            self._read_event(event)

smdebug/profiler/trace_event_file_parser.py

@@ -0,0 +1,127 @@
+# First Party
+from smdebug.core.logger import get_logger
+
+
+class ThreadInfo:
+    def __init__(self, tid, thread_name):
+        self.tid = tid
+        self.thread_name = thread_name
+
+
+class ProcessInfo:
+    def __init__(self, id, name):
+        self.id = id
+        self.name = name
+        self._threads = dict()
+
+    def add_thread(self, threadid, thread_name):
+        self._threads[threadid] = ThreadInfo(threadid, thread_name)
+
+    def get_thread_info(self, threadid):
+        return self._threads[threadid]
+
+
+class TraceEvent:
+    def __init__(self, ts, name, dur, pid, tid, event_args):
+        self.start_time = ts
+        self.event_name = name
+        self.duration = dur
+        self.end_time = self.start_time + self.duration
+        self.pid = pid
+        self.tid = tid
+        self.event_args = event_args
+
+
+class TraceEventParser:
+    def __init__(self):
+        self._processes = dict()
+        self._trace_events = list()
+        self._start_timestamp = 0
+        self._start_time_known = False
+        # The timestamp in trace events are in micro seconds, we multiply by 1000 to convert to ns
+        self._timescale_multiplier_for_ns = 1000
+        self.logger = get_logger("smdebug-profiler")
+
+    def read_trace_file(self):
+        pass
+
+    def _populate_process_info_for_metaevent(self, event):
+        id = event["pid"]
+        if event["name"] == "process_name":
+            name = event["args"]["name"] if "name" in event["args"] else "Unknown"
+            self._processes[id] = ProcessInfo(id, name)
+
+    def _populate_thread_info_for_metaevent(self, event):
+        pass
+
+    def _populate_start_time(self, event):
+        pass
+
+    def _read_event(self, event):
+        if "ph" not in event:
+            self.logger.error(f"In correctly formatted trace file. The 'ph' field is not present")
+            return
+        phase_type = event["ph"]
+        if phase_type == "M":
+            self._populate_process_info_for_metaevent(event)
+            self._populate_thread_info_for_metaevent(event)
+            self._populate_start_time(event)
+
+        if phase_type == "X":
+            # In nano seconds
+            start_time = (event["ts"] + self._start_timestamp) * self._timescale_multiplier_for_ns
+            # In nano seconds
+            dur = event["dur"] * self._timescale_multiplier_for_ns
+            name = event["name"]
+            id = event["pid"]
+            tid = event["tid"] if "tid" in event else "0"
+            event_args = event["args"] if "args" in event else None
+            t_event = TraceEvent(start_time, name, dur, id, tid, event_args)
+            self._trace_events.append(t_event)
+
+    def get_all_events(self):
+        return self._trace_events
+
+    def get_events_start_time_sorted(self):
+        return sorted(self._trace_events, key=lambda x: x.start_time)
+
+    def get_events_end_time_sorted(self):
+        return sorted(self._trace_events, key=lambda x: x.end_time)
+
+    """
+    Return the events that are in progress at the specified timestamp.
+    Performance of this function can be improved by implementing interval tree.
+    """
+
+    def get_events_at(self, timestamp_in_nanoseconds):
+        result_events = list()
+        for x_event in self._trace_events:
+            if x_event.start_time <= timestamp_in_nanoseconds <= x_event.end_time:
+                result_events.append(x_event)
+        return result_events
+
+    """
+    Return the events that have started and completed within the given start and end time boundaries.
+    The events that are in progress during these boundaries are not included.
+    """
+
+    def get_events_within_range(self, start_time, end_time):
+        result_events = list()
+        for x_event in self._trace_events:
+            if start_time <= x_event.start_time and end_time >= x_event.end_time:
+                result_events.append(x_event)
+        return result_events
+
+    def get_process_info(self, process_id):
+        return self._processes[process_id]
+
+    def get_processes(self):
+        return self._processes
+
+    # TODO
+    def get_events_for_process(self, pid, start_time, end_time):
+        pass
+
+    # TODO
+    def get_events_for_thread(self, tid, start_time, end_time):
+        pass