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(
299299 self , benchmark : TextGenerationBenchmark , end_time : float , max_number : float
300300 ) -> AsyncGenerator [Union [TextGenerationResult , TextGenerationError ], None ]:
301301 tasks = []
302- completed = 0
302+ pending = asyncio . Semaphore ( settings . max_concurrency )
303303
304304 for index , (request , submit_at ) in enumerate (
305305 zip (self .generator , self .load_generator .times ())
306306 ):
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 ( )
309309
310310 if index >= max_number or time .time () >= end_time or submit_at >= end_time :
311311 break
@@ -317,8 +317,8 @@ async def _run_async(
317317 )
318318
319319 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 ()
322322 _res = _task .result ()
323323
324324 if _res :
@@ -333,7 +333,7 @@ def _completed(_task: asyncio.Task) -> None:
333333 tasks .append (task )
334334
335335 # release control to the event loop for other tasks
336- await asyncio .sleep (0.001 )
336+ await asyncio .sleep (0 )
337337
338338 for compl_task in asyncio .as_completed (tasks ):
339339 task_res = await compl_task
You can’t perform that action at this time.
0 commit comments