Updates for @log_sparse with context manager and parameters for %%viz…

…tracer (#360)

docs/source/basic_usage.rst

@@ -110,7 +110,24 @@ If you are using Jupyter, you can use viztracer cell magics.
     %%viztracer
     # Your code after
 
-A ``Show VizTracer Report`` button will appear after the cell and you can click it to view the results
+.. code-block:: python
+
+    # you can define arguments of VizTracer in magic
+    %%viztracer -p 8888
+    # Your code after
+
+A ``Show VizTracer Report`` button will appear after the cell and you can click it to view the results.
+
+Cell magic ``%%viztracer`` supports some of the command line arguments:
+
+* ``--port``
+* ``--output_file``
+* ``--max_stack_depth``
+* ``--ignore_c_function``
+* ``--ignore_frozen``
+* ``--log_func_args``
+* ``--log_print``
+* ``--log_sparse``
 
 Display Report
 --------------

docs/source/decorator.rst

@@ -3,7 +3,7 @@ Decorator
 
 .. py:decorator:: ignore_function
 
-    ``@ignore_function`` can tell VizTracer to skip on functions you specified. 
+    ``@ignore_function`` can tell VizTracer to skip on functions you specified.
 
     .. code-block:: python
 
@@ -12,21 +12,21 @@ Decorator
         def function_you_want_to_ignore():
             # function body
 
-        # You can specify tracer if no global tracer is registered 
+        # You can specify tracer if no global tracer is registered
         @ignore_function(tracer=tracer)
         def function_you_want_to_ignore():
             # function body
 
 .. py:decorator:: trace_and_save(method=None, output_dir="./", **viztracer_kwargs)
-    
+
     :param function method: trick to make both ``@trace_and_save`` and ``@trace_and_save(**kwargs)`` work
     :param str output_dir: output directory you want to put your logs in
     :param dict viztracer_kwargs: kwargs for VizTracer
 
-    ``@trace_and_save`` can be used to trace a certain function and save the result as the program goes. 
-    This can be very helpful when you are running a very long program and just want to keep recording 
+    ``@trace_and_save`` can be used to trace a certain function and save the result as the program goes.
+    This can be very helpful when you are running a very long program and just want to keep recording
     something periodically. You won't drain your memory and the parsing/dumping will be done in a new process,
-    which can minimize the performance impact to your main process. 
+    which can minimize the performance impact to your main process.
 
     You can pass any argument you want to ``VizTracer`` by giving it to the decorator
 
@@ -39,4 +39,37 @@ Decorator
         # this works as well
         @trace_and_save
         def function_you_want_to_trace():
-            # function body
+            # function body
+
+.. py:decorator:: log_sparse(func=None, stack_depth=0, dynamic_tracer_check=False)
+
+    You can make VizTracer log only certain functions using ``--log_sparse`` mode.
+
+    :param function func: callable to decorate
+    :param int stack_depth: log the function and its descendants with a limit stack depth
+    :param bool dynamic_tracer_check: run time check of tracer
+
+    .. code-block:: python
+
+        from viztracer import log_sparse
+
+        # @log_sparse will only log this function
+        @log_sparse
+        def function_you_want_to_log():
+            # function body
+
+        # @log_sparse(stack_depth=5) will log this function and its descendants
+        # with a limit stack depth of 5
+        # Nested @log_sparse with stack_depth won't work
+        # (only the outermost function and its stack will be logged)
+        @log_sparse(stack_depth=5)
+        def function_you_want_to_log():
+            # function body
+
+        # Use dynamic_tracer_check=True if you use tracer as a context manager (or with %%viztracer).
+        @log_sparse(dynamic_tracer_check=True)
+        def function_you_want_to_log():
+            # function body
+
+        with VizTracer(log_sparse=True):
+            function_you_want_to_log()

docs/source/filter.rst

@@ -79,14 +79,14 @@ Ignore C Function
 
 By default, VizTracer will trace both python and C functions. You can turn off tracing C functions by
 
-.. code-block:: 
+.. code-block::
 
     viztracer --ignore_c_function my_script.py
 
 OR
 
 .. code-block:: python
-    
+
     tracer = VizTracer(ignore_c_function=True)
 
 Since most of the builtin functions(like ``append`` or ``len``) are C functions which are frequently called,
@@ -98,14 +98,14 @@ Ignore Non File
 
 You can ask VizTracer not to trace any functions that are not in a valid file(mostly import stuff) using ``ignore_frozen``
 
-.. code-block:: 
+.. code-block::
 
     viztracer --ignore_frozen my_script.py
 
 OR
 
 .. code-block:: python
-    
+
     tracer = VizTracer(ignore_frozen=True)
 
 
@@ -128,7 +128,7 @@ Log Sparse
 ----------
 
 You can make VizTracer log only certain functions using ``--log_sparse``. This is helpful when you are only interested in the time spent on
-specific functions for a big picture on larger projects. 
+specific functions for a big picture on larger projects.
 
 First, you need to add decorator ``@log_sparse`` on the function you want to log
 
@@ -149,15 +149,27 @@ First, you need to add decorator ``@log_sparse`` on the function you want to log
     def function_you_want_to_log():
         # function body
 
+    # Use dynamic_tracer_check=True if you use tracer as a context manager (or with %%viztracer).
+    @log_sparse(dynamic_tracer_check=True)
+    def function_you_want_to_log():
+        # function body
+
+    with VizTracer(log_sparse=True):
+        function_you_want_to_log()
+
 Then just call viztracer with ``--log_sparse``
 
-.. code-block:: 
+.. code-block::
 
     viztracer --log_sparse your_script.py
-    
+
 When you are using ``--log_sparse``, due to the nature of the recording, some advanced features may not work with it.
 
 You can leave ``@log_sparse`` as it is when you are not running the script with VizTracer. It will be like a no-op
 
 If you want to log a piece of code, rather than a full function, please check :ref:`duration_event_label`. Duration Event
-is compatible with ``log_sparse``
+is compatible with ``log_sparse``
+
+To use ``@log_sparse`` in conjunction with a context manager, you must define decorating functions within the created
+context, or set the `dynamic_tracer_check=True`` argument of decorator. The second option leads to runtime checks,
+so it increases the overhead.

src/viztracer/cellmagic.py

@@ -4,9 +4,27 @@
 try:
     from IPython.core.magic import (Magics, cell_magic,  # type: ignore
                                     magics_class, needs_local_scope)
+    from IPython.core.magic_arguments import (argument, magic_arguments, parse_argstring)  # type: ignore
 
     @magics_class
     class VizTracerMagics(Magics):
+        @magic_arguments()
+        @argument("--port", "-p", default=9001, type=int,
+                  help="specify the port vizviewer will use")
+        @argument("--output_file", default="./viztracer_report.json",
+                  help="output file path. End with .json or .html or .gz")
+        @argument("--max_stack_depth", type=int, default=-1,
+                  help="maximum stack depth you want to trace.")
+        @argument("--ignore_c_function", action="store_true", default=False,
+                  help="ignore all c functions including most builtin functions and libraries")
+        @argument("--ignore_frozen", action="store_true", default=False,
+                  help="ignore all functions that are frozen(like import)")
+        @argument("--log_func_args", action="store_true", default=False,
+                  help="log all function arguments, this will introduce large overhead")
+        @argument("--log_print", action="store_true", default=False,
+                  help="replace all print() function to adding an event to the result")
+        @argument("--log_sparse", action="store_true", default=False,
+                  help="log only selected functions with @log_sparse")
         @needs_local_scope
         @cell_magic
         def viztracer(self, line, cell, local_ns) -> None:
@@ -15,14 +33,25 @@ def viztracer(self, line, cell, local_ns) -> None:
 
             from .viewer import ServerThread
             from .viztracer import VizTracer
+            options = parse_argstring(self.viztracer, line)
             assert self.shell is not None
             code = self.shell.transform_cell(cell)
-            file_path = "./viztracer_report.json"
-            with VizTracer(verbose=0, output_file=file_path):
+            file_path = options.output_file
+            tracer_kwargs = {
+                "output_file": file_path,
+                "verbose": 0,
+                "max_stack_depth": options.max_stack_depth,
+                "ignore_c_function": options.ignore_c_function,
+                "ignore_frozen": options.ignore_frozen,
+                "log_func_args": options.log_func_args,
+                "log_print": options.log_print,
+                "log_sparse": options.log_sparse
+            }
+            with VizTracer(**tracer_kwargs):
                 exec(code, local_ns, local_ns)
 
             def view():  # pragma: no cover
-                server = ServerThread(file_path, once=True)
+                server = ServerThread(file_path, port=options.port, once=True)
                 server.start()
                 server.ready.wait()
                 import webbrowser

src/viztracer/decorator.py

@@ -57,47 +57,56 @@ def wrapper(*args, **kwargs) -> Any:
     return inner
 
 
-def log_sparse(func: Optional[Callable] = None, stack_depth: int = 0) -> Callable:
-    tracer = get_tracer()
-    if tracer is None or not tracer.log_sparse:
-        if func is None:
-            return lambda f: f
-        return func
-
+def _log_sparse_wrapper(func: Callable, stack_depth: int = 0,
+                        dynamic_tracer_check: bool = False) -> Callable:
+    if not dynamic_tracer_check:
+        tracer = get_tracer()
+        if tracer is None or not tracer.log_sparse:
+            return func
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs) -> Any:
+        local_tracer = get_tracer() if dynamic_tracer_check else tracer
+
+        if local_tracer is None:
+            return func(*args, **kwargs)
+        assert isinstance(local_tracer, VizTracer)
+
+        if local_tracer.log_sparse and not local_tracer.enable:
+            if stack_depth > 0:
+                orig_max_stack_depth = local_tracer.max_stack_depth
+                local_tracer.max_stack_depth = stack_depth
+                local_tracer.start()
+                ret = func(*args, **kwargs)
+                local_tracer.stop()
+                local_tracer.max_stack_depth = orig_max_stack_depth
+                return ret
+            else:
+                start = local_tracer.getts()
+                ret = func(*args, **kwargs)
+                dur = local_tracer.getts() - start
+                code = func.__code__
+                raw_data = {
+                    "ph": "X",
+                    "name": f"{code.co_name} ({code.co_filename}:{code.co_firstlineno})",
+                    "ts": start,
+                    "dur": dur,
+                    "cat": "FEE",
+                }
+                local_tracer.add_raw(raw_data)
+                return ret
+        elif local_tracer.enable and not local_tracer.log_sparse:
+            # The call is made from the module inside, so if `trace_self=False` it will be ignored.
+            # To avoid this behavior, we need to reset the counter `ignore_stack_depth`` and then
+            # recover it
+            return local_tracer.shield_ignore(func, *args, **kwargs)
+        else:
+            return func(*args, **kwargs)
+
+    return wrapper
+
+
+def log_sparse(func: Optional[Callable] = None, stack_depth: int = 0, dynamic_tracer_check: bool = False) -> Callable:
     if func is None:
-        def inner(dec_func: Callable) -> Callable:
-            @functools.wraps(dec_func)
-            def wrapper(*args, **kwargs) -> Any:
-                assert isinstance(tracer, VizTracer)
-                if not tracer.enable:
-                    orig_max_stack_depth = tracer.max_stack_depth
-                    tracer.max_stack_depth = stack_depth
-                    tracer.start()
-                    ret = dec_func(*args, **kwargs)
-                    tracer.stop()
-                    tracer.max_stack_depth = orig_max_stack_depth
-                    return ret
-                else:
-                    return dec_func(*args, **kwargs)
-            return wrapper
-        return inner
-    else:
-        @functools.wraps(func)
-        def wrapper(*args, **kwargs) -> Any:
-            assert callable(func)
-            assert isinstance(tracer, VizTracer)
-            start = tracer.getts()
-            ret = func(*args, **kwargs)
-            dur = tracer.getts() - start
-            code = func.__code__
-            raw_data = {
-                "ph": "X",
-                "name": f"{code.co_name} ({code.co_filename}:{code.co_firstlineno})",
-                "ts": start,
-                "dur": dur,
-                "cat": "FEE",
-            }
-            tracer.add_raw(raw_data)
-            return ret
-
-        return wrapper
+        return functools.partial(_log_sparse_wrapper, stack_depth=stack_depth, dynamic_tracer_check=dynamic_tracer_check)
+    return _log_sparse_wrapper(func=func, stack_depth=stack_depth, dynamic_tracer_check=dynamic_tracer_check)

src/viztracer/modules/snaptrace.c

@@ -43,6 +43,7 @@ static PyObject* snaptrace_addfunctionarg(TracerObject* self, PyObject* args);
 static PyObject* snaptrace_getfunctionarg(TracerObject* self, PyObject* args);
 static PyObject* snaptrace_getts(TracerObject* self, PyObject* args);
 static PyObject* snaptrace_setcurrstack(TracerObject* self, PyObject* args);
+static PyObject* snaptrace_setignorestackcounter(TracerObject* self, PyObject* args);
 static void snaptrace_threaddestructor(void* key);
 static struct ThreadInfo* snaptrace_createthreadinfo(TracerObject* self);
 static void log_func_args(struct FunctionNode* node, PyFrameObject* frame);
@@ -218,6 +219,7 @@ static PyMethodDef Tracer_methods[] = {
     {"setcurrstack", (PyCFunction)snaptrace_setcurrstack, METH_VARARGS, "set current stack depth"},
     {"pause", (PyCFunction)snaptrace_pause, METH_VARARGS, "pause profiling"},
     {"resume", (PyCFunction)snaptrace_resume, METH_VARARGS, "resume profiling"},
+    {"setignorestackcounter", (PyCFunction)snaptrace_setignorestackcounter, METH_VARARGS, "reset ignore stack depth"},
     {NULL, NULL, 0, NULL}
 };
 
@@ -301,7 +303,7 @@ snaptrace_tracefunc(PyObject* obj, PyFrameObject* frame, int what, PyObject* arg
                 goto cleanup;
             }
         }
