mp_pool.py

#
# A test of `multiprocessing.Pool` class
#
# Copyright (c) 2006-2008, R Oudkerk
# All rights reserved.
#

importmultiprocessing
importtime
importrandom
importsys

#
# Functions used by test code
#

defcalculate(func, args):
result=func(*args)
return'%s says that %s%s = %s'% (
multiprocessing.current_process().name,
func.__name__, args, result
 )

defcalculatestar(args):
returncalculate(*args)

defmul(a, b):
time.sleep(0.5*random.random())
returna*b

defplus(a, b):
time.sleep(0.5*random.random())
returna+b

deff(x):
return1.0/ (x-5.0)

defpow3(x):
returnx**3

defnoop(x):
pass

#
# Test code
#

deftest():
print'cpu_count() = %d\n'%multiprocessing.cpu_count()

#
# Create pool
#

PROCESSES=4
print'Creating pool with %d processes\n'%PROCESSES
pool=multiprocessing.Pool(PROCESSES)
print'pool = %s'%pool
print

#
# Tests
#

TASKS= [(mul, (i, 7)) foriinrange(10)] + \
 [(plus, (i, 8)) foriinrange(10)]

results= [pool.apply_async(calculate, t) fortinTASKS]
imap_it=pool.imap(calculatestar, TASKS)
imap_unordered_it=pool.imap_unordered(calculatestar, TASKS)

print'Ordered results using pool.apply_async():'
forrinresults:
print'\t', r.get()
print

print'Ordered results using pool.imap():'
forxinimap_it:
print'\t', x
print

print'Unordered results using pool.imap_unordered():'
forxinimap_unordered_it:
print'\t', x
print

print'Ordered results using pool.map() --- will block till complete:'
forxinpool.map(calculatestar, TASKS):
print'\t', x
print

#
# Simple benchmarks
#

N=100000
print'def pow3(x): return x**3'

t=time.time()
A=map(pow3, xrange(N))
print'\tmap(pow3, xrange(%d)):\n\t\t%s seconds'% \
 (N, time.time() -t)

t=time.time()
B=pool.map(pow3, xrange(N))
print'\tpool.map(pow3, xrange(%d)):\n\t\t%s seconds'% \
 (N, time.time() -t)

t=time.time()
C=list(pool.imap(pow3, xrange(N), chunksize=N//8))
print'\tlist(pool.imap(pow3, xrange(%d), chunksize=%d)):\n\t\t%s' \
' seconds'% (N, N//8, time.time() -t)

assertA==B==C, (len(A), len(B), len(C))
print

L= [None] *1000000
print'def noop(x): pass'
print'L = [None] * 1000000'

t=time.time()
A=map(noop, L)
print'\tmap(noop, L):\n\t\t%s seconds'% \
 (time.time() -t)

t=time.time()
B=pool.map(noop, L)
print'\tpool.map(noop, L):\n\t\t%s seconds'% \
 (time.time() -t)

t=time.time()
C=list(pool.imap(noop, L, chunksize=len(L)//8))
print'\tlist(pool.imap(noop, L, chunksize=%d)):\n\t\t%s seconds'% \
 (len(L)//8, time.time() -t)

assertA==B==C, (len(A), len(B), len(C))
print

delA, B, C, L

#
# Test error handling
#

print'Testing error handling:'

try:
printpool.apply(f, (5,))
exceptZeroDivisionError:
print'\tGot ZeroDivisionError as expected from pool.apply()'
else:
raiseAssertionError('expected ZeroDivisionError')

try:
printpool.map(f, range(10))
exceptZeroDivisionError:
print'\tGot ZeroDivisionError as expected from pool.map()'
else:
raiseAssertionError('expected ZeroDivisionError')

try:
printlist(pool.imap(f, range(10)))
exceptZeroDivisionError:
print'\tGot ZeroDivisionError as expected from list(pool.imap())'
else:
raiseAssertionError('expected ZeroDivisionError')

it=pool.imap(f, range(10))
foriinrange(10):
try:
x=it.next()
exceptZeroDivisionError:
ifi==5:
pass
exceptStopIteration:
break
else:
ifi==5:
raiseAssertionError('expected ZeroDivisionError')

asserti==9
print'\tGot ZeroDivisionError as expected from IMapIterator.next()'
print

#
# Testing timeouts
#

print'Testing ApplyResult.get() with timeout:',
res=pool.apply_async(calculate, TASKS[0])
while1:
sys.stdout.flush()
try:
sys.stdout.write('\n\t%s'%res.get(0.02))
break
exceptmultiprocessing.TimeoutError:
sys.stdout.write('.')
print
print

print'Testing IMapIterator.next() with timeout:',
it=pool.imap(calculatestar, TASKS)
while1:
sys.stdout.flush()
try:
sys.stdout.write('\n\t%s'%it.next(0.02))
exceptStopIteration:
break
exceptmultiprocessing.TimeoutError:
sys.stdout.write('.')
print
print

#
# Testing callback
#

print'Testing callback:'

A= []
B= [56, 0, 1, 8, 27, 64, 125, 216, 343, 512, 729]

r=pool.apply_async(mul, (7, 8), callback=A.append)
r.wait()

r=pool.map_async(pow3, range(10), callback=A.extend)
r.wait()

ifA==B:
print'\tcallbacks succeeded\n'
else:
print'\t*** callbacks failed\n\t\t%s != %s\n'% (A, B)

#
# Check there are no outstanding tasks
#

assertnotpool._cache, 'cache = %r'%pool._cache

#
# Check close() methods
#

print'Testing close():'

forworkerinpool._pool:
assertworker.is_alive()

result=pool.apply_async(time.sleep, [0.5])
pool.close()
pool.join()

assertresult.get() isNone

forworkerinpool._pool:
assertnotworker.is_alive()

print'\tclose() succeeded\n'

#
# Check terminate() method
#

print'Testing terminate():'

pool=multiprocessing.Pool(2)
DELTA=0.1
ignore=pool.apply(pow3, [2])
results= [pool.apply_async(time.sleep, [DELTA]) foriinrange(100)]
pool.terminate()
pool.join()

forworkerinpool._pool:
assertnotworker.is_alive()

print'\tterminate() succeeded\n'

#
# Check garbage collection
#

print'Testing garbage collection:'

pool=multiprocessing.Pool(2)
DELTA=0.1
processes=pool._pool
ignore=pool.apply(pow3, [2])
results= [pool.apply_async(time.sleep, [DELTA]) foriinrange(100)]

results=pool=None

time.sleep(DELTA*2)

forworkerinprocesses:
assertnotworker.is_alive()

print'\tgarbage collection succeeded\n'


if__name__=='__main__':
multiprocessing.freeze_support()

assertlen(sys.argv) in (1, 2)

iflen(sys.argv) ==1orsys.argv[1] =='processes':
print' Using processes '.center(79, '-')
elifsys.argv[1] =='threads':
print' Using threads '.center(79, '-')
importmultiprocessing.dummyasmultiprocessing
else:
print'Usage:\n\t%s [processes | threads]'%sys.argv[0]
raiseSystemExit(2)

test()