Skip to content

Commit 61eaec7

Browse files
committed
Replace completed count with a semaphore
Reduce sleep at end of run loop to 0
1 parent ecf2984 commit 61eaec7

File tree

1 file changed

+6
-6
lines changed

1 file changed

+6
-6
lines changed

src/guidellm/scheduler/base.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff 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

0 commit comments

Comments
 (0)