Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/joblib/executor.py: 37%

1"""Utility function to construct a loky.ReusableExecutor with custom pickler. 

2 

3This module provides efficient ways of working with data stored in 

4shared memory with numpy.memmap arrays without inducing any memory 

5copy between the parent and child processes. 

6""" 

7# Author: Thomas Moreau <thomas.moreau.2010@gmail.com> 

8# Copyright: 2017, Thomas Moreau 

9# License: BSD 3 clause 

10 

11import threading 

12 

13from ._memmapping_reducer import TemporaryResourcesManager, get_memmapping_reducers 

14from .externals.loky.reusable_executor import _ReusablePoolExecutor 

15 

16_local_executor_args = threading.local() 

17 

18 

19def get_memmapping_executor(n_jobs, **kwargs): 

20 return MemmappingExecutor.get_memmapping_executor(n_jobs, **kwargs) 

21 

22 

23class MemmappingExecutor(_ReusablePoolExecutor): 

24 @classmethod 

25 def get_memmapping_executor( 

26 cls, 

27 n_jobs, 

28 timeout=300, 

29 initializer=None, 

30 initargs=(), 

31 env=None, 

32 temp_folder=None, 

33 context_id=None, 

34 **backend_args, 

35 ): 

36 """Factory for ReusableExecutor with automatic memmapping for large 

37 numpy arrays. 

38 """ 

39 # Check if we can reuse the executor here instead of deferring the test 

40 # to loky as the reducers are objects that changes at each call. 

41 executor_args = backend_args.copy() 

42 executor_args.update(env if env else {}) 

43 executor_args.update( 

44 dict(timeout=timeout, initializer=initializer, initargs=initargs) 

45 ) 

46 current_args = getattr(_local_executor_args, "args", None) 

47 reuse = current_args is None or current_args == executor_args 

48 _local_executor_args.args = executor_args 

49 

50 # Propagate context_id to avoid registering an unused default context 

51 # folder. 

52 manager = TemporaryResourcesManager(temp_folder, context_id=context_id) 

53 

54 # reducers access the temporary folder in which to store temporary 

55 # pickles through a call to manager.resolve_temp_folder_name. resolving 

56 # the folder name dynamically is useful to use different folders across 

57 # calls of a same reusable executor 

58 job_reducers, result_reducers = get_memmapping_reducers( 

59 unlink_on_gc_collect=True, 

60 temp_folder_resolver=manager.resolve_temp_folder_name, 

61 **backend_args, 

62 ) 

63 _executor, executor_is_reused = super().get_reusable_executor( 

64 n_jobs, 

65 job_reducers=job_reducers, 

66 result_reducers=result_reducers, 

67 reuse=reuse, 

68 timeout=timeout, 

69 initializer=initializer, 

70 initargs=initargs, 

71 env=env, 

72 ) 

73 

74 if not executor_is_reused: 

75 # Only set a _temp_folder_manager for new executors. Reused 

76 # executors already have a _temporary_folder_manager that must not 

77 # be re-assigned like that because it is referenced in various 

78 # places in the reducing machinery of the executor. 

79 _executor._temp_folder_manager = manager 

80 else: 

81 # The discarded manager already registered a folder and finalizer 

82 manager._clean_temporary_resources() 

83 

84 if context_id is not None: 

85 # Only register the specified context once we know which manager 

86 # the current executor is using, in order to not register an atexit 

87 # finalizer twice for the same folder. 

88 _executor._temp_folder_manager.register_new_context(context_id) 

89 

90 return _executor 

91 

92 def terminate(self, kill_workers=False): 

93 self.shutdown(kill_workers=kill_workers) 

94 

95 # When workers are killed in a brutal manner, they cannot execute the 

96 # finalizer of their shared memmaps. The refcount of those memmaps may 

97 # be off by an unknown number, so instead of decref'ing them, we force 

98 # delete the whole temporary folder, and unregister them. There is no 

99 # risk of PermissionError at folder deletion because at this 

100 # point, all child processes are dead, so all references to temporary 

101 # memmaps are closed. Otherwise, just try to delete as much as possible 

102 # with allow_non_empty=True but if we can't, it will be clean up later 

103 # on by the resource_tracker. 

104 self._temp_folder_manager._clean_temporary_resources( 

105 force=kill_workers, allow_non_empty=True 

106 ) 

107 

108 @property 

109 def _temp_folder(self): 

110 # Legacy property in tests. could be removed if we refactored the 

111 # memmapping tests. SHOULD ONLY BE USED IN TESTS! 

112 # We cache this property because it is called late in the tests - at 

113 # this point, all context have been unregistered, and 

114 # resolve_temp_folder_name raises an error. 

115 if getattr(self, "_cached_temp_folder", None) is not None: 

116 return self._cached_temp_folder 

117 else: 

118 self._cached_temp_folder = ( 

119 self._temp_folder_manager.resolve_temp_folder_name() 

120 ) # noqa 

121 return self._cached_temp_folder 

122 

123 

124class _TestingMemmappingExecutor(MemmappingExecutor): 

125 """Wrapper around ReusableExecutor to ease memmapping testing with Pool 

126 and Executor. This is only for testing purposes. 

127 

128 """ 

129 

130 def apply_async(self, func, args): 

131 """Schedule a func to be run""" 

132 future = self.submit(func, *args) 

133 future.get = future.result 

134 return future 

135 

136 def map(self, f, *args): 

137 return list(super().map(f, *args))