From 407337fbf3109b297e21ee428d017274a83c254d Mon Sep 17 00:00:00 2001
From: "zhangzhicheng.zzc" <zhangzhicheng.zzc@alibaba-inc.com>
Date: Thu, 7 Jul 2022 20:21:52 +0800
Subject: [PATCH] [to #42322933]  add model profiling

---
 tests/run.py            |  6 +++++
 tests/utils/__init__.py |  1 +
 tests/utils/profiler.py | 59 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+)
 create mode 100644 tests/utils/profiler.py

diff --git a/tests/run.py b/tests/run.py
index 38c5a897..eaa45f60 100644
--- a/tests/run.py
+++ b/tests/run.py
@@ -62,7 +62,13 @@ if __name__ == '__main__':
         '--test_dir', default='tests', help='directory to be tested')
     parser.add_argument(
         '--level', default=0, type=int, help='2 -- all, 1 -- p1, 0 -- p0')
+    parser.add_argument(
+        '--disable_profile', action='store_true', help='disable profiling')
     args = parser.parse_args()
     set_test_level(args.level)
     logger.info(f'TEST LEVEL: {test_level()}')
+    if not args.disable_profile:
+        from utils import profiler
+        logger.info('enable profile ...')
+        profiler.enable()
     main(args)
diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py
index e69de29b..9166292f 100644
--- a/tests/utils/__init__.py
+++ b/tests/utils/__init__.py
@@ -0,0 +1 @@
+from .profiler import *  # noqa F403
diff --git a/tests/utils/profiler.py b/tests/utils/profiler.py
new file mode 100644
index 00000000..92708ad3
--- /dev/null
+++ b/tests/utils/profiler.py
@@ -0,0 +1,59 @@
+import importlib
+import sys
+from functools import wraps
+from typing import Any, Callable, Dict, Tuple, Type
+
+
+def reraise(tp, value, tb):
+    try:
+        if value is None:
+            value = tp()
+        if value.__traceback__ is not tb:
+            raise value.with_traceback(tb)
+        raise value
+    finally:
+        value = None
+        tb = None
+
+
+class Profiler:
+
+    def __init__(self) -> None:
+        import cProfile
+        self.pr = cProfile.Profile()
+
+    def __enter__(self):
+        self.pr.enable()
+
+    def __exit__(self, tp, exc, tb):
+        self.pr.disable()
+        if tp is not None:
+            reraise(tp, exc, tb)
+
+        import pstats
+        ps = pstats.Stats(self.pr, stream=sys.stderr).sort_stats('tottime')
+        ps.print_stats(20)
+
+
+def wrapper(tp: Type[Profiler]) -> Callable[[], Callable[..., Any]]:
+
+    def _inner(func: Callable[..., Any]) -> Callable[..., Any]:
+
+        @wraps(func)
+        def executor(*args: Tuple[Any, ...], **kwargs: Dict[str, Any]) -> Any:
+            with tp():
+                return func(*args, **kwargs)
+
+        return executor
+
+    return _inner
+
+
+PIPELINE_BASE_MODULE = 'modelscope.pipelines.base'
+PIPELINE_BASE_CLASS = 'Pipeline'
+
+
+def enable():
+    base = importlib.import_module(PIPELINE_BASE_MODULE)
+    Pipeline = getattr(base, PIPELINE_BASE_CLASS)
+    Pipeline.__call__ = wrapper(Profiler)(Pipeline.__call__)