Pytorch profiling - Break down python profiling, fixes in detailed profiling

1/ Break down python profiling steps into 2 parts (pre-forward to backend and backwardend to next preforward)
2/ detailed profiling enable profiler api changed for pytorch 1.6 . Fixed to check the version and then pass the parameter
3/ Function to generate count_parameter
4/ Enable cProfile to emit total time rather than just cpu time

Author: Vikas-kum

smdebug/core/locations.py

@@ -164,8 +164,11 @@ def get_python_profiling_stats_dir(
         step,
         start_time_since_epoch_in_micros,
         end_time_since_epoch_in_micros,
+        step_phase="",
     ):
         node_id = get_node_id()
+        if step_phase:
+            step = str(step) + "-" + step_phase
         folder_name = "{0}_{1}_{2}_{3}".format(
             start_time_since_epoch_in_micros, end_time_since_epoch_in_micros, node_id, step
         )

smdebug/profiler/analysis/python_stats_reader.py

@@ -28,6 +28,12 @@ def load_python_profile_stats(self):
         """Load the python profile stats. To be implemented in subclass.
         """
 
+    def _get_step_stepphase(self, step_phase_str):
+        splits = step_phase_str.split("-", 1)
+        step = splits[0]
+        step_phase = splits[1] if len(splits) > 1 else "full"
+        return step, step_phase
+
 
 class S3PythonStatsReader(PythonStatsReader):
     """Higher level stats reader to download python stats from s3.
@@ -93,7 +99,8 @@ def load_python_profile_stats(self):
             with open(stats_file_path, "wb") as f:
                 f.write(object_data)
 
-            start_time, end_time, node_id, step = stats_dir.split("_")
+            start_time, end_time, node_id, step_phase_str = stats_dir.split("_")
+            step, step_phase = self._get_step_stepphase(step_phase_str)
             python_profile_stats.append(
                 StepPythonProfileStats(
                     profiler_name,
@@ -102,6 +109,7 @@ def load_python_profile_stats(self):
                     float(end_time),
                     node_id,
                     stats_file_path,
+                    step_phase,
                 )
             )
         python_profile_stats.sort(
@@ -131,7 +139,9 @@ def load_python_profile_stats(self):
         """
         python_profile_stats = []
         for python_stat_dir in os.listdir(self.profile_dir):
-            start_time, end_time, node_id, step = python_stat_dir.split("_")
+            start_time, end_time, node_id, step_phase_str = python_stat_dir.split("_")
+            step, step_phase = self._get_step_stepphase(step_phase_str)
+
             stats_dir = os.path.join(self.profile_dir, python_stat_dir)
             if os.path.isfile(os.path.join(stats_dir, CPROFILE_STATS_FILENAME)):
                 profiler_name = CPROFILE_NAME

smdebug/profiler/analysis/utils/python_profile_analysis_utils.py

@@ -11,6 +11,7 @@ def __init__(
         end_time_since_epoch_in_micros,
         node_id,
         stats_path,
+        step_phase="",
     ):
         """Class that represents the metadata for profiling on a specific step (or before step 0).
         Used so that users can easily filter through which steps they want profiling stats of.
@@ -27,6 +28,7 @@ def __init__(
         self.end_time_since_epoch_in_micros = end_time_since_epoch_in_micros
         self.node_id = node_id
         self.stats_path = stats_path
+        self.step_phase = step_phase
 
     def in_time_interval(self, start_time_since_epoch_in_micros, end_time_since_epoch_in_micros):
         """Returns whether this step is in the provided time interval.

smdebug/profiler/python_profiler.py

@@ -15,9 +15,6 @@
 
 
 class PythonProfiler:
