WIP

Author: Dmytro Parfeniuk

src/domain/backend/base.py

@@ -93,7 +93,6 @@ def submit(self, request: TextGenerationRequest) -> TextGenerationResult:
         result = TextGenerationResult(request=request)
         result.start(request.prompt)
 
-        breakpoint()  # TODO: remove
         for response in self.make_request(request):
             if response.type_ == "token_iter" and response.add_token:
                 result.output_token(response.add_token)

src/domain/core/result.py

@@ -40,7 +40,6 @@ def __init__(self, request: TextGenerationRequest):
         self._first_token_time: Optional[float] = None
         self._decode_times = Distribution()
 
-        breakpoint()  # TODO: remove
         logger.debug(f"Initialized TextGenerationResult for request: {self._request}")
 
     def __repr__(self) -> str:

src/domain/scheduler/scheduler.py

@@ -72,26 +72,17 @@ async def _run_async(self) -> TextGenerationBenchmark:
         )
         load_gen = LoadGenerator(self._load_gen_mode, self._load_gen_rate)
 
-        coroutines = []
+        tasks = []
         start_time = time.time()
         counter = 0
-
         try:
-            for text_generation_request, task_start_time in zip(
-                self._request_generator, load_gen.times()
-            ):
-                coro = Task(
-                    func=self._backend.submit,
-                    params={"request": text_generation_request.prompt},
-                    err_container=TextGenerationError,
-                )
-
+            for task, task_start_time in zip(self._task_iterator(), load_gen.times()):
                 pending_time = task_start_time - time.time()
 
                 if pending_time > 0:
                     await asyncio.sleep(pending_time)
 
-                coroutines.append(self._run_task_async(coro, result_set))
+                tasks.append(self._run_task_async(task, result_set))
                 counter += 1
 
                 if (
@@ -108,17 +99,24 @@ async def _run_async(self) -> TextGenerationBenchmark:
                     await asyncio.sleep(pending_duration)
                 raise asyncio.CancelledError()
 
-            await asyncio.gather(*coroutines)
-
+            await asyncio.gather(*tasks)
         except asyncio.CancelledError:
             # Cancel all pending tasks
-            for coro in coroutines:
-                if not coro.done():
-                    coro.cancel()
+            for task in tasks:
+                if not task.done():
+                    task.cancel()
 
         return result_set
 
     async def _run_task_async(self, task: Task, result_set: TextGenerationBenchmark):
         result_set.request_started()
         res = await task.run_async()
         result_set.request_completed(res)
+
+    def _task_iterator(self) -> Iterable[Task]:
+        for request in self._request_generator:
+            yield Task(
+                func=self._backend.submit,
+                params={"request": request},
+                err_container=TextGenerationError,
+            )