diff --git a/docs/source/basic_usage.rst b/docs/source/basic_usage.rst index ca719982..d199afb3 100644 --- a/docs/source/basic_usage.rst +++ b/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 -------------- diff --git a/docs/source/decorator.rst b/docs/source/decorator.rst index a1b9c9f2..34b06088 100644 --- a/docs/source/decorator.rst +++ b/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 \ No newline at end of file + # 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() diff --git a/docs/source/filter.rst b/docs/source/filter.rst index d8a78435..18632792 100644 --- a/docs/source/filter.rst +++ b/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`` \ No newline at end of file +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. diff --git a/src/viztracer/cellmagic.py b/src/viztracer/cellmagic.py index d89853b1..6c18b7e1 100644 --- a/src/viztracer/cellmagic.py +++ b/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 diff --git a/src/viztracer/decorator.py b/src/viztracer/decorator.py index 8067e5a3..6560b8fc 100644 --- a/src/viztracer/decorator.py +++ b/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) diff --git a/src/viztracer/modules/snaptrace.c b/src/viztracer/modules/snaptrace.c index 5afef02d..32d3d809 100644 --- a/src/viztracer/modules/snaptrace.c +++ b/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) { diff --git a/src/viztracer/tracer.py b/src/viztracer/tracer.py index f72569b7..56ea3a5e 100644 --- a/src/viztracer/tracer.py +++ b/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"]: diff --git a/src/viztracer/viztracer.py b/src/viztracer/viztracer.py index 3dccb267..0c41dbdd 100644 --- a/src/viztracer/viztracer.py +++ b/src/viztracer/viztracer.py @@ -121,11 +121,13 @@ def init_kwargs(self) -> Dict: } def __enter__(self) -> "VizTracer": - self.start() + if not self.log_sparse: + self.start() return self def __exit__(self, type, value, trace) -> None: - self.stop() + if not self.log_sparse: + self.stop() self.save() self.terminate() @@ -162,6 +164,12 @@ def log_event(self, event_name: str) -> VizEvent: call_frame = sys._getframe(1) return VizEvent(self, event_name, call_frame.f_code.co_filename, call_frame.f_lineno) + def shield_ignore(self, func, *args, **kwargs): + prev_ignore_stack = self.setignorestackcounter(0) + res = func(*args, **kwargs) + self.setignorestackcounter(prev_ignore_stack) + return res + def set_afterfork(self, callback: Callable, *args, **kwargs) -> None: self._afterfork_cb = callback self._afterfork_args = args diff --git a/tests/test_logsparse.py b/tests/test_logsparse.py index 534a5329..d85705f2 100644 --- a/tests/test_logsparse.py +++ b/tests/test_logsparse.py @@ -1,6 +1,7 @@ # Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0 # For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt +import functools import multiprocessing import os @@ -75,32 +76,110 @@ def f(x): time.sleep(0.1) """ +file_context_manager = """ +from viztracer import VizTracer, log_sparse + +@log_sparse(dynamic_tracer_check=True) +def f(): + return 1 + +def g(): + return f() + +@log_sparse +def h(): + return 2 + +@log_sparse(dynamic_tracer_check=True, stack_depth=1) +def q(): + return 3 + +if __name__ == "__main__": + with VizTracer(output_file="result.json"): + assert g() == 1 + assert h() == 2 + assert q() == 3 +""" + +file_context_manager_logsparse = """ +from viztracer import VizTracer, log_sparse + +@log_sparse(dynamic_tracer_check=True) +def f(): + return 1 + +def g(): + return f() + +@log_sparse +def h(): + return 2 + +@log_sparse(dynamic_tracer_check=True, stack_depth=1) +def q(): + return 3 + +if __name__ == "__main__": + with VizTracer(output_file="result.json", log_sparse=True): + assert g() == 1 + assert h() == 2 + assert q() == 3 +""" + +file_context_manager_logsparse_stack = """ +from viztracer import VizTracer, log_sparse + +@log_sparse(dynamic_tracer_check=True) +def f(): + return 1 + +@log_sparse(dynamic_tracer_check=True, stack_depth=1) +def g(): + return f() + +@log_sparse(dynamic_tracer_check=True) +def h(): + return 2 + +if __name__ == "__main__": + assert g() == 1 + assert h() == 2 + + with VizTracer(output_file="result.json", log_sparse=True): + assert g() == 1 + assert h() == 2 +""" + class TestLogSparse(CmdlineTmpl): - def test_basic(self): - def check_func(data): - for entry in data["traceEvents"]: - self.assertNotEqual(entry["name"], "f") + def check_func(self, data, target): + names = [entry["name"] for entry in data["traceEvents"]] + function_names = [name.split(' ')[0] for name in names if name not in ['process_name', 'thread_name']] + self.assertEqual(function_names, target) + + def test_basic(self): self.template(["viztracer", "-o", "result.json", "--log_sparse", "cmdline_test.py"], script=file_basic, expected_output_file="result.json", - expected_entries=1) + expected_entries=1, + check_func=functools.partial(self.check_func, target=['f'])) self.template(["viztracer", "-o", "result.json", "cmdline_test.py"], script=file_basic, - expected_output_file="result.json", - check_func=check_func) + expected_output_file="result.json") def test_stack(self): self.template(["viztracer", "-o", "result.json", "--log_sparse", "cmdline_test.py"], script=file_stack, expected_output_file="result.json", - expected_entries=4) + expected_entries=4, + check_func=functools.partial(self.check_func, target=['f', 'g', 'f', 'g'])) self.template(["viztracer", "-o", "result.json", "--log_sparse", "cmdline_test.py"], script=file_stack_nested, expected_output_file="result.json", - expected_entries=4) + expected_entries=4, + check_func=functools.partial(self.check_func, target=['f', 'g', 'f', 'g'])) def test_without_tracer(self): self.template(["python", "cmdline_test.py"], script=file_basic, expected_output_file=None) @@ -113,8 +192,22 @@ def test_multiprocess(self): script=file_multiprocess, expected_output_file="result.json", expected_entries=3, + check_func=functools.partial(self.check_func, target=['f', 'f', 'f']), concurrency="multiprocessing") except Exception as e: # coveragepy has some issue with multiprocess pool if not os.getenv("COVERAGE_RUN"): raise e + + def test_context_manager(self): + self.template(["python", "cmdline_test.py"], script=file_context_manager, + expected_output_file="result.json", expected_entries=4, + check_func=functools.partial(self.check_func, target=['f', 'g', 'h', 'q'])) + + self.template(["python", "cmdline_test.py"], script=file_context_manager_logsparse, + expected_output_file="result.json", expected_entries=2, + check_func=functools.partial(self.check_func, target=['f', 'q'])) + + self.template(["python", "cmdline_test.py"], script=file_context_manager_logsparse_stack, + expected_output_file="result.json", expected_entries=2, + check_func=functools.partial(self.check_func, target=['g', 'h']))