More functional fixes to worker

Signed-off-by: Samuel Monson <[email protected]>

Author: sjmonson

src/guidellm/scheduler/worker.py

@@ -30,6 +30,7 @@
 from guidellm.scheduler.result import (
     MPQueues,
     SchedulerRequestInfo,
+    WorkerProcessRequestTime,
     WorkerProcessResult,
 )
 from guidellm.scheduler.types import RequestT, ResponseT
@@ -107,6 +108,11 @@ async def resolve(
         """
         ...
 
+    async def get_request_time(
+        self, times_queue: Queue[WorkerProcessRequestTime]
+    ) -> WorkerProcessRequestTime:
+        return await asyncio.to_thread(times_queue.get)  # type: ignore[attr-defined]
+
     async def send_result(
         self,
         results_queue: Queue[WorkerProcessResult[RequestT, ResponseT]],
@@ -170,7 +176,7 @@ async def resolve_scheduler_request(
         asyncio.create_task(self.send_result(results_queue, result))
 
         request_session.push_response(response)
-        return request
+        return request_session
 
     def process_loop_asynchronous(
         self,
@@ -194,7 +200,7 @@ async def _process_runner():
                         else queues.requests.get_nowait()
                     )
                     dequeued_time = time.time()
-                    request_times = queues.times.get()
+                    request_times = await self.get_request_time(queues.times)
                 except (QueueEmpty, IndexError):
                     # Requeue the session if we don't have a next time yet
                     if request_session is not None:
@@ -215,7 +221,7 @@ async def wait_then_requeue(
                         # Release the lock with the session on top of the stack
                         lock.release()
 
-                async def _request_callback(
+                def _request_callback(
                     session_future: asyncio.Future[RequestSession[RequestT, ResponseT]],
                 ):
                     # If we are prioritizing sessions, hold
@@ -224,7 +230,7 @@ async def _request_callback(
                     if not prioritize_sessions:
                         lock.release()
 
-                    session = await session_future
+                    session = session_future.result()
                     if not session.complete:
                         asyncio.create_task(wait_then_requeue(session))
                     elif prioritize_sessions: