Skip to content
This repository has been archived by the owner on Jan 23, 2024. It is now read-only.

Commit

Permalink
Fix argument preparation for Pool.starmap
Browse files Browse the repository at this point in the history
Both Pool.startmap and Pool.apply_async set `starmap=True` when creating
tasks. Previously Pool.starmap didn't have correct argument list. The
correct format should be:

    [(args, kwargs), ...] (used by Pool.apply_async) or
    [(args,), ...] (used by Pool.starmap)
  • Loading branch information
calio committed Aug 21, 2020
1 parent 3c8d7e1 commit 62c085c
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 1 deletion.
4 changes: 3 additions & 1 deletion fiber/pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -1293,9 +1293,11 @@ def starmap_async(self, func, iterable, chunksize=None, callback=None,

seq = self._inventory.add(len(iterable))

iterable = [(item,) for item in iterable]

chunks = self.__class__._chunks(iterable, chunksize)
for batch, chunk in enumerate(chunks):
self._task_put((seq, batch * chunksize, func, (chunk, ), True))
self._task_put((seq, batch * chunksize, func, chunk, True))

res = MapResult(seq, self._inventory)

Expand Down
14 changes: 14 additions & 0 deletions tests/test_pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@ def client():
def f(x):
return x * x

def f2(x, y):
return x * y

def fy(x, y=1):
return x * x * y
Expand Down Expand Up @@ -143,6 +145,18 @@ def test_pool_starmap(self):
pool.terminate()
pool.join()

def test_pool_starmap2(self):
pool = Pool(4)
res = pool.starmap(f2, [(x, x) for x in range(100)], 10)
assert res == [x * x for x in range(100)]

async_res = pool.starmap_async(f, [(x,) for x in range(100)], 10)
res = async_res.get()
assert res == [x * x for x in range(100)]

pool.terminate()
pool.join()

def test_pool_multiple_workers_inside_one_job(self):
old_val = fiber_config.cpu_per_job
try:
Expand Down

0 comments on commit 62c085c

Please sign in to comment.