Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/joblib/externals/loky/initializers.py: 21%

1import warnings 

2 

3 

4def _viztracer_init(init_kwargs): 

5 """Initialize viztracer's profiler in worker processes""" 

6 from viztracer import VizTracer 

7 

8 tracer = VizTracer(**init_kwargs) 

9 tracer.register_exit() 

10 tracer.start() 

11 

12 

13def _make_viztracer_initializer_and_initargs(): 

14 try: 

15 import viztracer 

16 

17 tracer = viztracer.get_tracer() 

18 if tracer is not None and getattr(tracer, "enable", False): 

19 # Profiler is active: introspect its configuration to 

20 # initialize the workers with the same configuration. 

21 return _viztracer_init, (tracer.init_kwargs,) 

22 except ImportError: 

23 # viztracer is not installed: nothing to do 

24 pass 

25 except Exception as e: 

26 # In case viztracer's API evolve, we do not want to crash loky but 

27 # we want to know about it to be able to update loky. 

28 warnings.warn(f"Unable to introspect viztracer state: {e}") 

29 return None, () 

30 

31 

32class _ChainedInitializer: 

33 """Compound worker initializer 

34 

35 This is meant to be used in conjunction with _chain_initializers to 

36 produce the necessary chained_args list to be passed to __call__. 

37 """ 

38 

39 def __init__(self, initializers): 

40 self._initializers = initializers 

41 

42 def __call__(self, *chained_args): 

43 for initializer, args in zip(self._initializers, chained_args): 

44 initializer(*args) 

45 

46 

47def _chain_initializers(initializer_and_args): 

48 """Convenience helper to combine a sequence of initializers. 

49 

50 If some initializers are None, they are filtered out. 

51 """ 

52 filtered_initializers = [] 

53 filtered_initargs = [] 

54 for initializer, initargs in initializer_and_args: 

55 if initializer is not None: 

56 filtered_initializers.append(initializer) 

57 filtered_initargs.append(initargs) 

58 

59 if not filtered_initializers: 

60 return None, () 

61 elif len(filtered_initializers) == 1: 

62 return filtered_initializers[0], filtered_initargs[0] 

63 else: 

64 return _ChainedInitializer(filtered_initializers), filtered_initargs 

65 

66 

67def _prepare_initializer(initializer, initargs): 

68 if initializer is not None and not callable(initializer): 

69 raise TypeError( 

70 f"initializer must be a callable, got: {initializer!r}" 

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 [ 

77 (initializer, initargs), 

78 _make_viztracer_initializer_and_initargs(), 

79 ] 

80 )