1import warnings
2
3
4def _viztracer_init(init_kwargs):
5    """Initialize viztracer's profiler in worker processes"""
6    from viztracer import VizTracer
7    tracer = VizTracer(**init_kwargs)
8    tracer.register_exit()
9    tracer.start()
10
11
12def _make_viztracer_initializer_and_initargs():
13    try:
14        import viztracer
15        tracer = viztracer.get_tracer()
16        if tracer is not None and getattr(tracer, 'enable', False):
17            # Profiler is active: introspect its configuration to
18            # initialize the workers with the same configuration.
19            return _viztracer_init, (tracer.init_kwargs,)
20    except ImportError:
21        # viztracer is not installed: nothing to do
22        pass
23    except Exception as e:
24        # In case viztracer's API evolve, we do not want to crash loky but
25        # we want to know about it to be able to update loky.
26        warnings.warn("Unable to introspect viztracer state: {}"
27                      .format(e))
28    return None, ()
29
30
31class _ChainedInitializer():
32    """Compound worker initializer
33
34    This is meant to be used in conjunction with _chain_initializers to
35    produce  the necessary chained_args list to be passed to __call__.
36    """
37
38    def __init__(self, initializers):
39        self._initializers = initializers
40
41    def __call__(self, *chained_args):
42        for initializer, args in zip(self._initializers, chained_args):
43            initializer(*args)
44
45
46def _chain_initializers(initializer_and_args):
47    """Convenience helper to combine a sequence of initializers.
48
49    If some initializers are None, they are filtered out.
50    """
51    filtered_initializers = []
52    filtered_initargs = []
53    for initializer, initargs in initializer_and_args:
54        if initializer is not None:
55            filtered_initializers.append(initializer)
56            filtered_initargs.append(initargs)
57
58    if len(filtered_initializers) == 0:
59        return None, ()
60    elif len(filtered_initializers) == 1:
61        return filtered_initializers[0], filtered_initargs[0]
62    else:
63        return _ChainedInitializer(filtered_initializers), filtered_initargs
64
65
66def _prepare_initializer(initializer, initargs):
67    if initializer is not None and not callable(initializer):
68        raise TypeError(
69            "initializer must be a callable, got: {!r}"
70            .format(initializer)
71        )
72
73    # Introspect runtime to determine if we need to propagate the viztracer
74    # profiler information to the workers:
75    return _chain_initializers([
76        (initializer, initargs),
77        _make_viztracer_initializer_and_initargs(),
78    ])
79