Skip to content

Commit

Permalink
Limit coroutines using a pool instead of chunks (#1544)
Browse files Browse the repository at this point in the history

---------

Co-authored-by: Martin Durant <martin.durant@alumni.utoronto.ca>
  • Loading branch information
sisp and martindurant committed Mar 15, 2024
1 parent 29c06ea commit f2f4c26
Showing 1 changed file with 29 additions and 14 deletions.
43 changes: 29 additions & 14 deletions fsspec/asyn.py
Original file line number Diff line number Diff line change
Expand Up @@ -239,20 +239,35 @@ async def _run_coros_in_chunks(
batch_size = len(coros)

assert batch_size > 0
results = []
for start in range(0, len(coros), batch_size):
chunk = [
asyncio.Task(asyncio.wait_for(c, timeout=timeout))
for c in coros[start : start + batch_size]
]
if callback is not DEFAULT_CALLBACK:
[
t.add_done_callback(lambda *_, **__: callback.relative_update(1))
for t in chunk
]
results.extend(
await asyncio.gather(*chunk, return_exceptions=return_exceptions),
)

async def _run_coro(coro, i):
try:
return await asyncio.wait_for(coro, timeout=timeout), i
except Exception as e:
if not return_exceptions:
raise
return e, i
finally:
callback.relative_update(1)

i = 0
n = len(coros)
results = [None] * n
pending = set()

while pending or i < n:
while len(pending) < batch_size and i < n:
pending.add(asyncio.ensure_future(_run_coro(coros[i], i)))
i += 1

if not pending:
break

done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
while done:
result, k = await done.pop()
results[k] = result

return results


Expand Down

0 comments on commit f2f4c26

Please sign in to comment.