PY-74184 Add tool window support with flame graph functionality and thread information support for the profiler

GitOrigin-RevId: 8ac062fbc3112cb01d810863b8328dda43c35040

python/helpers/profiler/load_pstat.py

@@ -5,7 +5,14 @@ from _prof_imports import IS_PY3K
 from _prof_imports import ProfilerResponse
 from _prof_imports import TBinaryProtocolFactory
 from _prof_imports import serialize
-from prof_util import stats_to_response
+from prof_util import ystats_to_response, stats_to_response
+
+try:
+    import yappi
+    yappi_installed = True
+except ImportError:
+    yappi_installed = False
+
 
 if __name__ == '__main__':
 
@@ -17,8 +24,12 @@ if __name__ == '__main__':
         import vmprof_profiler
         vmprof_profiler.tree_stats_to_response(filename, m)
     else:
-        stats = pstats.Stats(filename)
-        stats_to_response(stats.stats, m)
+        if yappi_installed:
+            ystats = yappi.YFuncStats(filename)
+            ystats_to_response(ystats, m)
+        else:
+            stats = pstats.Stats(filename)
+            stats_to_response(stats.stats, m)
 
     data = serialize(m, TBinaryProtocolFactory())
 

python/helpers/profiler/prof_util.py

@@ -18,7 +18,7 @@ except NameError:
             glob = sys._getframe().f_back.f_globals
         if loc is None:
             loc = glob
-    
+
         # It seems that the best way is using tokenize.open(): http://code.activestate.com/lists/python-dev/131251/
         import tokenize
         stream = tokenize.open(file)  # @UndefinedVariable
@@ -26,9 +26,9 @@ except NameError:
             contents = stream.read()
         finally:
             stream.close()
-    
+
         # execute the script (note: it's important to compile first to have the filename set in debug mode)
-        exec(compile(contents+"\n", file, 'exec'), glob, loc)
+        exec(compile(contents + "\n", file, 'exec'), glob, loc)
 
 
 def save_main_module(file, module_name):
@@ -86,17 +86,19 @@ def get_snapshot_basepath(basepath, local_temp_dir):
     if basepath is None:
         basepath = 'snapshot'
     if local_temp_dir:
-        basepath = os.path.join(tempfile.gettempdir(), os.path.basename(basepath.replace('\\', '/')))
+        basepath = os.path.join(tempfile.gettempdir(),
+                                os.path.basename(basepath.replace('\\', '/')))
     return basepath
 
 
 def stats_to_response(stats, m):
     if stats is None:
         return
-    
-    ystats = Stats()
-    ystats.func_stats = []
-    m.ystats = ystats
+
+    pstats = Stats()
+    pstats.func_stats = []
+    pstats.profiler = "cprofile"
+    m.ystats = pstats
 
     for func, stat in stats.items():
         path, line, func_name = func
@@ -104,7 +106,7 @@ def stats_to_response(stats, m):
         func = Function()
         func_stat = FuncStat()
         func.func_stat = func_stat
-        ystats.func_stats.append(func)
+        pstats.func_stats.append(func)
 
         func_stat.file = path
         func_stat.line = line
@@ -125,3 +127,58 @@ def stats_to_response(stats, m):
             caller_stat.calls_count = cc
             caller_stat.total_time = ct
             caller_stat.own_time = tt
+
+
+def create_func_stat(file_path, line_num, func_name, calls_count, total_time, own_time,
+                     thread_name, thread_id=None):
+    func_stat = FuncStat()
+    func_stat.file = str(file_path)
+    func_stat.line = int(line_num) if line_num is not None else None
+    func_stat.func_name = str(func_name)
+    func_stat.calls_count = int(calls_count)
+    func_stat.total_time = float(total_time)
+    func_stat.own_time = float(own_time)
+    func_stat.threadName = str(thread_name)
+    if thread_id is not None:
+        func_stat.threadId = str(thread_id)
+    return func_stat
+
+
+def ystats_to_response(ystats_data, m):
+    if ystats_data is None:
+        return
+
+    ystats = Stats()
+    ystats.func_stats = []
+    ystats.profiler = "yappi"
+    m.ystats = ystats
+
+    for stat in ystats_data:
+        func_stat = create_func_stat(
+            file_path=stat[1],
+            line_num=stat[2],
+            func_name=stat[0],
+            calls_count=stat[4],
+            total_time=stat[6],
+            own_time=stat[7],
+            thread_name=stat[11],
+            thread_id=stat[10]
+        )
+
+        func = Function()
+        func.func_stat = func_stat
+        ystats.func_stats.append(func)
+
+        func.callers = []
+        if stat[9]:
+            for child_func in stat[9]:
+                caller_stat = create_func_stat(
+                    file_path=child_func[8],
+                    line_num=child_func[9],
+                    func_name=child_func[10],
+                    calls_count=child_func[0],
+                    total_time=child_func[3],
+                    own_time=child_func[4],
+                    thread_name=stat[11]
+                )
+                func.callers.append(caller_stat)

python/helpers/profiler/profiler.thrift

@@ -7,6 +7,8 @@ struct FuncStat {
   4: required i32 calls_count,	  // number of times the executed function is called.
   5: required double total_time,	  // total time spent in the executed function. See Clock Types to interpret this value correctly.
   6: required double own_time,     // total time spent in the executed function, excluding subcalls. See Clock Types to interpret this value correctly.
+  7: optional string threadName, // Name of thread
+  8: optional string threadId, // ID of thread
 }
 
 struct Function {
@@ -16,7 +18,8 @@ struct Function {
 
 
 struct Stats {
-  1: required list<Function> func_stats
+  1: required list<Function> func_stats,
+  2: required string profiler // profiler type (currently yappi/cProfile)
 }
 
 struct CallTreeStat {

python/helpers/profiler/yappi_profiler.py

@@ -6,6 +6,7 @@ class YappiProfile(object):
     """
     def __init__(self):
         self.stats = None
+        yappi.set_clock_type("wall")
 
     def runcall(self, func, *args, **kw):
         self.enable()
@@ -20,17 +21,28 @@ class YappiProfile(object):
     def disable(self):
         yappi.stop()
 
+    def convert_stats_to_dict(self, stats):
+        result = []
+        for stat in stats:
+            func_dict = {
+                'filename': stat.module,
+                'name': stat.name,
+                'line': stat.lineno,
+                'calls': stat.ncall,
+                'total_time': stat.ttot,
+                'cumulative_time': stat.tsub,
+                'thread_id': stat.ctx_id
+            }
+            result.append(func_dict)
+        return result
+
     def create_stats(self):
-        self.stats = yappi.convert2pstats(yappi.get_func_stats()).stats
+        self.stats = yappi.get_func_stats()
 
     def getstats(self):
         self.create_stats()
-
         return self.stats
 
     def dump_stats(self, file):
-        import marshal
-        f = open(file, 'wb')
-        marshal.dump(self.getstats(), f)
-        f.close()
+        self.getstats().save(file, type="ystat")