Updating analysis utils (#63)

* Modify step stats util to compute stats for multiproc data
* Modify utils to handle multi-node data
* Modify notebook utils to handle multi-node data

Co-authored-by: Neelesh Dodda <[email protected]>

Author: vandanavk

smdebug/profiler/analysis/notebook_utils/heatmap.py

@@ -23,10 +23,12 @@ def __init__(
         select_dimensions=[".*CPU", ".*GPU", ".*Memory"],
         select_events=[".*"],
         plot_height=350,
+        show_workers=True,
     ):
 
         self.select_dimensions = select_dimensions
         self.select_events = select_events
+        self.show_workers = show_workers
         self.metrics_reader = metrics_reader
         self.available_dimensions = []
         self.available_events = []
@@ -54,13 +56,17 @@ def preprocess_system_metrics(self, events, system_metrics):
 
         # read all available system metric events and store them in dict
         for event in events:
-            if event.dimension not in system_metrics:
-                system_metrics[event.dimension] = {}
-                self.available_dimensions.append(event.dimension)
-            if event.name not in system_metrics[event.dimension]:
-                system_metrics[event.dimension][event.name] = []
+            if self.show_workers is True:
+                event_unique_id = f"{event.dimension}-nodeid:{str(event.node_id)}"
+            else:
+                event_unique_id = event.dimension
+            if event_unique_id not in system_metrics:
+                system_metrics[event_unique_id] = {}
+                self.available_dimensions.append(event_unique_id)
+            if event.name not in system_metrics[event_unique_id]:
+                system_metrics[event_unique_id][event.name] = []
                 self.available_events.append(event.name)
-            system_metrics[event.dimension][event.name].append([event.timestamp, event.value])
+            system_metrics[event_unique_id][event.name].append([event.timestamp, event.value])
 
         for dimension in system_metrics:
             for event in system_metrics[dimension]:
@@ -104,7 +110,14 @@ def create_plot(self):
         yaxis = {}
 
         # number of datapoints
-        self.width = self.system_metrics["CPUUtilization"]["total"].shape[0]
+        max_width = 0
+        for key in self.system_metrics.keys():
+            if key.startswith("CPUUtilization"):
+                width = self.system_metrics[key]["total"].shape[0]
+                if width >= max_width:
+                    max_width = width
+
+        self.width = max_width
 
         for dimension in self.filtered_dimensions:
             for event in self.filtered_events:
@@ -183,7 +196,14 @@ def update_data(self, current_timestamp):
                                 new_system_metrics[dimension][event],
                             ]
                         )
-            self.width = self.system_metrics["CPUUtilization"]["cpu0"].shape[0]
+            max_width = 0
+            for key in self.system_metrics.keys():
+                if key.startswith("CPUUtilization"):
+                    width = self.system_metrics[key]["cpu0"].shape[0]
+                    if width >= max_width:
+                        max_width = width
+
+            self.width = max_width
 
             tmp = []
             metric_names = []

smdebug/profiler/analysis/notebook_utils/metrics_histogram.py

