Skip to content

Commit 7f611a5

Browse files
authored
Replace busy wait in async loop with a Semaphore (#80)
Use a semaphore to concurrency limit the scheduling of new tasks rather than a busy wait. Additionally reduce async sleep at the end of the loop iteration as that is sufficient to cause preemption. As a side effect #70 is fixed by this PR since the offending code has been reworked. Closes #70
1 parent ecf2984 commit 7f611a5

File tree

1 file changed

+7
-6
lines changed

1 file changed

+7
-6
lines changed

src/guidellm/scheduler/base.py

Lines changed: 7 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,9 @@ async def _run_async(
317317
)
318318

319319
def _completed(_task: asyncio.Task) -> None:
320-
nonlocal completed
321-
completed += 1
320+
# NOTE: this is only ok because we don't use threads/processes
321+
nonlocal pending
322+
pending.release()
322323
_res = _task.result()
323324

324325
if _res:
@@ -333,7 +334,7 @@ def _completed(_task: asyncio.Task) -> None:
333334
tasks.append(task)
334335

335336
# release control to the event loop for other tasks
336-
await asyncio.sleep(0.001)
337+
await asyncio.sleep(0)
337338

338339
for compl_task in asyncio.as_completed(tasks):
339340
task_res = await compl_task

0 commit comments

Comments
 (0)