def test_get_tracer(self):
with self.assertRaises(NameError):
tmp = __viz_tracer__ # noqa: F821
self.assertIs(get_tracer(), None)
tracer = VizTracer
builtins.__dict__["__viz_tracer__"] = tracer
self.assertIs(get_tracer(), tracer)
builtins.__dict__.pop("__viz_tracer__")
def log_sparse(func):
tracer = get_tracer()
if not tracer or not tracer.log_sparse:
return func
@functools.wraps(func)
def wrapper(*args, **kwargs):
tracer = get_tracer()
if tracer:
# This should always be true
start = tracer._tracer.getts()
ret = func(*args, **kwargs)
dur = tracer._tracer.getts() - start
raw_data = {
"ph": "X",
"name": func.__qualname__,
"ts": start,
"dur": dur,
"cat": "FEE"
}
tracer._tracer.addraw(raw_data)
else: # pragma: no cover
raise Exception("This should not be possible")
return ret
return wrapper
def log_sparse(func: Callable) -> Callable:
tracer = get_tracer()
if not tracer or not tracer.log_sparse:
return func
@functools.wraps(func)
def wrapper(*args, **kwargs) -> Any:
tracer = get_tracer()
if tracer:
# This should always be true
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)
else: # pragma: no cover
raise Exception("This should not be possible")
return ret
return wrapper
def exit_routine():
atexit.unregister(exit_routine)
if not self._exiting:
self._exiting = True
tracer = viztracer.get_tracer()
tracer.stop()
tracer.save()
tracer.terminate()
exit(0)
def wrapper(*args, **kwargs):
tracer = viztracer.get_tracer()
if tracer and tracer.log_sparse:
tracer.resume()
try:
return func(*args, **kwargs)
finally:
tracer.pause()
return func(*args, **kwargs)
def ignore_wrapper(*args, **kwargs):
# We need this to keep trace a local variable
t = tracer
if not t:
t = get_tracer()
if not t:
raise NameError("ignore_function only works with global tracer")
t.pause()
ret = func(*args, **kwargs)
t.resume()
return ret
def _make_viztracer_initializer_and_initargs():
try:
import viztracer
tracer = viztracer.get_tracer()
if tracer is not None and getattr(tracer, 'enable', False):
# Profiler is active: introspect its configuration to
# initialize the workers with the same configuration.
return _viztracer_init, (tracer.init_kwargs, )
except ImportError:
# viztracer is not installed: nothing to do
pass
except Exception as e:
# In case viztracer's API evolve, we do not want to crash loky but
# we want to know about it to be able to update loky.
warnings.warn("Unable to introspect viztracer state: {}".format(e))
return None, ()
def ignore_function(method=None, tracer=None):
if not tracer:
tracer = get_tracer()
def inner(func):
@functools.wraps(func)
def ignore_wrapper(*args, **kwargs):
tracer.pause()
ret = func(*args, **kwargs)
tracer.resume()
return ret
return ignore_wrapper
if method:
return inner(method)
return inner
def wrapper(*args, **kwargs):
tracer = get_tracer()
if tracer:
start = tracer._tracer.getts()
ret = func(*args, **kwargs)
dur = tracer._tracer.getts() - start
raw_data = {
"ph": "X",
"name": func.__qualname__,
"ts": start,
"dur": dur,
"cat": "FEE"
}
tracer._tracer.addraw(raw_data)
else:
ret = func(*args, **kwargs)
return ret
def wrapper(*args, **kwargs):
tracer = get_tracer()
if tracer:
# This should always be true
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)
else: # pragma: no cover
raise Exception("This should not be possible")
return ret
def trace_sparse(func):
"Only collect traces while `func` is running"
tracer = viztracer.get_tracer()
if not tracer or not tracer.log_sparse:
return func
if inspect.iscoroutinefunction(func):
# TODO this async mode doesn't work that well, since it also profiles whatever other
# coroutines get context-switched in while awaiting this one.
@functools.wraps(func)
# Function declaration "wrapper" is obscured by a declaration of the same name
async def wrapper(*args, **kwargs): # type: ignore
tracer = viztracer.get_tracer()
if tracer and tracer.log_sparse:
tracer.resume()
try:
return await func(*args, **kwargs)
finally:
tracer.pause()
return func(*args, **kwargs)
else:
@functools.wraps(func)
def wrapper(*args, **kwargs):
tracer = viztracer.get_tracer()
if tracer and tracer.log_sparse:
tracer.resume()
try:
return func(*args, **kwargs)
finally:
tracer.pause()
return func(*args, **kwargs)
return wrapper
import logging
from viztracer import get_tracer, VizLoggingHandler
def fib(n):
if n < 2:
logging.warn("Base case, return 1")
return 1
logging.info("Recursive, working on {}".format(n))
return fib(n - 1) + fib(n - 2)
handler = VizLoggingHandler()
handler.setTracer(get_tracer())
logging.basicConfig(handlers=[handler], level=logging.INFO)
fib(7)