-        
+
         // IMPORTANT: the C function will always be called from our python methods, 
         // so this check is redundant. However, the user should never be allowed 
         // to call our C methods directly! Otherwise C functions will be recorded.
@@ -1386,6 +1388,24 @@ snaptrace_addraw(TracerObject* self, PyObject* args)
     Py_RETURN_NONE;
 }
 
+static PyObject*
+snaptrace_setignorestackcounter(TracerObject* self, PyObject* args)
+{
+    int value = 0;
+    int current_value = 0;
+    struct ThreadInfo* info = get_thread_info(self);
+
+    if (!PyArg_ParseTuple(args, "i", &value)) {
+        printf("Error when parsing arguments!\n");
+        exit(1);
+    }
+
+    current_value = info->ignore_stack_depth;
+    info->ignore_stack_depth = value;
+
+    return Py_BuildValue("i", current_value);
+}
+
 static PyObject*
 snaptrace_getfunctionarg(TracerObject* self, PyObject* args)
 {

src/viztracer/tracer.py

@@ -277,6 +277,10 @@ def enable_thread_tracing(self) -> None:
     def getts(self) -> float:
         return self._tracer.getts()
 
+    def setignorestackcounter(self, value) -> int:
+        # +2 and -2 to compensate for two calls: the current one and the call in the next line
+        return self._tracer.setignorestackcounter(value + 2) - 2
+
     def add_instant(self, name: str, args: Any = None, scope: str = "g") -> None:
         if self.enable:
             if scope not in ["g", "p", "t"]: