/Lib/multiprocessing/util.py

http://unladen-swallow.googlecode.com/ · Python · 291 lines · 180 code · 56 blank · 55 comment · 43 complexity · 959e3c9c78b85289a6ef2d20700fde56 MD5 · raw file

  1. #
  2. # Module providing various facilities to other parts of the package
  3. #
  4. # multiprocessing/util.py
  5. #
  6. # Copyright (c) 2006-2008, R Oudkerk --- see COPYING.txt
  7. #
  8. import itertools
  9. import weakref
  10. import atexit
  11. import threading # we want threading to install it's
  12. # cleanup function before multiprocessing does
  13. from multiprocessing.process import current_process, active_children
  14. __all__ = [
  15. 'sub_debug', 'debug', 'info', 'sub_warning', 'get_logger',
  16. 'log_to_stderr', 'get_temp_dir', 'register_after_fork',
  17. 'is_exiting', 'Finalize', 'ForkAwareThreadLock', 'ForkAwareLocal',
  18. 'SUBDEBUG', 'SUBWARNING',
  19. ]
  20. #
  21. # Logging
  22. #
  23. NOTSET = 0
  24. SUBDEBUG = 5
  25. DEBUG = 10
  26. INFO = 20
  27. SUBWARNING = 25
  28. LOGGER_NAME = 'multiprocessing'
  29. DEFAULT_LOGGING_FORMAT = '[%(levelname)s/%(processName)s] %(message)s'
  30. _logger = None
  31. _log_to_stderr = False
  32. def sub_debug(msg, *args):
  33. if _logger:
  34. _logger.log(SUBDEBUG, msg, *args)
  35. def debug(msg, *args):
  36. if _logger:
  37. _logger.log(DEBUG, msg, *args)
  38. def info(msg, *args):
  39. if _logger:
  40. _logger.log(INFO, msg, *args)
  41. def sub_warning(msg, *args):
  42. if _logger:
  43. _logger.log(SUBWARNING, msg, *args)
  44. def get_logger():
  45. '''
  46. Returns logger used by multiprocessing
  47. '''
  48. global _logger
  49. import logging, atexit
  50. logging._acquireLock()
  51. try:
  52. if not _logger:
  53. _logger = logging.getLogger(LOGGER_NAME)
  54. _logger.propagate = 0
  55. logging.addLevelName(SUBDEBUG, 'SUBDEBUG')
  56. logging.addLevelName(SUBWARNING, 'SUBWARNING')
  57. # XXX multiprocessing should cleanup before logging
  58. if hasattr(atexit, 'unregister'):
  59. atexit.unregister(_exit_function)
  60. atexit.register(_exit_function)
  61. else:
  62. atexit._exithandlers.remove((_exit_function, (), {}))
  63. atexit._exithandlers.append((_exit_function, (), {}))
  64. finally:
  65. logging._releaseLock()
  66. return _logger
  67. def log_to_stderr(level=None):
  68. '''
  69. Turn on logging and add a handler which prints to stderr
  70. '''
  71. global _log_to_stderr
  72. import logging
  73. logger = get_logger()
  74. formatter = logging.Formatter(DEFAULT_LOGGING_FORMAT)
  75. handler = logging.StreamHandler()
  76. handler.setFormatter(formatter)
  77. logger.addHandler(handler)
  78. if level:
  79. logger.setLevel(level)
  80. _log_to_stderr = True
  81. return _logger
  82. #
  83. # Function returning a temp directory which will be removed on exit
  84. #
  85. def get_temp_dir():
  86. # get name of a temp directory which will be automatically cleaned up
  87. if current_process()._tempdir is None:
  88. import shutil, tempfile
  89. tempdir = tempfile.mkdtemp(prefix='pymp-')
  90. info('created temp directory %s', tempdir)
  91. Finalize(None, shutil.rmtree, args=[tempdir], exitpriority=-100)
  92. current_process()._tempdir = tempdir
  93. return current_process()._tempdir
  94. #
  95. # Support for reinitialization of objects when bootstrapping a child process
  96. #
  97. _afterfork_registry = weakref.WeakValueDictionary()
  98. _afterfork_counter = itertools.count()
  99. def _run_after_forkers():
  100. items = list(_afterfork_registry.items())
  101. items.sort()
  102. for (index, ident, func), obj in items:
  103. try:
  104. func(obj)
  105. except Exception, e:
  106. info('after forker raised exception %s', e)
  107. def register_after_fork(obj, func):
  108. _afterfork_registry[(_afterfork_counter.next(), id(obj), func)] = obj
  109. #
  110. # Finalization using weakrefs
  111. #
  112. _finalizer_registry = {}
  113. _finalizer_counter = itertools.count()
  114. class Finalize(object):
  115. '''
  116. Class which supports object finalization using weakrefs
  117. '''
  118. def __init__(self, obj, callback, args=(), kwargs=None, exitpriority=None):
  119. assert exitpriority is None or type(exitpriority) is int
  120. if obj is not None:
  121. self._weakref = weakref.ref(obj, self)
  122. else:
  123. assert exitpriority is not None
  124. self._callback = callback
  125. self._args = args
  126. self._kwargs = kwargs or {}
  127. self._key = (exitpriority, _finalizer_counter.next())
  128. _finalizer_registry[self._key] = self
  129. def __call__(self, wr=None):
  130. '''
  131. Run the callback unless it has already been called or cancelled
  132. '''
  133. try:
  134. del _finalizer_registry[self._key]
  135. except KeyError:
  136. sub_debug('finalizer no longer registered')
  137. else:
  138. sub_debug('finalizer calling %s with args %s and kwargs %s',
  139. self._callback, self._args, self._kwargs)
  140. res = self._callback(*self._args, **self._kwargs)
  141. self._weakref = self._callback = self._args = \
  142. self._kwargs = self._key = None
  143. return res
  144. def cancel(self):
  145. '''
  146. Cancel finalization of the object
  147. '''
  148. try:
  149. del _finalizer_registry[self._key]
  150. except KeyError:
  151. pass
  152. else:
  153. self._weakref = self._callback = self._args = \
  154. self._kwargs = self._key = None
  155. def still_active(self):
  156. '''
  157. Return whether this finalizer is still waiting to invoke callback
  158. '''
  159. return self._key in _finalizer_registry
  160. def __repr__(self):
  161. try:
  162. obj = self._weakref()
  163. except (AttributeError, TypeError):
  164. obj = None
  165. if obj is None:
  166. return '<Finalize object, dead>'
  167. x = '<Finalize object, callback=%s' % \
  168. getattr(self._callback, '__name__', self._callback)
  169. if self._args:
  170. x += ', args=' + str(self._args)
  171. if self._kwargs:
  172. x += ', kwargs=' + str(self._kwargs)
  173. if self._key[0] is not None:
  174. x += ', exitprority=' + str(self._key[0])
  175. return x + '>'
  176. def _run_finalizers(minpriority=None):
  177. '''
  178. Run all finalizers whose exit priority is not None and at least minpriority
  179. Finalizers with highest priority are called first; finalizers with
  180. the same priority will be called in reverse order of creation.
  181. '''
  182. if minpriority is None:
  183. f = lambda p : p[0][0] is not None
  184. else:
  185. f = lambda p : p[0][0] is not None and p[0][0] >= minpriority
  186. items = [x for x in _finalizer_registry.items() if f(x)]
  187. items.sort(reverse=True)
  188. for key, finalizer in items:
  189. sub_debug('calling %s', finalizer)
  190. try:
  191. finalizer()
  192. except Exception:
  193. import traceback
  194. traceback.print_exc()
  195. if minpriority is None:
  196. _finalizer_registry.clear()
  197. #
  198. # Clean up on exit
  199. #
  200. def is_exiting():
  201. '''
  202. Returns true if the process is shutting down
  203. '''
  204. return _exiting or _exiting is None
  205. _exiting = False
  206. def _exit_function():
  207. global _exiting
  208. info('process shutting down')
  209. debug('running all "atexit" finalizers with priority >= 0')
  210. _run_finalizers(0)
  211. for p in active_children():
  212. if p._daemonic:
  213. info('calling terminate() for daemon %s', p.name)
  214. p._popen.terminate()
  215. for p in active_children():
  216. info('calling join() for process %s', p.name)
  217. p.join()
  218. debug('running the remaining "atexit" finalizers')
  219. _run_finalizers()
  220. atexit.register(_exit_function)
  221. #
  222. # Some fork aware types
  223. #
  224. class ForkAwareThreadLock(object):
  225. def __init__(self):
  226. self._lock = threading.Lock()
  227. self.acquire = self._lock.acquire
  228. self.release = self._lock.release
  229. register_after_fork(self, ForkAwareThreadLock.__init__)
  230. class ForkAwareLocal(threading.local):
  231. def __init__(self):
  232. register_after_fork(self, lambda obj : obj.__dict__.clear())
  233. def __reduce__(self):
  234. return type(self), ()