diff --git a/benchmark/hicache/bench_mix.py b/benchmark/hicache/bench_mix.py index 833dbf780add..2a65574ea882 100644 --- a/benchmark/hicache/bench_mix.py +++ b/benchmark/hicache/bench_mix.py @@ -426,11 +426,13 @@ async def handle_request(self, user_data): def request_sender(self): async def request_loop(): + tasks = [] while True: if self.sent_requests - self.completed_requests < self.max_parallel: new_request = self.user_generator.pop() if new_request: - asyncio.create_task(self.handle_request(new_request)) + task = asyncio.create_task(self.handle_request(new_request)) + tasks.append(task) self.sent_requests += 1 else: await asyncio.sleep(0.05) @@ -440,6 +442,11 @@ async def request_loop(): self.done = True break + # Cancel all pending tasks and wait for them to finish + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) loop.run_until_complete(request_loop())