pool.py 65 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006
  1. # -*- coding: utf-8 -*-
  2. #
  3. # Module providing the `Pool` class for managing a process pool
  4. #
  5. # multiprocessing/pool.py
  6. #
  7. # Copyright (c) 2006-2008, R Oudkerk
  8. # Licensed to PSF under a Contributor Agreement.
  9. #
  10. from __future__ import absolute_import
  11. #
  12. # Imports
  13. #
  14. import copy
  15. import errno
  16. import itertools
  17. import os
  18. import platform
  19. import signal
  20. import sys
  21. import threading
  22. import time
  23. import warnings
  24. from collections import deque
  25. from functools import partial
  26. from . import cpu_count, get_context
  27. from . import util
  28. from .common import (
  29. TERM_SIGNAL, human_status, pickle_loads, reset_signals, restart_state,
  30. )
  31. from .compat import get_errno, mem_rss, send_offset
  32. from .einfo import ExceptionInfo
  33. from .dummy import DummyProcess
  34. from .exceptions import (
  35. CoroStop,
  36. RestartFreqExceeded,
  37. SoftTimeLimitExceeded,
  38. Terminated,
  39. TimeLimitExceeded,
  40. TimeoutError,
  41. WorkerLostError,
  42. )
  43. from .five import Empty, Queue, range, values, reraise, monotonic
  44. from .util import Finalize, debug
  45. MAXMEM_USED_FMT = """\
  46. child process exiting after exceeding memory limit ({0}KiB / {1}KiB)
  47. """
  48. PY3 = sys.version_info[0] == 3
  49. if platform.system() == 'Windows': # pragma: no cover
  50. # On Windows os.kill calls TerminateProcess which cannot be
  51. # handled by # any process, so this is needed to terminate the task
  52. # *and its children* (if any).
  53. from ._win import kill_processtree as _kill # noqa
  54. SIGKILL = TERM_SIGNAL
  55. else:
  56. from os import kill as _kill # noqa
  57. SIGKILL = signal.SIGKILL
  58. try:
  59. TIMEOUT_MAX = threading.TIMEOUT_MAX
  60. except AttributeError: # pragma: no cover
  61. TIMEOUT_MAX = 1e10 # noqa
  62. if sys.version_info >= (3, 3):
  63. _Semaphore = threading.Semaphore
  64. else:
  65. # Semaphore is a factory function pointing to _Semaphore
  66. _Semaphore = threading._Semaphore # noqa
  67. #
  68. # Constants representing the state of a pool
  69. #
  70. RUN = 0
  71. CLOSE = 1
  72. TERMINATE = 2
  73. #
  74. # Constants representing the state of a job
  75. #
  76. ACK = 0
  77. READY = 1
  78. TASK = 2
  79. NACK = 3
  80. DEATH = 4
  81. #
  82. # Exit code constants
  83. #
  84. EX_OK = 0
  85. EX_FAILURE = 1
  86. EX_RECYCLE = 0x9B
  87. # Signal used for soft time limits.
  88. SIG_SOFT_TIMEOUT = getattr(signal, "SIGUSR1", None)
  89. #
  90. # Miscellaneous
  91. #
  92. LOST_WORKER_TIMEOUT = 10.0
  93. EX_OK = getattr(os, "EX_OK", 0)
  94. job_counter = itertools.count()
  95. Lock = threading.Lock
  96. def _get_send_offset(connection):
  97. try:
  98. native = connection.send_offset
  99. except AttributeError:
  100. native = None
  101. if native is None:
  102. return partial(send_offset, connection.fileno())
  103. return native
  104. def mapstar(args):
  105. return list(map(*args))
  106. def starmapstar(args):
  107. return list(itertools.starmap(args[0], args[1]))
  108. def error(msg, *args, **kwargs):
  109. util.get_logger().error(msg, *args, **kwargs)
  110. def stop_if_not_current(thread, timeout=None):
  111. if thread is not threading.current_thread():
  112. thread.stop(timeout)
  113. class LaxBoundedSemaphore(_Semaphore):
  114. """Semaphore that checks that # release is <= # acquires,
  115. but ignores if # releases >= value."""
  116. def shrink(self):
  117. self._initial_value -= 1
  118. self.acquire()
  119. if PY3:
  120. def __init__(self, value=1, verbose=None):
  121. _Semaphore.__init__(self, value)
  122. self._initial_value = value
  123. def grow(self):
  124. with self._cond:
  125. self._initial_value += 1
  126. self._value += 1
  127. self._cond.notify()
  128. def release(self):
  129. cond = self._cond
  130. with cond:
  131. if self._value < self._initial_value:
  132. self._value += 1
  133. cond.notify_all()
  134. def clear(self):
  135. while self._value < self._initial_value:
  136. _Semaphore.release(self)
  137. else:
  138. def __init__(self, value=1, verbose=None):
  139. _Semaphore.__init__(self, value, verbose)
  140. self._initial_value = value
  141. def grow(self):
  142. cond = self._Semaphore__cond
  143. with cond:
  144. self._initial_value += 1
  145. self._Semaphore__value += 1
  146. cond.notify()
  147. def release(self): # noqa
  148. cond = self._Semaphore__cond
  149. with cond:
  150. if self._Semaphore__value < self._initial_value:
  151. self._Semaphore__value += 1
  152. cond.notifyAll()
  153. def clear(self): # noqa
  154. while self._Semaphore__value < self._initial_value:
  155. _Semaphore.release(self)
  156. #
  157. # Exceptions
  158. #
  159. class MaybeEncodingError(Exception):
  160. """Wraps possible unpickleable errors, so they can be
  161. safely sent through the socket."""
  162. def __init__(self, exc, value):
  163. self.exc = repr(exc)
  164. self.value = repr(value)
  165. super(MaybeEncodingError, self).__init__(self.exc, self.value)
  166. def __repr__(self):
  167. return "<%s: %s>" % (self.__class__.__name__, str(self))
  168. def __str__(self):
  169. return "Error sending result: '%r'. Reason: '%r'." % (
  170. self.value, self.exc)
  171. class WorkersJoined(Exception):
  172. """All workers have terminated."""
  173. def soft_timeout_sighandler(signum, frame):
  174. raise SoftTimeLimitExceeded()
  175. #
  176. # Code run by worker processes
  177. #
  178. class Worker(object):
  179. def __init__(self, inq, outq, synq=None, initializer=None, initargs=(),
  180. maxtasks=None, sentinel=None, on_exit=None,
  181. sigprotection=True, wrap_exception=True,
  182. max_memory_per_child=None):
  183. assert maxtasks is None or (type(maxtasks) == int and maxtasks > 0)
  184. self.initializer = initializer
  185. self.initargs = initargs
  186. self.maxtasks = maxtasks
  187. self.max_memory_per_child = max_memory_per_child
  188. self._shutdown = sentinel
  189. self.on_exit = on_exit
  190. self.sigprotection = sigprotection
  191. self.inq, self.outq, self.synq = inq, outq, synq
  192. self.wrap_exception = wrap_exception # XXX cannot disable yet
  193. self.contribute_to_object(self)
  194. def contribute_to_object(self, obj):
  195. obj.inq, obj.outq, obj.synq = self.inq, self.outq, self.synq
  196. obj.inqW_fd = self.inq._writer.fileno() # inqueue write fd
  197. obj.outqR_fd = self.outq._reader.fileno() # outqueue read fd
  198. if self.synq:
  199. obj.synqR_fd = self.synq._reader.fileno() # synqueue read fd
  200. obj.synqW_fd = self.synq._writer.fileno() # synqueue write fd
  201. obj.send_syn_offset = _get_send_offset(self.synq._writer)
  202. else:
  203. obj.synqR_fd = obj.synqW_fd = obj._send_syn_offset = None
  204. obj._quick_put = self.inq._writer.send
  205. obj._quick_get = self.outq._reader.recv
  206. obj.send_job_offset = _get_send_offset(self.inq._writer)
  207. return obj
  208. def __reduce__(self):
  209. return self.__class__, (
  210. self.inq, self.outq, self.synq, self.initializer,
  211. self.initargs, self.maxtasks, self._shutdown, self.on_exit,
  212. self.sigprotection, self.wrap_exception,
  213. )
  214. def __call__(self):
  215. _exit = sys.exit
  216. _exitcode = [None]
  217. def exit(status=None):
  218. _exitcode[0] = status
  219. return _exit()
  220. sys.exit = exit
  221. pid = os.getpid()
  222. self._make_child_methods()
  223. self.after_fork()
  224. self.on_loop_start(pid=pid) # callback on loop start
  225. try:
  226. sys.exit(self.workloop(pid=pid))
  227. except Exception as exc:
  228. error('Pool process %r error: %r', self, exc, exc_info=1)
  229. self._do_exit(pid, _exitcode[0], exc)
  230. finally:
  231. self._do_exit(pid, _exitcode[0], None)
  232. def _do_exit(self, pid, exitcode, exc=None):
  233. if exitcode is None:
  234. exitcode = EX_FAILURE if exc else EX_OK
  235. if self.on_exit is not None:
  236. self.on_exit(pid, exitcode)
  237. if sys.platform != 'win32':
  238. try:
  239. self.outq.put((DEATH, (pid, exitcode)))
  240. time.sleep(1)
  241. finally:
  242. os._exit(exitcode)
  243. else:
  244. os._exit(exitcode)
  245. def on_loop_start(self, pid):
  246. pass
  247. def prepare_result(self, result):
  248. return result
  249. def workloop(self, debug=debug, now=monotonic, pid=None):
  250. pid = pid or os.getpid()
  251. put = self.outq.put
  252. inqW_fd = self.inqW_fd
  253. synqW_fd = self.synqW_fd
  254. maxtasks = self.maxtasks
  255. max_memory_per_child = self.max_memory_per_child or 0
  256. prepare_result = self.prepare_result
  257. wait_for_job = self.wait_for_job
  258. _wait_for_syn = self.wait_for_syn
  259. def wait_for_syn(jid):
  260. i = 0
  261. while 1:
  262. if i > 60:
  263. error('!!!WAIT FOR ACK TIMEOUT: job:%r fd:%r!!!',
  264. jid, self.synq._reader.fileno(), exc_info=1)
  265. req = _wait_for_syn()
  266. if req:
  267. type_, args = req
  268. if type_ == NACK:
  269. return False
  270. assert type_ == ACK
  271. return True
  272. i += 1
  273. completed = 0
  274. while maxtasks is None or (maxtasks and completed < maxtasks):
  275. req = wait_for_job()
  276. if req:
  277. type_, args_ = req
  278. assert type_ == TASK
  279. job, i, fun, args, kwargs = args_
  280. put((ACK, (job, i, now(), pid, synqW_fd)))
  281. if _wait_for_syn:
  282. confirm = wait_for_syn(job)
  283. if not confirm:
  284. continue # received NACK
  285. try:
  286. result = (True, prepare_result(fun(*args, **kwargs)))
  287. except Exception:
  288. result = (False, ExceptionInfo())
  289. try:
  290. put((READY, (job, i, result, inqW_fd)))
  291. except Exception as exc:
  292. _, _, tb = sys.exc_info()
  293. try:
  294. wrapped = MaybeEncodingError(exc, result[1])
  295. einfo = ExceptionInfo((
  296. MaybeEncodingError, wrapped, tb,
  297. ))
  298. put((READY, (job, i, (False, einfo), inqW_fd)))
  299. finally:
  300. del(tb)
  301. completed += 1
  302. if max_memory_per_child > 0:
  303. used_kb = mem_rss()
  304. if used_kb <= 0:
  305. error('worker unable to determine memory usage')
  306. if used_kb > 0 and used_kb > max_memory_per_child:
  307. error(MAXMEM_USED_FMT.format(
  308. used_kb, max_memory_per_child))
  309. return EX_RECYCLE
  310. debug('worker exiting after %d tasks', completed)
  311. if maxtasks:
  312. return EX_RECYCLE if completed == maxtasks else EX_FAILURE
  313. return EX_OK
  314. def after_fork(self):
  315. if hasattr(self.inq, '_writer'):
  316. self.inq._writer.close()
  317. if hasattr(self.outq, '_reader'):
  318. self.outq._reader.close()
  319. if self.initializer is not None:
  320. self.initializer(*self.initargs)
  321. # Make sure all exiting signals call finally: blocks.
  322. # This is important for the semaphore to be released.
  323. reset_signals(full=self.sigprotection)
  324. # install signal handler for soft timeouts.
  325. if SIG_SOFT_TIMEOUT is not None:
  326. signal.signal(SIG_SOFT_TIMEOUT, soft_timeout_sighandler)
  327. try:
  328. signal.signal(signal.SIGINT, signal.SIG_IGN)
  329. except AttributeError:
  330. pass
  331. def _make_recv_method(self, conn):
  332. get = conn.get
  333. if hasattr(conn, '_reader'):
  334. _poll = conn._reader.poll
  335. if hasattr(conn, 'get_payload') and conn.get_payload:
  336. get_payload = conn.get_payload
  337. def _recv(timeout, loads=pickle_loads):
  338. return True, loads(get_payload())
  339. else:
  340. def _recv(timeout): # noqa
  341. if _poll(timeout):
  342. return True, get()
  343. return False, None
  344. else:
  345. def _recv(timeout): # noqa
  346. try:
  347. return True, get(timeout=timeout)
  348. except Queue.Empty:
  349. return False, None
  350. return _recv
  351. def _make_child_methods(self, loads=pickle_loads):
  352. self.wait_for_job = self._make_protected_receive(self.inq)
  353. self.wait_for_syn = (self._make_protected_receive(self.synq)
  354. if self.synq else None)
  355. def _make_protected_receive(self, conn):
  356. _receive = self._make_recv_method(conn)
  357. should_shutdown = self._shutdown.is_set if self._shutdown else None
  358. def receive(debug=debug):
  359. if should_shutdown and should_shutdown():
  360. debug('worker got sentinel -- exiting')
  361. raise SystemExit(EX_OK)
  362. try:
  363. ready, req = _receive(1.0)
  364. if not ready:
  365. return None
  366. except (EOFError, IOError) as exc:
  367. if get_errno(exc) == errno.EINTR:
  368. return None # interrupted, maybe by gdb
  369. debug('worker got %s -- exiting', type(exc).__name__)
  370. raise SystemExit(EX_FAILURE)
  371. if req is None:
  372. debug('worker got sentinel -- exiting')
  373. raise SystemExit(EX_FAILURE)
  374. return req
  375. return receive
  376. #
  377. # Class representing a process pool
  378. #
  379. class PoolThread(DummyProcess):
  380. def __init__(self, *args, **kwargs):
  381. DummyProcess.__init__(self)
  382. self._state = RUN
  383. self._was_started = False
  384. self.daemon = True
  385. def run(self):
  386. try:
  387. return self.body()
  388. except RestartFreqExceeded as exc:
  389. error("Thread %r crashed: %r", type(self).__name__, exc,
  390. exc_info=1)
  391. _kill(os.getpid(), TERM_SIGNAL)
  392. sys.exit()
  393. except Exception as exc:
  394. error("Thread %r crashed: %r", type(self).__name__, exc,
  395. exc_info=1)
  396. os._exit(1)
  397. def start(self, *args, **kwargs):
  398. self._was_started = True
  399. super(PoolThread, self).start(*args, **kwargs)
  400. def on_stop_not_started(self):
  401. pass
  402. def stop(self, timeout=None):
  403. if self._was_started:
  404. self.join(timeout)
  405. return
  406. self.on_stop_not_started()
  407. def terminate(self):
  408. self._state = TERMINATE
  409. def close(self):
  410. self._state = CLOSE
  411. class Supervisor(PoolThread):
  412. def __init__(self, pool):
  413. self.pool = pool
  414. super(Supervisor, self).__init__()
  415. def body(self):
  416. debug('worker handler starting')
  417. time.sleep(0.8)
  418. pool = self.pool
  419. try:
  420. # do a burst at startup to verify that we can start
  421. # our pool processes, and in that time we lower
  422. # the max restart frequency.
  423. prev_state = pool.restart_state
  424. pool.restart_state = restart_state(10 * pool._processes, 1)
  425. for _ in range(10):
  426. if self._state == RUN and pool._state == RUN:
  427. pool._maintain_pool()
  428. time.sleep(0.1)
  429. # Keep maintaing workers until the cache gets drained, unless
  430. # the pool is termianted
  431. pool.restart_state = prev_state
  432. while self._state == RUN and pool._state == RUN:
  433. pool._maintain_pool()
  434. time.sleep(0.8)
  435. except RestartFreqExceeded:
  436. pool.close()
  437. pool.join()
  438. raise
  439. debug('worker handler exiting')
  440. class TaskHandler(PoolThread):
  441. def __init__(self, taskqueue, put, outqueue, pool, cache):
  442. self.taskqueue = taskqueue
  443. self.put = put
  444. self.outqueue = outqueue
  445. self.pool = pool
  446. self.cache = cache
  447. super(TaskHandler, self).__init__()
  448. def body(self):
  449. cache = self.cache
  450. taskqueue = self.taskqueue
  451. put = self.put
  452. for taskseq, set_length in iter(taskqueue.get, None):
  453. task = None
  454. i = -1
  455. try:
  456. for i, task in enumerate(taskseq):
  457. if self._state:
  458. debug('task handler found thread._state != RUN')
  459. break
  460. try:
  461. put(task)
  462. except IOError:
  463. debug('could not put task on queue')
  464. break
  465. except Exception:
  466. job, ind = task[:2]
  467. try:
  468. cache[job]._set(ind, (False, ExceptionInfo()))
  469. except KeyError:
  470. pass
  471. else:
  472. if set_length:
  473. debug('doing set_length()')
  474. set_length(i + 1)
  475. continue
  476. break
  477. except Exception:
  478. job, ind = task[:2] if task else (0, 0)
  479. if job in cache:
  480. cache[job]._set(ind + 1, (False, ExceptionInfo()))
  481. if set_length:
  482. util.debug('doing set_length()')
  483. set_length(i + 1)
  484. else:
  485. debug('task handler got sentinel')
  486. self.tell_others()
  487. def tell_others(self):
  488. outqueue = self.outqueue
  489. put = self.put
  490. pool = self.pool
  491. try:
  492. # tell result handler to finish when cache is empty
  493. debug('task handler sending sentinel to result handler')
  494. outqueue.put(None)
  495. # tell workers there is no more work
  496. debug('task handler sending sentinel to workers')
  497. for p in pool:
  498. put(None)
  499. except IOError:
  500. debug('task handler got IOError when sending sentinels')
  501. debug('task handler exiting')
  502. def on_stop_not_started(self):
  503. self.tell_others()
  504. class TimeoutHandler(PoolThread):
  505. def __init__(self, processes, cache, t_soft, t_hard):
  506. self.processes = processes
  507. self.cache = cache
  508. self.t_soft = t_soft
  509. self.t_hard = t_hard
  510. self._it = None
  511. super(TimeoutHandler, self).__init__()
  512. def _process_by_pid(self, pid):
  513. return next((
  514. (proc, i) for i, proc in enumerate(self.processes)
  515. if proc.pid == pid
  516. ), (None, None))
  517. def on_soft_timeout(self, job):
  518. debug('soft time limit exceeded for %r', job)
  519. process, _index = self._process_by_pid(job._worker_pid)
  520. if not process:
  521. return
  522. # Run timeout callback
  523. job.handle_timeout(soft=True)
  524. try:
  525. _kill(job._worker_pid, SIG_SOFT_TIMEOUT)
  526. except OSError as exc:
  527. if get_errno(exc) != errno.ESRCH:
  528. raise
  529. def on_hard_timeout(self, job):
  530. if job.ready():
  531. return
  532. debug('hard time limit exceeded for %r', job)
  533. # Remove from cache and set return value to an exception
  534. try:
  535. raise TimeLimitExceeded(job._timeout)
  536. except TimeLimitExceeded:
  537. job._set(job._job, (False, ExceptionInfo()))
  538. else: # pragma: no cover
  539. pass
  540. # Remove from _pool
  541. process, _index = self._process_by_pid(job._worker_pid)
  542. # Run timeout callback
  543. job.handle_timeout(soft=False)
  544. if process:
  545. self._trywaitkill(process)
  546. def _trywaitkill(self, worker):
  547. debug('timeout: sending TERM to %s', worker._name)
  548. try:
  549. if os.getpgid(worker.pid) == worker.pid:
  550. debug("worker %s is a group leader. It is safe to kill (SIGTERM) the whole group", worker.pid)
  551. os.killpg(os.getpgid(worker.pid), signal.SIGTERM)
  552. else:
  553. worker.terminate()
  554. except OSError:
  555. pass
  556. else:
  557. if worker._popen.wait(timeout=0.1):
  558. return
  559. debug('timeout: TERM timed-out, now sending KILL to %s', worker._name)
  560. try:
  561. if os.getpgid(worker.pid) == worker.pid:
  562. debug("worker %s is a group leader. It is safe to kill (SIGKILL) the whole group", worker.pid)
  563. os.killpg(os.getpgid(worker.pid), signal.SIGKILL)
  564. else:
  565. _kill(worker.pid, SIGKILL)
  566. except OSError:
  567. pass
  568. def handle_timeouts(self):
  569. cache = copy.deepcopy(self.cache)
  570. t_hard, t_soft = self.t_hard, self.t_soft
  571. dirty = set()
  572. on_soft_timeout = self.on_soft_timeout
  573. on_hard_timeout = self.on_hard_timeout
  574. def _timed_out(start, timeout):
  575. if not start or not timeout:
  576. return False
  577. if monotonic() >= start + timeout:
  578. return True
  579. # Inner-loop
  580. while self._state == RUN:
  581. # Remove dirty items not in cache anymore
  582. if dirty:
  583. dirty = set(k for k in dirty if k in cache)
  584. for i, job in list(cache.items()):
  585. ack_time = job._time_accepted
  586. soft_timeout = job._soft_timeout
  587. if soft_timeout is None:
  588. soft_timeout = t_soft
  589. hard_timeout = job._timeout
  590. if hard_timeout is None:
  591. hard_timeout = t_hard
  592. if _timed_out(ack_time, hard_timeout):
  593. on_hard_timeout(job)
  594. elif i not in dirty and _timed_out(ack_time, soft_timeout):
  595. on_soft_timeout(job)
  596. dirty.add(i)
  597. yield
  598. def body(self):
  599. while self._state == RUN:
  600. try:
  601. for _ in self.handle_timeouts():
  602. time.sleep(1.0) # don't spin
  603. except CoroStop:
  604. break
  605. debug('timeout handler exiting')
  606. def handle_event(self, *args):
  607. if self._it is None:
  608. self._it = self.handle_timeouts()
  609. try:
  610. next(self._it)
  611. except StopIteration:
  612. self._it = None
  613. class ResultHandler(PoolThread):
  614. def __init__(self, outqueue, get, cache, poll,
  615. join_exited_workers, putlock, restart_state,
  616. check_timeouts, on_job_ready):
  617. self.outqueue = outqueue
  618. self.get = get
  619. self.cache = cache
  620. self.poll = poll
  621. self.join_exited_workers = join_exited_workers
  622. self.putlock = putlock
  623. self.restart_state = restart_state
  624. self._it = None
  625. self._shutdown_complete = False
  626. self.check_timeouts = check_timeouts
  627. self.on_job_ready = on_job_ready
  628. self._make_methods()
  629. super(ResultHandler, self).__init__()
  630. def on_stop_not_started(self):
  631. # used when pool started without result handler thread.
  632. self.finish_at_shutdown(handle_timeouts=True)
  633. def _make_methods(self):
  634. cache = self.cache
  635. putlock = self.putlock
  636. restart_state = self.restart_state
  637. on_job_ready = self.on_job_ready
  638. def on_ack(job, i, time_accepted, pid, synqW_fd):
  639. restart_state.R = 0
  640. try:
  641. cache[job]._ack(i, time_accepted, pid, synqW_fd)
  642. except (KeyError, AttributeError):
  643. # Object gone or doesn't support _ack (e.g. IMAPIterator).
  644. pass
  645. def on_ready(job, i, obj, inqW_fd):
  646. if on_job_ready is not None:
  647. on_job_ready(job, i, obj, inqW_fd)
  648. try:
  649. item = cache[job]
  650. except KeyError:
  651. return
  652. if not item.ready():
  653. if putlock is not None:
  654. putlock.release()
  655. try:
  656. item._set(i, obj)
  657. except KeyError:
  658. pass
  659. def on_death(pid, exitcode):
  660. try:
  661. os.kill(pid, TERM_SIGNAL)
  662. except OSError as exc:
  663. if get_errno(exc) != errno.ESRCH:
  664. raise
  665. state_handlers = self.state_handlers = {
  666. ACK: on_ack, READY: on_ready, DEATH: on_death
  667. }
  668. def on_state_change(task):
  669. state, args = task
  670. try:
  671. state_handlers[state](*args)
  672. except KeyError:
  673. debug("Unknown job state: %s (args=%s)", state, args)
  674. self.on_state_change = on_state_change
  675. def _process_result(self, timeout=1.0):
  676. poll = self.poll
  677. on_state_change = self.on_state_change
  678. while 1:
  679. try:
  680. ready, task = poll(timeout)
  681. except (IOError, EOFError) as exc:
  682. debug('result handler got %r -- exiting', exc)
  683. raise CoroStop()
  684. if self._state:
  685. assert self._state == TERMINATE
  686. debug('result handler found thread._state=TERMINATE')
  687. raise CoroStop()
  688. if ready:
  689. if task is None:
  690. debug('result handler got sentinel')
  691. raise CoroStop()
  692. on_state_change(task)
  693. if timeout != 0: # blocking
  694. break
  695. else:
  696. break
  697. yield
  698. def handle_event(self, fileno=None, events=None):
  699. if self._state == RUN:
  700. if self._it is None:
  701. self._it = self._process_result(0) # non-blocking
  702. try:
  703. next(self._it)
  704. except (StopIteration, CoroStop):
  705. self._it = None
  706. def body(self):
  707. debug('result handler starting')
  708. try:
  709. while self._state == RUN:
  710. try:
  711. for _ in self._process_result(1.0): # blocking
  712. pass
  713. except CoroStop:
  714. break
  715. finally:
  716. self.finish_at_shutdown()
  717. def finish_at_shutdown(self, handle_timeouts=False):
  718. self._shutdown_complete = True
  719. get = self.get
  720. outqueue = self.outqueue
  721. cache = self.cache
  722. poll = self.poll
  723. join_exited_workers = self.join_exited_workers
  724. check_timeouts = self.check_timeouts
  725. on_state_change = self.on_state_change
  726. time_terminate = None
  727. while cache and self._state != TERMINATE:
  728. if check_timeouts is not None:
  729. check_timeouts()
  730. try:
  731. ready, task = poll(1.0)
  732. except (IOError, EOFError) as exc:
  733. debug('result handler got %r -- exiting', exc)
  734. return
  735. if ready:
  736. if task is None:
  737. debug('result handler ignoring extra sentinel')
  738. continue
  739. on_state_change(task)
  740. try:
  741. join_exited_workers(shutdown=True)
  742. except WorkersJoined:
  743. now = monotonic()
  744. if not time_terminate:
  745. time_terminate = now
  746. else:
  747. if now - time_terminate > 5.0:
  748. debug('result handler exiting: timed out')
  749. break
  750. debug('result handler: all workers terminated, '
  751. 'timeout in %ss',
  752. abs(min(now - time_terminate - 5.0, 0)))
  753. if hasattr(outqueue, '_reader'):
  754. debug('ensuring that outqueue is not full')
  755. # If we don't make room available in outqueue then
  756. # attempts to add the sentinel (None) to outqueue may
  757. # block. There is guaranteed to be no more than 2 sentinels.
  758. try:
  759. for i in range(10):
  760. if not outqueue._reader.poll():
  761. break
  762. get()
  763. except (IOError, EOFError):
  764. pass
  765. debug('result handler exiting: len(cache)=%s, thread._state=%s',
  766. len(cache), self._state)
  767. class Pool(object):
  768. '''
  769. Class which supports an async version of applying functions to arguments.
  770. '''
  771. _wrap_exception = True
  772. Worker = Worker
  773. Supervisor = Supervisor
  774. TaskHandler = TaskHandler
  775. TimeoutHandler = TimeoutHandler
  776. ResultHandler = ResultHandler
  777. SoftTimeLimitExceeded = SoftTimeLimitExceeded
  778. def __init__(self, processes=None, initializer=None, initargs=(),
  779. maxtasksperchild=None, timeout=None, soft_timeout=None,
  780. lost_worker_timeout=None,
  781. max_restarts=None, max_restart_freq=1,
  782. on_process_up=None,
  783. on_process_down=None,
  784. on_timeout_set=None,
  785. on_timeout_cancel=None,
  786. threads=True,
  787. semaphore=None,
  788. putlocks=False,
  789. allow_restart=False,
  790. synack=False,
  791. on_process_exit=None,
  792. context=None,
  793. max_memory_per_child=None,
  794. enable_timeouts=False,
  795. **kwargs):
  796. self._ctx = context or get_context()
  797. self.synack = synack
  798. self._setup_queues()
  799. self._taskqueue = Queue()
  800. self._cache = {}
  801. self._state = RUN
  802. self.timeout = timeout
  803. self.soft_timeout = soft_timeout
  804. self._maxtasksperchild = maxtasksperchild
  805. self._max_memory_per_child = max_memory_per_child
  806. self._initializer = initializer
  807. self._initargs = initargs
  808. self._on_process_exit = on_process_exit
  809. self.lost_worker_timeout = lost_worker_timeout or LOST_WORKER_TIMEOUT
  810. self.on_process_up = on_process_up
  811. self.on_process_down = on_process_down
  812. self.on_timeout_set = on_timeout_set
  813. self.on_timeout_cancel = on_timeout_cancel
  814. self.threads = threads
  815. self.readers = {}
  816. self.allow_restart = allow_restart
  817. self.enable_timeouts = bool(
  818. enable_timeouts or
  819. self.timeout is not None or
  820. self.soft_timeout is not None
  821. )
  822. if soft_timeout and SIG_SOFT_TIMEOUT is None:
  823. warnings.warn(UserWarning(
  824. "Soft timeouts are not supported: "
  825. "on this platform: It does not have the SIGUSR1 signal.",
  826. ))
  827. soft_timeout = None
  828. self._processes = self.cpu_count() if processes is None else processes
  829. self.max_restarts = max_restarts or round(self._processes * 100)
  830. self.restart_state = restart_state(max_restarts, max_restart_freq or 1)
  831. if initializer is not None and not callable(initializer):
  832. raise TypeError('initializer must be a callable')
  833. if on_process_exit is not None and not callable(on_process_exit):
  834. raise TypeError('on_process_exit must be callable')
  835. self._Process = self._ctx.Process
  836. self._pool = []
  837. self._poolctrl = {}
  838. self.putlocks = putlocks
  839. self._putlock = semaphore or LaxBoundedSemaphore(self._processes)
  840. for i in range(self._processes):
  841. self._create_worker_process(i)
  842. self._worker_handler = self.Supervisor(self)
  843. if threads:
  844. self._worker_handler.start()
  845. self._task_handler = self.TaskHandler(self._taskqueue,
  846. self._quick_put,
  847. self._outqueue,
  848. self._pool,
  849. self._cache)
  850. if threads:
  851. self._task_handler.start()
  852. self.check_timeouts = None
  853. # Thread killing timedout jobs.
  854. if self.enable_timeouts:
  855. self._timeout_handler = self.TimeoutHandler(
  856. self._pool, self._cache,
  857. self.soft_timeout, self.timeout,
  858. )
  859. self._timeout_handler_mutex = Lock()
  860. self._timeout_handler_started = False
  861. self._start_timeout_handler()
  862. # If running without threads, we need to check for timeouts
  863. # while waiting for unfinished work at shutdown.
  864. if not threads:
  865. self.check_timeouts = self._timeout_handler.handle_event
  866. else:
  867. self._timeout_handler = None
  868. self._timeout_handler_started = False
  869. self._timeout_handler_mutex = None
  870. # Thread processing results in the outqueue.
  871. self._result_handler = self.create_result_handler()
  872. self.handle_result_event = self._result_handler.handle_event
  873. if threads:
  874. self._result_handler.start()
  875. self._terminate = Finalize(
  876. self, self._terminate_pool,
  877. args=(self._taskqueue, self._inqueue, self._outqueue,
  878. self._pool, self._worker_handler, self._task_handler,
  879. self._result_handler, self._cache,
  880. self._timeout_handler,
  881. self._help_stuff_finish_args()),
  882. exitpriority=15,
  883. )
  884. def Process(self, *args, **kwds):
  885. return self._Process(*args, **kwds)
  886. def WorkerProcess(self, worker):
  887. return worker.contribute_to_object(self.Process(target=worker))
  888. def create_result_handler(self, **extra_kwargs):
  889. return self.ResultHandler(
  890. self._outqueue, self._quick_get, self._cache,
  891. self._poll_result, self._join_exited_workers,
  892. self._putlock, self.restart_state, self.check_timeouts,
  893. self.on_job_ready, **extra_kwargs
  894. )
  895. def on_job_ready(self, job, i, obj, inqW_fd):
  896. pass
  897. def _help_stuff_finish_args(self):
  898. return self._inqueue, self._task_handler, self._pool
  899. def cpu_count(self):
  900. try:
  901. return cpu_count()
  902. except NotImplementedError:
  903. return 1
  904. def handle_result_event(self, *args):
  905. return self._result_handler.handle_event(*args)
  906. def _process_register_queues(self, worker, queues):
  907. pass
  908. def _process_by_pid(self, pid):
  909. return next((
  910. (proc, i) for i, proc in enumerate(self._pool)
  911. if proc.pid == pid
  912. ), (None, None))
  913. def get_process_queues(self):
  914. return self._inqueue, self._outqueue, None
  915. def _create_worker_process(self, i):
  916. sentinel = self._ctx.Event() if self.allow_restart else None
  917. inq, outq, synq = self.get_process_queues()
  918. w = self.WorkerProcess(self.Worker(
  919. inq, outq, synq, self._initializer, self._initargs,
  920. self._maxtasksperchild, sentinel, self._on_process_exit,
  921. # Need to handle all signals if using the ipc semaphore,
  922. # to make sure the semaphore is released.
  923. sigprotection=self.threads,
  924. wrap_exception=self._wrap_exception,
  925. max_memory_per_child=self._max_memory_per_child,
  926. ))
  927. self._pool.append(w)
  928. self._process_register_queues(w, (inq, outq, synq))
  929. w.name = w.name.replace('Process', 'PoolWorker')
  930. w.daemon = True
  931. w.index = i
  932. w.start()
  933. self._poolctrl[w.pid] = sentinel
  934. if self.on_process_up:
  935. self.on_process_up(w)
  936. return w
  937. def process_flush_queues(self, worker):
  938. pass
  939. def _join_exited_workers(self, shutdown=False):
  940. """Cleanup after any worker processes which have exited due to
  941. reaching their specified lifetime. Returns True if any workers were
  942. cleaned up.
  943. """
  944. now = None
  945. # The worker may have published a result before being terminated,
  946. # but we have no way to accurately tell if it did. So we wait for
  947. # _lost_worker_timeout seconds before we mark the job with
  948. # WorkerLostError.
  949. for job in [job for job in list(self._cache.values())
  950. if not job.ready() and job._worker_lost]:
  951. now = now or monotonic()
  952. lost_time, lost_ret = job._worker_lost
  953. if now - lost_time > job._lost_worker_timeout:
  954. self.mark_as_worker_lost(job, lost_ret)
  955. if shutdown and not len(self._pool):
  956. raise WorkersJoined()
  957. cleaned, exitcodes = {}, {}
  958. for i in reversed(range(len(self._pool))):
  959. worker = self._pool[i]
  960. exitcode = worker.exitcode
  961. popen = worker._popen
  962. if popen is None or exitcode is not None:
  963. # worker exited
  964. debug('Supervisor: cleaning up worker %d', i)
  965. if popen is not None:
  966. worker.join()
  967. debug('Supervisor: worked %d joined', i)
  968. cleaned[worker.pid] = worker
  969. exitcodes[worker.pid] = exitcode
  970. if exitcode not in (EX_OK, EX_RECYCLE) and \
  971. not getattr(worker, '_controlled_termination', False):
  972. error(
  973. 'Process %r pid:%r exited with %r',
  974. worker.name, worker.pid, human_status(exitcode),
  975. exc_info=0,
  976. )
  977. self.process_flush_queues(worker)
  978. del self._pool[i]
  979. del self._poolctrl[worker.pid]
  980. if cleaned:
  981. all_pids = [w.pid for w in self._pool]
  982. for job in list(self._cache.values()):
  983. acked_by_gone = next(
  984. (pid for pid in job.worker_pids()
  985. if pid in cleaned or pid not in all_pids),
  986. None
  987. )
  988. # already accepted by process
  989. if acked_by_gone:
  990. self.on_job_process_down(job, acked_by_gone)
  991. if not job.ready():
  992. exitcode = exitcodes.get(acked_by_gone) or 0
  993. proc = cleaned.get(acked_by_gone)
  994. if proc and getattr(proc, '_job_terminated', False):
  995. job._set_terminated(exitcode)
  996. else:
  997. self.on_job_process_lost(
  998. job, acked_by_gone, exitcode,
  999. )
  1000. else:
  1001. # started writing to
  1002. write_to = job._write_to
  1003. # was scheduled to write to
  1004. sched_for = job._scheduled_for
  1005. if write_to and not write_to._is_alive():
  1006. self.on_job_process_down(job, write_to.pid)
  1007. elif sched_for and not sched_for._is_alive():
  1008. self.on_job_process_down(job, sched_for.pid)
  1009. for worker in values(cleaned):
  1010. if self.on_process_down:
  1011. if not shutdown:
  1012. self._process_cleanup_queues(worker)
  1013. self.on_process_down(worker)
  1014. return list(exitcodes.values())
  1015. return []
  1016. def on_partial_read(self, job, worker):
  1017. pass
  1018. def _process_cleanup_queues(self, worker):
  1019. pass
  1020. def on_job_process_down(self, job, pid_gone):
  1021. pass
  1022. def on_job_process_lost(self, job, pid, exitcode):
  1023. job._worker_lost = (monotonic(), exitcode)
  1024. def mark_as_worker_lost(self, job, exitcode):
  1025. try:
  1026. raise WorkerLostError(
  1027. 'Worker exited prematurely: {0}.'.format(
  1028. human_status(exitcode)),
  1029. )
  1030. except WorkerLostError:
  1031. job._set(None, (False, ExceptionInfo()))
  1032. else: # pragma: no cover
  1033. pass
  1034. def __enter__(self):
  1035. return self
  1036. def __exit__(self, *exc_info):
  1037. return self.terminate()
  1038. def on_grow(self, n):
  1039. pass
  1040. def on_shrink(self, n):
  1041. pass
  1042. def shrink(self, n=1):
  1043. for i, worker in enumerate(self._iterinactive()):
  1044. self._processes -= 1
  1045. if self._putlock:
  1046. self._putlock.shrink()
  1047. worker.terminate_controlled()
  1048. self.on_shrink(1)
  1049. if i >= n - 1:
  1050. break
  1051. else:
  1052. raise ValueError("Can't shrink pool. All processes busy!")
  1053. def grow(self, n=1):
  1054. for i in range(n):
  1055. self._processes += 1
  1056. if self._putlock:
  1057. self._putlock.grow()
  1058. self.on_grow(n)
  1059. def _iterinactive(self):
  1060. for worker in self._pool:
  1061. if not self._worker_active(worker):
  1062. yield worker
  1063. def _worker_active(self, worker):
  1064. for job in values(self._cache):
  1065. if worker.pid in job.worker_pids():
  1066. return True
  1067. return False
  1068. def _repopulate_pool(self, exitcodes):
  1069. """Bring the number of pool processes up to the specified number,
  1070. for use after reaping workers which have exited.
  1071. """
  1072. for i in range(self._processes - len(self._pool)):
  1073. if self._state != RUN:
  1074. return
  1075. try:
  1076. if exitcodes and exitcodes[i] not in (EX_OK, EX_RECYCLE):
  1077. self.restart_state.step()
  1078. except IndexError:
  1079. self.restart_state.step()
  1080. self._create_worker_process(self._avail_index())
  1081. debug('added worker')
  1082. def _avail_index(self):
  1083. assert len(self._pool) < self._processes
  1084. indices = set(p.index for p in self._pool)
  1085. return next(i for i in range(self._processes) if i not in indices)
  1086. def did_start_ok(self):
  1087. return not self._join_exited_workers()
  1088. def _maintain_pool(self):
  1089. """"Clean up any exited workers and start replacements for them.
  1090. """
  1091. joined = self._join_exited_workers()
  1092. self._repopulate_pool(joined)
  1093. for i in range(len(joined)):
  1094. if self._putlock is not None:
  1095. self._putlock.release()
  1096. def maintain_pool(self):
  1097. if self._worker_handler._state == RUN and self._state == RUN:
  1098. try:
  1099. self._maintain_pool()
  1100. except RestartFreqExceeded:
  1101. self.close()
  1102. self.join()
  1103. raise
  1104. except OSError as exc:
  1105. if get_errno(exc) == errno.ENOMEM:
  1106. reraise(MemoryError,
  1107. MemoryError(str(exc)),
  1108. sys.exc_info()[2])
  1109. raise
  1110. def _setup_queues(self):
  1111. self._inqueue = self._ctx.SimpleQueue()
  1112. self._outqueue = self._ctx.SimpleQueue()
  1113. self._quick_put = self._inqueue._writer.send
  1114. self._quick_get = self._outqueue._reader.recv
  1115. def _poll_result(timeout):
  1116. if self._outqueue._reader.poll(timeout):
  1117. return True, self._quick_get()
  1118. return False, None
  1119. self._poll_result = _poll_result
  1120. def _start_timeout_handler(self):
  1121. # ensure more than one thread does not start the timeout handler
  1122. # thread at once.
  1123. if self.threads and self._timeout_handler is not None:
  1124. with self._timeout_handler_mutex:
  1125. if not self._timeout_handler_started:
  1126. self._timeout_handler_started = True
  1127. self._timeout_handler.start()
  1128. def apply(self, func, args=(), kwds={}):
  1129. '''
  1130. Equivalent of `func(*args, **kwargs)`.
  1131. '''
  1132. if self._state == RUN:
  1133. return self.apply_async(func, args, kwds).get()
  1134. def starmap(self, func, iterable, chunksize=None):
  1135. '''
  1136. Like `map()` method but the elements of the `iterable` are expected to
  1137. be iterables as well and will be unpacked as arguments. Hence
  1138. `func` and (a, b) becomes func(a, b).
  1139. '''
  1140. if self._state == RUN:
  1141. return self._map_async(func, iterable,
  1142. starmapstar, chunksize).get()
  1143. def starmap_async(self, func, iterable, chunksize=None,
  1144. callback=None, error_callback=None):
  1145. '''
  1146. Asynchronous version of `starmap()` method.
  1147. '''
  1148. if self._state == RUN:
  1149. return self._map_async(func, iterable, starmapstar, chunksize,
  1150. callback, error_callback)
  1151. def map(self, func, iterable, chunksize=None):
  1152. '''
  1153. Apply `func` to each element in `iterable`, collecting the results
  1154. in a list that is returned.
  1155. '''
  1156. if self._state == RUN:
  1157. return self.map_async(func, iterable, chunksize).get()
  1158. def imap(self, func, iterable, chunksize=1, lost_worker_timeout=None):
  1159. '''
  1160. Equivalent of `map()` -- can be MUCH slower than `Pool.map()`.
  1161. '''
  1162. if self._state != RUN:
  1163. return
  1164. lost_worker_timeout = lost_worker_timeout or self.lost_worker_timeout
  1165. if chunksize == 1:
  1166. result = IMapIterator(self._cache,
  1167. lost_worker_timeout=lost_worker_timeout)
  1168. self._taskqueue.put((
  1169. ((TASK, (result._job, i, func, (x,), {}))
  1170. for i, x in enumerate(iterable)),
  1171. result._set_length,
  1172. ))
  1173. return result
  1174. else:
  1175. assert chunksize > 1
  1176. task_batches = Pool._get_tasks(func, iterable, chunksize)
  1177. result = IMapIterator(self._cache,
  1178. lost_worker_timeout=lost_worker_timeout)
  1179. self._taskqueue.put((
  1180. ((TASK, (result._job, i, mapstar, (x,), {}))
  1181. for i, x in enumerate(task_batches)),
  1182. result._set_length,
  1183. ))
  1184. return (item for chunk in result for item in chunk)
  1185. def imap_unordered(self, func, iterable, chunksize=1,
  1186. lost_worker_timeout=None):
  1187. '''
  1188. Like `imap()` method but ordering of results is arbitrary.
  1189. '''
  1190. if self._state != RUN:
  1191. return
  1192. lost_worker_timeout = lost_worker_timeout or self.lost_worker_timeout
  1193. if chunksize == 1:
  1194. result = IMapUnorderedIterator(
  1195. self._cache, lost_worker_timeout=lost_worker_timeout,
  1196. )
  1197. self._taskqueue.put((
  1198. ((TASK, (result._job, i, func, (x,), {}))
  1199. for i, x in enumerate(iterable)),
  1200. result._set_length,
  1201. ))
  1202. return result
  1203. else:
  1204. assert chunksize > 1
  1205. task_batches = Pool._get_tasks(func, iterable, chunksize)
  1206. result = IMapUnorderedIterator(
  1207. self._cache, lost_worker_timeout=lost_worker_timeout,
  1208. )
  1209. self._taskqueue.put((
  1210. ((TASK, (result._job, i, mapstar, (x,), {}))
  1211. for i, x in enumerate(task_batches)),
  1212. result._set_length,
  1213. ))
  1214. return (item for chunk in result for item in chunk)
  1215. def apply_async(self, func, args=(), kwds={},
  1216. callback=None, error_callback=None, accept_callback=None,
  1217. timeout_callback=None, waitforslot=None,
  1218. soft_timeout=None, timeout=None, lost_worker_timeout=None,
  1219. callbacks_propagate=(),
  1220. correlation_id=None):
  1221. '''
  1222. Asynchronous equivalent of `apply()` method.
  1223. Callback is called when the functions return value is ready.
  1224. The accept callback is called when the job is accepted to be executed.
  1225. Simplified the flow is like this:
  1226. >>> def apply_async(func, args, kwds, callback, accept_callback):
  1227. ... if accept_callback:
  1228. ... accept_callback()
  1229. ... retval = func(*args, **kwds)
  1230. ... if callback:
  1231. ... callback(retval)
  1232. '''
  1233. if self._state != RUN:
  1234. return
  1235. soft_timeout = soft_timeout or self.soft_timeout
  1236. timeout = timeout or self.timeout
  1237. lost_worker_timeout = lost_worker_timeout or self.lost_worker_timeout
  1238. if soft_timeout and SIG_SOFT_TIMEOUT is None:
  1239. warnings.warn(UserWarning(
  1240. "Soft timeouts are not supported: "
  1241. "on this platform: It does not have the SIGUSR1 signal.",
  1242. ))
  1243. soft_timeout = None
  1244. if self._state == RUN:
  1245. waitforslot = self.putlocks if waitforslot is None else waitforslot
  1246. if waitforslot and self._putlock is not None:
  1247. self._putlock.acquire()
  1248. result = ApplyResult(
  1249. self._cache, callback, accept_callback, timeout_callback,
  1250. error_callback, soft_timeout, timeout, lost_worker_timeout,
  1251. on_timeout_set=self.on_timeout_set,
  1252. on_timeout_cancel=self.on_timeout_cancel,
  1253. callbacks_propagate=callbacks_propagate,
  1254. send_ack=self.send_ack if self.synack else None,
  1255. correlation_id=correlation_id,
  1256. )
  1257. if timeout or soft_timeout:
  1258. # start the timeout handler thread when required.
  1259. self._start_timeout_handler()
  1260. if self.threads:
  1261. self._taskqueue.put(([(TASK, (result._job, None,
  1262. func, args, kwds))], None))
  1263. else:
  1264. self._quick_put((TASK, (result._job, None, func, args, kwds)))
  1265. return result
  1266. def send_ack(self, response, job, i, fd):
  1267. pass
  1268. def terminate_job(self, pid, sig=None):
  1269. proc, _ = self._process_by_pid(pid)
  1270. if proc is not None:
  1271. try:
  1272. _kill(pid, sig or TERM_SIGNAL)
  1273. except OSError as exc:
  1274. if get_errno(exc) != errno.ESRCH:
  1275. raise
  1276. else:
  1277. proc._controlled_termination = True
  1278. proc._job_terminated = True
  1279. def map_async(self, func, iterable, chunksize=None,
  1280. callback=None, error_callback=None):
  1281. '''
  1282. Asynchronous equivalent of `map()` method.
  1283. '''
  1284. return self._map_async(
  1285. func, iterable, mapstar, chunksize, callback, error_callback,
  1286. )
  1287. def _map_async(self, func, iterable, mapper, chunksize=None,
  1288. callback=None, error_callback=None):
  1289. '''
  1290. Helper function to implement map, starmap and their async counterparts.
  1291. '''
  1292. if self._state != RUN:
  1293. return
  1294. if not hasattr(iterable, '__len__'):
  1295. iterable = list(iterable)
  1296. if chunksize is None:
  1297. chunksize, extra = divmod(len(iterable), len(self._pool) * 4)
  1298. if extra:
  1299. chunksize += 1
  1300. if len(iterable) == 0:
  1301. chunksize = 0
  1302. task_batches = Pool._get_tasks(func, iterable, chunksize)
  1303. result = MapResult(self._cache, chunksize, len(iterable), callback,
  1304. error_callback=error_callback)
  1305. self._taskqueue.put((((TASK, (result._job, i, mapper, (x,), {}))
  1306. for i, x in enumerate(task_batches)), None))
  1307. return result
  1308. @staticmethod
  1309. def _get_tasks(func, it, size):
  1310. it = iter(it)
  1311. while 1:
  1312. x = tuple(itertools.islice(it, size))
  1313. if not x:
  1314. return
  1315. yield (func, x)
  1316. def __reduce__(self):
  1317. raise NotImplementedError(
  1318. 'pool objects cannot be passed between processes or pickled',
  1319. )
  1320. def close(self):
  1321. debug('closing pool')
  1322. if self._state == RUN:
  1323. self._state = CLOSE
  1324. if self._putlock:
  1325. self._putlock.clear()
  1326. self._worker_handler.close()
  1327. self._taskqueue.put(None)
  1328. stop_if_not_current(self._worker_handler)
  1329. def terminate(self):
  1330. debug('terminating pool')
  1331. self._state = TERMINATE
  1332. self._worker_handler.terminate()
  1333. self._terminate()
  1334. @staticmethod
  1335. def _stop_task_handler(task_handler):
  1336. stop_if_not_current(task_handler)
  1337. def join(self):
  1338. assert self._state in (CLOSE, TERMINATE)
  1339. debug('joining worker handler')
  1340. stop_if_not_current(self._worker_handler)
  1341. debug('joining task handler')
  1342. self._stop_task_handler(self._task_handler)
  1343. debug('joining result handler')
  1344. stop_if_not_current(self._result_handler)
  1345. debug('result handler joined')
  1346. for i, p in enumerate(self._pool):
  1347. debug('joining worker %s/%s (%r)', i + 1, len(self._pool), p)
  1348. if p._popen is not None: # process started?
  1349. p.join()
  1350. debug('pool join complete')
  1351. def restart(self):
  1352. for e in values(self._poolctrl):
  1353. e.set()
  1354. @staticmethod
  1355. def _help_stuff_finish(inqueue, task_handler, _pool):
  1356. # task_handler may be blocked trying to put items on inqueue
  1357. debug('removing tasks from inqueue until task handler finished')
  1358. inqueue._rlock.acquire()
  1359. while task_handler.is_alive() and inqueue._reader.poll():
  1360. inqueue._reader.recv()
  1361. time.sleep(0)
  1362. @classmethod
  1363. def _set_result_sentinel(cls, outqueue, pool):
  1364. outqueue.put(None)
  1365. @classmethod
  1366. def _terminate_pool(cls, taskqueue, inqueue, outqueue, pool,
  1367. worker_handler, task_handler,
  1368. result_handler, cache, timeout_handler,
  1369. help_stuff_finish_args):
  1370. # this is guaranteed to only be called once
  1371. debug('finalizing pool')
  1372. worker_handler.terminate()
  1373. task_handler.terminate()
  1374. taskqueue.put(None) # sentinel
  1375. debug('helping task handler/workers to finish')
  1376. cls._help_stuff_finish(*help_stuff_finish_args)
  1377. result_handler.terminate()
  1378. cls._set_result_sentinel(outqueue, pool)
  1379. if timeout_handler is not None:
  1380. timeout_handler.terminate()
  1381. # Terminate workers which haven't already finished
  1382. if pool and hasattr(pool[0], 'terminate'):
  1383. debug('terminating workers')
  1384. for p in pool:
  1385. if p._is_alive():
  1386. p.terminate()
  1387. debug('joining task handler')
  1388. cls._stop_task_handler(task_handler)
  1389. debug('joining result handler')
  1390. result_handler.stop()
  1391. if timeout_handler is not None:
  1392. debug('joining timeout handler')
  1393. timeout_handler.stop(TIMEOUT_MAX)
  1394. if pool and hasattr(pool[0], 'terminate'):
  1395. debug('joining pool workers')
  1396. for p in pool:
  1397. if p.is_alive():
  1398. # worker has not yet exited
  1399. debug('cleaning up worker %d', p.pid)
  1400. if p._popen is not None:
  1401. p.join()
  1402. debug('pool workers joined')
  1403. @property
  1404. def process_sentinels(self):
  1405. return [w._popen.sentinel for w in self._pool]
  1406. #
  1407. # Class whose instances are returned by `Pool.apply_async()`
  1408. #
  1409. class ApplyResult(object):
  1410. _worker_lost = None
  1411. _write_to = None
  1412. _scheduled_for = None
  1413. def __init__(self, cache, callback, accept_callback=None,
  1414. timeout_callback=None, error_callback=None, soft_timeout=None,
  1415. timeout=None, lost_worker_timeout=LOST_WORKER_TIMEOUT,
  1416. on_timeout_set=None, on_timeout_cancel=None,
  1417. callbacks_propagate=(), send_ack=None,
  1418. correlation_id=None):
  1419. self.correlation_id = correlation_id
  1420. self._mutex = Lock()
  1421. self._event = threading.Event()
  1422. self._job = next(job_counter)
  1423. self._cache = cache
  1424. self._callback = callback
  1425. self._accept_callback = accept_callback
  1426. self._error_callback = error_callback
  1427. self._timeout_callback = timeout_callback
  1428. self._timeout = timeout
  1429. self._soft_timeout = soft_timeout
  1430. self._lost_worker_timeout = lost_worker_timeout
  1431. self._on_timeout_set = on_timeout_set
  1432. self._on_timeout_cancel = on_timeout_cancel
  1433. self._callbacks_propagate = callbacks_propagate or ()
  1434. self._send_ack = send_ack
  1435. self._accepted = False
  1436. self._cancelled = False
  1437. self._worker_pid = None
  1438. self._time_accepted = None
  1439. self._terminated = None
  1440. cache[self._job] = self
  1441. def __repr__(self):
  1442. return '<%s: {id} ack:{ack} ready:{ready}>'.format(
  1443. self.__class__.__name__,
  1444. id=self._job, ack=self._accepted, ready=self.ready(),
  1445. )
  1446. def ready(self):
  1447. return self._event.isSet()
  1448. def accepted(self):
  1449. return self._accepted
  1450. def successful(self):
  1451. assert self.ready()
  1452. return self._success
  1453. def _cancel(self):
  1454. """Only works if synack is used."""
  1455. self._cancelled = True
  1456. def discard(self):
  1457. self._cache.pop(self._job, None)
  1458. def terminate(self, signum):
  1459. self._terminated = signum
  1460. def _set_terminated(self, signum=None):
  1461. try:
  1462. raise Terminated(-(signum or 0))
  1463. except Terminated:
  1464. self._set(None, (False, ExceptionInfo()))
  1465. def worker_pids(self):
  1466. return [self._worker_pid] if self._worker_pid else []
  1467. def wait(self, timeout=None):
  1468. self._event.wait(timeout)
  1469. def get(self, timeout=None):
  1470. self.wait(timeout)
  1471. if not self.ready():
  1472. raise TimeoutError
  1473. if self._success:
  1474. return self._value
  1475. else:
  1476. raise self._value.exception
  1477. def safe_apply_callback(self, fun, *args, **kwargs):
  1478. if fun:
  1479. try:
  1480. fun(*args, **kwargs)
  1481. except self._callbacks_propagate:
  1482. raise
  1483. except Exception as exc:
  1484. error('Pool callback raised exception: %r', exc,
  1485. exc_info=1)
  1486. def handle_timeout(self, soft=False):
  1487. if self._timeout_callback is not None:
  1488. self.safe_apply_callback(
  1489. self._timeout_callback, soft=soft,
  1490. timeout=self._soft_timeout if soft else self._timeout,
  1491. )
  1492. def _set(self, i, obj):
  1493. with self._mutex:
  1494. if self._on_timeout_cancel:
  1495. self._on_timeout_cancel(self)
  1496. self._success, self._value = obj
  1497. self._event.set()
  1498. if self._accepted:
  1499. # if not accepted yet, then the set message
  1500. # was received before the ack, which means
  1501. # the ack will remove the entry.
  1502. self._cache.pop(self._job, None)
  1503. # apply callbacks last
  1504. if self._callback and self._success:
  1505. self.safe_apply_callback(
  1506. self._callback, self._value)
  1507. if (self._value is not None and
  1508. self._error_callback and not self._success):
  1509. self.safe_apply_callback(
  1510. self._error_callback, self._value)
  1511. def _ack(self, i, time_accepted, pid, synqW_fd):
  1512. with self._mutex:
  1513. if self._cancelled and self._send_ack:
  1514. self._accepted = True
  1515. if synqW_fd:
  1516. return self._send_ack(NACK, pid, self._job, synqW_fd)
  1517. return
  1518. self._accepted = True
  1519. self._time_accepted = time_accepted
  1520. self._worker_pid = pid
  1521. if self.ready():
  1522. # ack received after set()
  1523. self._cache.pop(self._job, None)
  1524. if self._on_timeout_set:
  1525. self._on_timeout_set(self, self._soft_timeout, self._timeout)
  1526. response = ACK
  1527. if self._accept_callback:
  1528. try:
  1529. self._accept_callback(pid, time_accepted)
  1530. except self._propagate_errors:
  1531. response = NACK
  1532. raise
  1533. except Exception:
  1534. response = NACK
  1535. # ignore other errors
  1536. finally:
  1537. if self._send_ack and synqW_fd:
  1538. return self._send_ack(
  1539. response, pid, self._job, synqW_fd
  1540. )
  1541. if self._send_ack and synqW_fd:
  1542. self._send_ack(response, pid, self._job, synqW_fd)
  1543. #
  1544. # Class whose instances are returned by `Pool.map_async()`
  1545. #
  1546. class MapResult(ApplyResult):
  1547. def __init__(self, cache, chunksize, length, callback, error_callback):
  1548. ApplyResult.__init__(
  1549. self, cache, callback, error_callback=error_callback,
  1550. )
  1551. self._success = True
  1552. self._length = length
  1553. self._value = [None] * length
  1554. self._accepted = [False] * length
  1555. self._worker_pid = [None] * length
  1556. self._time_accepted = [None] * length
  1557. self._chunksize = chunksize
  1558. if chunksize <= 0:
  1559. self._number_left = 0
  1560. self._event.set()
  1561. del cache[self._job]
  1562. else:
  1563. self._number_left = length // chunksize + bool(length % chunksize)
  1564. def _set(self, i, success_result):
  1565. success, result = success_result
  1566. if success:
  1567. self._value[i * self._chunksize:(i + 1) * self._chunksize] = result
  1568. self._number_left -= 1
  1569. if self._number_left == 0:
  1570. if self._callback:
  1571. self._callback(self._value)
  1572. if self._accepted:
  1573. self._cache.pop(self._job, None)
  1574. self._event.set()
  1575. else:
  1576. self._success = False
  1577. self._value = result
  1578. if self._error_callback:
  1579. self._error_callback(self._value)
  1580. if self._accepted:
  1581. self._cache.pop(self._job, None)
  1582. self._event.set()
  1583. def _ack(self, i, time_accepted, pid, *args):
  1584. start = i * self._chunksize
  1585. stop = min((i + 1) * self._chunksize, self._length)
  1586. for j in range(start, stop):
  1587. self._accepted[j] = True
  1588. self._worker_pid[j] = pid
  1589. self._time_accepted[j] = time_accepted
  1590. if self.ready():
  1591. self._cache.pop(self._job, None)
  1592. def accepted(self):
  1593. return all(self._accepted)
  1594. def worker_pids(self):
  1595. return [pid for pid in self._worker_pid if pid]
  1596. #
  1597. # Class whose instances are returned by `Pool.imap()`
  1598. #
  1599. class IMapIterator(object):
  1600. _worker_lost = None
  1601. def __init__(self, cache, lost_worker_timeout=LOST_WORKER_TIMEOUT):
  1602. self._cond = threading.Condition(threading.Lock())
  1603. self._job = next(job_counter)
  1604. self._cache = cache
  1605. self._items = deque()
  1606. self._index = 0
  1607. self._length = None
  1608. self._ready = False
  1609. self._unsorted = {}
  1610. self._worker_pids = []
  1611. self._lost_worker_timeout = lost_worker_timeout
  1612. cache[self._job] = self
  1613. def __iter__(self):
  1614. return self
  1615. def next(self, timeout=None):
  1616. with self._cond:
  1617. try:
  1618. item = self._items.popleft()
  1619. except IndexError:
  1620. if self._index == self._length:
  1621. self._ready = True
  1622. raise StopIteration
  1623. self._cond.wait(timeout)
  1624. try:
  1625. item = self._items.popleft()
  1626. except IndexError:
  1627. if self._index == self._length:
  1628. self._ready = True
  1629. raise StopIteration
  1630. raise TimeoutError
  1631. success, value = item
  1632. if success:
  1633. return value
  1634. raise Exception(value)
  1635. __next__ = next # XXX
  1636. def _set(self, i, obj):
  1637. with self._cond:
  1638. if self._index == i:
  1639. self._items.append(obj)
  1640. self._index += 1
  1641. while self._index in self._unsorted:
  1642. obj = self._unsorted.pop(self._index)
  1643. self._items.append(obj)
  1644. self._index += 1
  1645. self._cond.notify()
  1646. else:
  1647. self._unsorted[i] = obj
  1648. if self._index == self._length:
  1649. self._ready = True
  1650. del self._cache[self._job]
  1651. def _set_length(self, length):
  1652. with self._cond:
  1653. self._length = length
  1654. if self._index == self._length:
  1655. self._ready = True
  1656. self._cond.notify()
  1657. del self._cache[self._job]
  1658. def _ack(self, i, time_accepted, pid, *args):
  1659. self._worker_pids.append(pid)
  1660. def ready(self):
  1661. return self._ready
  1662. def worker_pids(self):
  1663. return self._worker_pids
  1664. #
  1665. # Class whose instances are returned by `Pool.imap_unordered()`
  1666. #
  1667. class IMapUnorderedIterator(IMapIterator):
  1668. def _set(self, i, obj):
  1669. with self._cond:
  1670. self._items.append(obj)
  1671. self._index += 1
  1672. self._cond.notify()
  1673. if self._index == self._length:
  1674. self._ready = True
  1675. del self._cache[self._job]
  1676. #
  1677. #
  1678. #
  1679. class ThreadPool(Pool):
  1680. from .dummy import Process as DummyProcess
  1681. Process = DummyProcess
  1682. def __init__(self, processes=None, initializer=None, initargs=()):
  1683. Pool.__init__(self, processes, initializer, initargs)
  1684. def _setup_queues(self):
  1685. self._inqueue = Queue()
  1686. self._outqueue = Queue()
  1687. self._quick_put = self._inqueue.put
  1688. self._quick_get = self._outqueue.get
  1689. def _poll_result(timeout):
  1690. try:
  1691. return True, self._quick_get(timeout=timeout)
  1692. except Empty:
  1693. return False, None
  1694. self._poll_result = _poll_result
  1695. @staticmethod
  1696. def _help_stuff_finish(inqueue, task_handler, pool):
  1697. # put sentinels at head of inqueue to make workers finish
  1698. with inqueue.not_empty:
  1699. inqueue.queue.clear()
  1700. inqueue.queue.extend([None] * len(pool))
  1701. inqueue.not_empty.notify_all()