@@ -29,7 +29,14 @@ def __init__(self, metrics_reader):
     @param select_metrics is array of metrics to be selected, Default ["cpu", "gpu"]
     """
 
-    def plot(self, starttime=0, endtime=None, select_dimensions=[".*"], select_events=[".*"]):
+    def plot(
+        self,
+        starttime=0,
+        endtime=None,
+        select_dimensions=[".*"],
+        select_events=[".*"],
+        show_workers=True,
+    ):
         if endtime == None:
             endtime = self.metrics_reader.get_timestamp_of_latest_available_file()
         all_events = self.metrics_reader.get_events(starttime, endtime)
@@ -39,6 +46,7 @@ def plot(self, starttime=0, endtime=None, select_dimensions=[".*"], select_event
         self.last_timestamp = endtime
         self.select_dimensions = select_dimensions
         self.select_events = select_events
+        self.show_workers = show_workers
         self.system_metrics = self.preprocess_system_metrics(
             all_events=all_events, system_metrics={}
         )
@@ -52,13 +60,17 @@ def preprocess_system_metrics(self, all_events=[], system_metrics={}):
 
         # read all available system metric events and store them in dict
         for event in all_events:
-            if event.dimension not in system_metrics:
-                system_metrics[event.dimension] = {}
-                self.available_dimensions.append(event.dimension)
-            if event.name not in system_metrics[event.dimension]:
-                system_metrics[event.dimension][event.name] = []
+            if self.show_workers is True:
+                event_unique_id = f"{event.dimension}-nodeid:{str(event.node_id)}"
+            else:
+                event_unique_id = event.dimension
+            if event_unique_id not in system_metrics:
+                system_metrics[event_unique_id] = {}
+                self.available_dimensions.append(event_unique_id)
+            if event.name not in system_metrics[event_unique_id]:
+                system_metrics[event_unique_id][event.name] = []
                 self.available_events.append(event.name)
-            system_metrics[event.dimension][event.name].append(event.value)
+            system_metrics[event_unique_id][event.name].append(event.value)
 
         # compute total utilization per event dimension
         for event_dimension in system_metrics:

smdebug/profiler/analysis/notebook_utils/timeline_charts.py

@@ -25,10 +25,12 @@ def __init__(
         select_dimensions=[".*"],
         select_events=[".*"],
         x=1000,
+        show_workers=True,
     ):
 
         self.select_dimensions = select_dimensions
         self.select_events = select_events
+        self.show_workers = show_workers
 
         # placeholder
         self.sources = {}
@@ -52,10 +54,16 @@ def __init__(
         self.start = 0  # replace with system_metrics_reader.get_first_available_timestamp()/1000000
         self.system_metrics = self.preprocess_system_metrics(events, system_metrics={})
 
-        if x < self.system_metrics["CPUUtilization"]["total"].shape[0]:
+        min_width = float("inf")
+        for key in self.system_metrics.keys():
+            if key.startswith("CPUUtilization"):
+                width = self.system_metrics[key]["total"].shape[0]
+                if width <= min_width:
+                    min_width = width
+        if x < min_width:
             self.width = x
         else:
-            self.width = self.system_metrics["CPUUtilization"]["total"].shape[0] - 1
+            self.width = min_width - 1
 
         # create plot
         self.create_plot()
@@ -64,13 +72,17 @@ def preprocess_system_metrics(self, events, system_metrics):
 
         # read all available system metric events and store them in dict
         for event in events:
-            if event.dimension not in system_metrics:
-                system_metrics[event.dimension] = {}
-                self.available_dimensions.append(event.dimension)
-            if event.name not in system_metrics[event.dimension]:
-                system_metrics[event.dimension][event.name] = []
+            if self.show_workers is True:
+                event_unique_id = f"{event.dimension}-nodeid:{str(event.node_id)}"
+            else:
+                event_unique_id = event.dimension
+            if event_unique_id not in system_metrics:
+                system_metrics[event_unique_id] = {}
+                self.available_dimensions.append(event_unique_id)
+            if event.name not in system_metrics[event_unique_id]:
+                system_metrics[event_unique_id][event.name] = []
                 self.available_events.append(event.name)
-            system_metrics[event.dimension][event.name].append([event.timestamp, event.value])
+            system_metrics[event_unique_id][event.name].append([event.timestamp, event.value])
 
         for dimension in system_metrics:
             for event in system_metrics[dimension]:
@@ -200,8 +212,15 @@ def create_plot(self):
     def find_time_annotations(self, indexes):
 
         if len(indexes) > 0:
-            begin_timestamp = self.system_metrics["CPUUtilization"]["total"][np.min(indexes), 0]
-            end_timestamp = self.system_metrics["CPUUtilization"]["total"][np.max(indexes), 0]
+            cpu_util = None
+            for key in self.system_metrics.keys():
+                if key.startswith("CPUUtilization"):
+                    width = self.system_metrics[key]["total"].shape[0]
+                    if cpu_util is None or np.min(indexes) <= width <= np.max(indexes):
+                        cpu_util = self.system_metrics[key]
+
+            begin_timestamp = cpu_util["total"][np.min(indexes), 0]
+            end_timestamp = cpu_util["total"][np.max(indexes), 0]
             total_time = end_timestamp - begin_timestamp
             print(
                 f"Selected timerange: {begin_timestamp + self.start} to {end_timestamp + self.start}"
@@ -281,8 +300,15 @@ def plot_dataloaders(self, events, begin_timestamp, end_timestamp):
     def plot_detailed_profiler_data(self, indexes):
 
         if len(indexes) > 0:
-            begin_timestamp = self.system_metrics["CPUUtilization"]["cpu0"][np.min(indexes), 0]
-            end_timestamp = self.system_metrics["CPUUtilization"]["cpu0"][np.max(indexes), 0]
+            cpu_util = None
+            for key in self.system_metrics.keys():
+                if key.startswith("CPUUtilization"):
+                    width = self.system_metrics[key]["cpu0"].shape[0]
+                    if cpu_util is None or np.min(indexes) <= width <= np.max(indexes):
+                        cpu_util = self.system_metrics[key]
+
+            begin_timestamp = cpu_util["cpu0"][np.min(indexes), 0]
+            end_timestamp = cpu_util["cpu0"][np.max(indexes), 0]
             print(
                 f"Selected timerange: {begin_timestamp + self.start} to {end_timestamp + self.start}"
             )
@@ -408,13 +434,22 @@ def update_data(self, current_timestamp):
                             event
                         ][self.system_metrics[dimension][event][:, 0].argsort()]
 
-            self.width = self.system_metrics["CPUUtilization"]["total"].shape[0] - 1
+            max_width = 0
+            cpu_util = None
+            for key in self.system_metrics.keys():
+                if key.startswith("CPUUtilization"):
+                    width = self.system_metrics[key]["total"].shape[0]
+                    if cpu_util is None or width >= max_width:
+                        max_width = width
+                        cpu_util = self.system_metrics[key]
+
+            self.width = max_width - 1
 
             if self.width > 1000:
-                min_value = self.system_metrics["CPUUtilization"]["total"][-1000, 0]
+                min_value = cpu_util["total"][-1000, 0]
             else:
-                min_value = self.system_metrics["CPUUtilization"]["total"][-self.width, 0]
-            max_value = self.system_metrics["CPUUtilization"]["total"][-1, 0]
+                min_value = cpu_util["total"][-self.width, 0]
+            max_value = cpu_util["total"][-1, 0]
 
             for figure in self.figures:
                 figure.x_range.start = int(min_value)