File tree Expand file tree Collapse file tree 1 file changed +6
-6
lines changed Expand file tree Collapse file tree 1 file changed +6
-6
lines changed Original file line number Diff line number Diff line change @@ -299,13 +299,13 @@ async def _run_async(
299
299
self , benchmark : TextGenerationBenchmark , end_time : float , max_number : float
300
300
) -> AsyncGenerator [Union [TextGenerationResult , TextGenerationError ], None ]:
301
301
tasks = []
302
- completed = 0
302
+ pending = asyncio . Semaphore ( settings . max_concurrency )
303
303
304
304
for index , (request , submit_at ) in enumerate (
305
305
zip (self .generator , self .load_generator .times ())
306
306
):
307
- while ( index + 1 - completed ) >= settings . max_concurrency :
308
- await asyncio . sleep ( 0.1 )
307
+ # wait for number of pending tasks to be >= max_concurrency
308
+ await pending . acquire ( )
309
309
310
310
if index >= max_number or time .time () >= end_time or submit_at >= end_time :
311
311
break
@@ -317,8 +317,8 @@ async def _run_async(
317
317
)
318
318
319
319
def _completed (_task : asyncio .Task ) -> None :
320
- nonlocal completed
321
- completed += 1
320
+ nonlocal pending # NOTE: this is only ok because we don't use threads/processes
321
+ pending . release ()
322
322
_res = _task .result ()
323
323
324
324
if _res :
@@ -333,7 +333,7 @@ def _completed(_task: asyncio.Task) -> None:
333
333
tasks .append (task )
334
334
335
335
# release control to the event loop for other tasks
336
- await asyncio .sleep (0.001 )
336
+ await asyncio .sleep (0 )
337
337
338
338
for compl_task in asyncio .as_completed (tasks ):
339
339
task_res = await compl_task
You can’t perform that action at this time.
0 commit comments