-    name = ""  # placeholder
-    stats_filename = ""  # placeholder
-
     def __init__(self, base_folder, framework):
         """Higher level class to manage execution of python profiler, dumping of python stats, and retrieval
         of stats based on time or step intervals.
@@ -51,6 +48,7 @@ def _reset_profiler(self):
         """Reset attributes to defaults
         """
         self._step, self._start_time_since_epoch_in_micros, self._is_profiling = None, None, False
+        self._start_stepphase = ""
 
     def _enable_profiler(self):
         """Enable the profiler (to be implemented in subclass, where the actual profiler is defined).
@@ -64,16 +62,24 @@ def _dump_stats(self, stats_path):
         """Dump the stats to the provided path (to be implemented in subclass, where the actual profiler is defined).
         """
 
-    def start_profiling(self, start_step=-1):
+    def _name(self):
+        return "default"
+
+    def _stats_filename(self):
+        # this is default value
+        return "python_stats"
+
+    def start_profiling(self, start_step=-1, step_phase=""):
         """Start the python profiler with the provided start step.
         If start step is -1, then this is profiling from import time to step 0.
         """
         self._step = start_step
         self._start_time_since_epoch_in_micros = time.time() * CONVERT_TO_MICROSECS
         self._is_profiling = True
+        self._start_stepphase = step_phase
         self._enable_profiler()
 
-    def stop_profiling(self):
+    def stop_profiling(self, step_phase=""):
         """Stop the python profiler.
         Dump the python stats for this step with a file path dependent on the base folder, framework, time and step.
         Append a record of this step's profiling with the corresponding metadata.
@@ -85,15 +91,19 @@ def stop_profiling(self):
         self._disable_profiler()
 
         current_time_since_epoch_in_micros = time.time() * CONVERT_TO_MICROSECS
+        step_phase_string = ""
+        if self._start_stepphase != "" or step_phase != "":
+            step_phase_string = f"{self._start_stepphase}-{step_phase}"
         stats_dir = TraceFileLocation.get_python_profiling_stats_dir(
             self._base_folder,
             self._framework,
-            self.name,
+            self._name(),
             self._step,
             self._start_time_since_epoch_in_micros,
             current_time_since_epoch_in_micros,
+            step_phase_string,
         )
-        self._dump_stats(os.path.join(stats_dir, self.stats_filename))
+        self._dump_stats(os.path.join(stats_dir, self._stats_filename()))
 
         self._reset_profiler()
 
@@ -110,14 +120,26 @@ class cProfilePythonProfiler(PythonProfiler):
     This is also the default Python profiler used if profiling is enabled.
     """
 
-    name = "cProfile"
-    stats_filename = "python_stats"
-
     def _reset_profiler(self):
         """Reset profiler and corresponding attributes to defaults
         """
         super()._reset_profiler()
-        self._profiler = cProfileProfiler()
+        self._profiler = cProfileProfiler(self._total_time)
+
+    def _name(self):
+        return "cProfile"
+
+    def _total_time(self):
+        times = os.times()
+        return times.elapsed
+
+    def _off_cpu_time(self):
+        times = os.times()
+        return times.elapsed - (times.system + times.user)
+
+    def _stats_filename(self):
+        # this is default value
+        return "python_stats"
 
     def _enable_profiler(self):
         """Enable the cProfile profiler.
@@ -135,20 +157,22 @@ def _dump_stats(self, stats_file_path):
         get_logger("smdebug-profiler").info(f"Dumping cProfile stats to {stats_file_path}.")
         pstats.Stats(self._profiler).dump_stats(stats_file_path)
 
-
 class PyinstrumentPythonProfiler(PythonProfiler):
     """Higher level class to oversee profiling specific to Pyinstrument, a third party Python profiler.
     """
 
-    name = "pyinstrument"
-    stats_filename = "python_stats.json"
-
     def _reset_profiler(self):
         """Reset profiler and corresponding attributes to defaults
         """
         super()._reset_profiler()
         self._profiler = PyinstrumentProfiler()
 
+    def _name(self):
+        return "pyinstrument"
+
+    def _stats_filename(self):
+        return "python_stats.json"
+
     def _enable_profiler(self):
         """Enable the pyinstrument profiler.
         """