/Doc/includes/mp_pool.py

http://unladen-swallow.googlecode.com/ · Python · 314 lines · 198 code · 74 blank · 42 comment · 43 complexity · 3a224a1590f0f60b5553019ae147db13 MD5 · raw file

  1. #
  2. # A test of `multiprocessing.Pool` class
  3. #
  4. # Copyright (c) 2006-2008, R Oudkerk
  5. # All rights reserved.
  6. #
  7. import multiprocessing
  8. import time
  9. import random
  10. import sys
  11. #
  12. # Functions used by test code
  13. #
  14. def calculate(func, args):
  15. result = func(*args)
  16. return '%s says that %s%s = %s' % (
  17. multiprocessing.current_process().name,
  18. func.__name__, args, result
  19. )
  20. def calculatestar(args):
  21. return calculate(*args)
  22. def mul(a, b):
  23. time.sleep(0.5*random.random())
  24. return a * b
  25. def plus(a, b):
  26. time.sleep(0.5*random.random())
  27. return a + b
  28. def f(x):
  29. return 1.0 / (x-5.0)
  30. def pow3(x):
  31. return x**3
  32. def noop(x):
  33. pass
  34. #
  35. # Test code
  36. #
  37. def test():
  38. print 'cpu_count() = %d\n' % multiprocessing.cpu_count()
  39. #
  40. # Create pool
  41. #
  42. PROCESSES = 4
  43. print 'Creating pool with %d processes\n' % PROCESSES
  44. pool = multiprocessing.Pool(PROCESSES)
  45. print 'pool = %s' % pool
  46. print
  47. #
  48. # Tests
  49. #
  50. TASKS = [(mul, (i, 7)) for i in range(10)] + \
  51. [(plus, (i, 8)) for i in range(10)]
  52. results = [pool.apply_async(calculate, t) for t in TASKS]
  53. imap_it = pool.imap(calculatestar, TASKS)
  54. imap_unordered_it = pool.imap_unordered(calculatestar, TASKS)
  55. print 'Ordered results using pool.apply_async():'
  56. for r in results:
  57. print '\t', r.get()
  58. print
  59. print 'Ordered results using pool.imap():'
  60. for x in imap_it:
  61. print '\t', x
  62. print
  63. print 'Unordered results using pool.imap_unordered():'
  64. for x in imap_unordered_it:
  65. print '\t', x
  66. print
  67. print 'Ordered results using pool.map() --- will block till complete:'
  68. for x in pool.map(calculatestar, TASKS):
  69. print '\t', x
  70. print
  71. #
  72. # Simple benchmarks
  73. #
  74. N = 100000
  75. print 'def pow3(x): return x**3'
  76. t = time.time()
  77. A = map(pow3, xrange(N))
  78. print '\tmap(pow3, xrange(%d)):\n\t\t%s seconds' % \
  79. (N, time.time() - t)
  80. t = time.time()
  81. B = pool.map(pow3, xrange(N))
  82. print '\tpool.map(pow3, xrange(%d)):\n\t\t%s seconds' % \
  83. (N, time.time() - t)
  84. t = time.time()
  85. C = list(pool.imap(pow3, xrange(N), chunksize=N//8))
  86. print '\tlist(pool.imap(pow3, xrange(%d), chunksize=%d)):\n\t\t%s' \
  87. ' seconds' % (N, N//8, time.time() - t)
  88. assert A == B == C, (len(A), len(B), len(C))
  89. print
  90. L = [None] * 1000000
  91. print 'def noop(x): pass'
  92. print 'L = [None] * 1000000'
  93. t = time.time()
  94. A = map(noop, L)
  95. print '\tmap(noop, L):\n\t\t%s seconds' % \
  96. (time.time() - t)
  97. t = time.time()
  98. B = pool.map(noop, L)
  99. print '\tpool.map(noop, L):\n\t\t%s seconds' % \
  100. (time.time() - t)
  101. t = time.time()
  102. C = list(pool.imap(noop, L, chunksize=len(L)//8))
  103. print '\tlist(pool.imap(noop, L, chunksize=%d)):\n\t\t%s seconds' % \
  104. (len(L)//8, time.time() - t)
  105. assert A == B == C, (len(A), len(B), len(C))
  106. print
  107. del A, B, C, L
  108. #
  109. # Test error handling
  110. #
  111. print 'Testing error handling:'
  112. try:
  113. print pool.apply(f, (5,))
  114. except ZeroDivisionError:
  115. print '\tGot ZeroDivisionError as expected from pool.apply()'
  116. else:
  117. raise AssertionError, 'expected ZeroDivisionError'
  118. try:
  119. print pool.map(f, range(10))
  120. except ZeroDivisionError:
  121. print '\tGot ZeroDivisionError as expected from pool.map()'
  122. else:
  123. raise AssertionError, 'expected ZeroDivisionError'
  124. try:
  125. print list(pool.imap(f, range(10)))
  126. except ZeroDivisionError:
  127. print '\tGot ZeroDivisionError as expected from list(pool.imap())'
  128. else:
  129. raise AssertionError, 'expected ZeroDivisionError'
  130. it = pool.imap(f, range(10))
  131. for i in range(10):
  132. try:
  133. x = it.next()
  134. except ZeroDivisionError:
  135. if i == 5:
  136. pass
  137. except StopIteration:
  138. break
  139. else:
  140. if i == 5:
  141. raise AssertionError, 'expected ZeroDivisionError'
  142. assert i == 9
  143. print '\tGot ZeroDivisionError as expected from IMapIterator.next()'
  144. print
  145. #
  146. # Testing timeouts
  147. #
  148. print 'Testing ApplyResult.get() with timeout:',
  149. res = pool.apply_async(calculate, TASKS[0])
  150. while 1:
  151. sys.stdout.flush()
  152. try:
  153. sys.stdout.write('\n\t%s' % res.get(0.02))
  154. break
  155. except multiprocessing.TimeoutError:
  156. sys.stdout.write('.')
  157. print
  158. print
  159. print 'Testing IMapIterator.next() with timeout:',
  160. it = pool.imap(calculatestar, TASKS)
  161. while 1:
  162. sys.stdout.flush()
  163. try:
  164. sys.stdout.write('\n\t%s' % it.next(0.02))
  165. except StopIteration:
  166. break
  167. except multiprocessing.TimeoutError:
  168. sys.stdout.write('.')
  169. print
  170. print
  171. #
  172. # Testing callback
  173. #
  174. print 'Testing callback:'
  175. A = []
  176. B = [56, 0, 1, 8, 27, 64, 125, 216, 343, 512, 729]
  177. r = pool.apply_async(mul, (7, 8), callback=A.append)
  178. r.wait()
  179. r = pool.map_async(pow3, range(10), callback=A.extend)
  180. r.wait()
  181. if A == B:
  182. print '\tcallbacks succeeded\n'
  183. else:
  184. print '\t*** callbacks failed\n\t\t%s != %s\n' % (A, B)
  185. #
  186. # Check there are no outstanding tasks
  187. #
  188. assert not pool._cache, 'cache = %r' % pool._cache
  189. #
  190. # Check close() methods
  191. #
  192. print 'Testing close():'
  193. for worker in pool._pool:
  194. assert worker.is_alive()
  195. result = pool.apply_async(time.sleep, [0.5])
  196. pool.close()
  197. pool.join()
  198. assert result.get() is None
  199. for worker in pool._pool:
  200. assert not worker.is_alive()
  201. print '\tclose() succeeded\n'
  202. #
  203. # Check terminate() method
  204. #
  205. print 'Testing terminate():'
  206. pool = multiprocessing.Pool(2)
  207. DELTA = 0.1
  208. ignore = pool.apply(pow3, [2])
  209. results = [pool.apply_async(time.sleep, [DELTA]) for i in range(100)]
  210. pool.terminate()
  211. pool.join()
  212. for worker in pool._pool:
  213. assert not worker.is_alive()
  214. print '\tterminate() succeeded\n'
  215. #
  216. # Check garbage collection
  217. #
  218. print 'Testing garbage collection:'
  219. pool = multiprocessing.Pool(2)
  220. DELTA = 0.1
  221. processes = pool._pool
  222. ignore = pool.apply(pow3, [2])
  223. results = [pool.apply_async(time.sleep, [DELTA]) for i in range(100)]
  224. results = pool = None
  225. time.sleep(DELTA * 2)
  226. for worker in processes:
  227. assert not worker.is_alive()
  228. print '\tgarbage collection succeeded\n'
  229. if __name__ == '__main__':
  230. multiprocessing.freeze_support()
  231. assert len(sys.argv) in (1, 2)
  232. if len(sys.argv) == 1 or sys.argv[1] == 'processes':
  233. print ' Using processes '.center(79, '-')
  234. elif sys.argv[1] == 'threads':
  235. print ' Using threads '.center(79, '-')
  236. import multiprocessing.dummy as multiprocessing
  237. else:
  238. print 'Usage:\n\t%s [processes | threads]' % sys.argv[0]
  239. raise SystemExit(2)
  240. test()