81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
|
import warnings
|
||
|
|
|
||
|
|
|
||
|
|
def _viztracer_init(init_kwargs):
|
||
|
|
"""Initialize viztracer's profiler in worker processes"""
|
||
|
|
from viztracer import VizTracer
|
||
|
|
|
||
|
|
tracer = VizTracer(**init_kwargs)
|
||
|
|
tracer.register_exit()
|
||
|
|
tracer.start()
|
||
|
|
|
||
|
|
|
||
|
|
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(f"Unable to introspect viztracer state: {e}")
|
||
|
|
return None, ()
|
||
|
|
|
||
|
|
|
||
|
|
class _ChainedInitializer:
|
||
|
|
"""Compound worker initializer
|
||
|
|
|
||
|
|
This is meant to be used in conjunction with _chain_initializers to
|
||
|
|
produce the necessary chained_args list to be passed to __call__.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, initializers):
|
||
|
|
self._initializers = initializers
|
||
|
|
|
||
|
|
def __call__(self, *chained_args):
|
||
|
|
for initializer, args in zip(self._initializers, chained_args):
|
||
|
|
initializer(*args)
|
||
|
|
|
||
|
|
|
||
|
|
def _chain_initializers(initializer_and_args):
|
||
|
|
"""Convenience helper to combine a sequence of initializers.
|
||
|
|
|
||
|
|
If some initializers are None, they are filtered out.
|
||
|
|
"""
|
||
|
|
filtered_initializers = []
|
||
|
|
filtered_initargs = []
|
||
|
|
for initializer, initargs in initializer_and_args:
|
||
|
|
if initializer is not None:
|
||
|
|
filtered_initializers.append(initializer)
|
||
|
|
filtered_initargs.append(initargs)
|
||
|
|
|
||
|
|
if not filtered_initializers:
|
||
|
|
return None, ()
|
||
|
|
elif len(filtered_initializers) == 1:
|
||
|
|
return filtered_initializers[0], filtered_initargs[0]
|
||
|
|
else:
|
||
|
|
return _ChainedInitializer(filtered_initializers), filtered_initargs
|
||
|
|
|
||
|
|
|
||
|
|
def _prepare_initializer(initializer, initargs):
|
||
|
|
if initializer is not None and not callable(initializer):
|
||
|
|
raise TypeError(
|
||
|
|
f"initializer must be a callable, got: {initializer!r}"
|
||
|
|
)
|
||
|
|
|
||
|
|
# Introspect runtime to determine if we need to propagate the viztracer
|
||
|
|
# profiler information to the workers:
|
||
|
|
return _chain_initializers(
|
||
|
|
[
|
||
|
|
(initializer, initargs),
|
||
|
|
_make_viztracer_initializer_and_initargs(),
|
||
|
|
]
|
||
|
|
)
|