[to #42322933] add model profiling

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)

tests/utils/__init__.py

@@ -0,0 +1 @@
+from .profiler import *  # noqa F403

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__)