Avoid circular import in WorkerProcessResult

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

Author: sjmonson

src/guidellm/scheduler/result.py

@@ -1,18 +1,24 @@
+from dataclasses import dataclass
+from multiprocessing import Queue
 from typing import (
     Generic,
     Literal,
     Optional,
 )
 
 from guidellm.objects import StandardBaseModel
+from guidellm.request.session import RequestSession
 from guidellm.scheduler.strategy import SchedulingStrategy
 from guidellm.scheduler.types import RequestT, ResponseT
 
 __all__ = [
+    "MPQueues",
     "SchedulerRequestInfo",
     "SchedulerRequestResult",
     "SchedulerResult",
     "SchedulerRunInfo",
+    "WorkerProcessRequestTime",
+    "WorkerProcessResult",
 ]
 
 
@@ -135,3 +141,28 @@ class SchedulerRequestResult(
     request: RequestT
     request_info: SchedulerRequestInfo
     response: Optional[ResponseT] = None
+
+
+# TODO: Move dataclasses somewhere else
+
+
+@dataclass
+class WorkerProcessRequestTime:
+    start_time: float
+    timeout_time: float
+    queued_time: float
+
+
+@dataclass
+class WorkerProcessResult(Generic[RequestT, ResponseT]):
+    type_: Literal["request_scheduled", "request_start", "request_complete"]
+    request: RequestT
+    response: Optional[ResponseT]
+    info: SchedulerRequestInfo
+
+
+@dataclass
+class MPQueues(Generic[RequestT, ResponseT]):
+    requests: Queue[RequestSession[RequestT, ResponseT]]
+    times: Queue[WorkerProcessRequestTime]
+    responses: Queue[WorkerProcessResult[RequestT, ResponseT]]

src/guidellm/scheduler/scheduler.py

@@ -17,17 +17,17 @@
 from guidellm.config import settings
 from guidellm.request.session import RequestSession
 from guidellm.scheduler.result import (
+    MPQueues,
     SchedulerRequestResult,
     SchedulerResult,
     SchedulerRunInfo,
+    WorkerProcessRequestTime,
+    WorkerProcessResult,
 )
 from guidellm.scheduler.strategy import SchedulingStrategy
 from guidellm.scheduler.types import (
-    MPQueues,
     RequestT,
     ResponseT,
-    WorkerProcessRequestTime,
-    WorkerProcessResult,
 )
 from guidellm.scheduler.worker import (
     RequestsWorker,

src/guidellm/scheduler/types.py

@@ -1,43 +1,10 @@
-from dataclasses import dataclass
-from multiprocessing import Queue
-from typing import Generic, Literal, Optional, TypeVar
-
-from guidellm.request.session import RequestSession
-from guidellm.scheduler.result import SchedulerRequestInfo
+from typing import TypeVar
 
 __all__ = [
-    "MPQueues",
     "RequestT",
     "ResponseT",
-    "WorkerProcessRequestTime",
-    "WorkerProcessResult",
 ]
 
 
 RequestT = TypeVar("RequestT")
 ResponseT = TypeVar("ResponseT")
-
-
-# TODO: Move dataclasses somewhere else
-
-
-@dataclass
-class WorkerProcessRequestTime:
-    start_time: float
-    timeout_time: float
-    queued_time: float
-
-
-@dataclass
-class WorkerProcessResult(Generic[RequestT, ResponseT]):
-    type_: Literal["request_scheduled", "request_start", "request_complete"]
-    request: RequestT
-    response: Optional[ResponseT]
-    info: SchedulerRequestInfo
-
-
-@dataclass
-class MPQueues(Generic[RequestT, ResponseT]):
-    requests: Queue[RequestSession[RequestT, ResponseT]]
-    times: Queue[WorkerProcessRequestTime]
-    responses: Queue[WorkerProcessResult[RequestT, ResponseT]]

src/guidellm/scheduler/worker.py

@@ -27,8 +27,12 @@
 from guidellm.objects import StandardBaseModel
 from guidellm.request import GenerationRequest
 from guidellm.request.session import RequestSession
-from guidellm.scheduler.result import SchedulerRequestInfo
-from guidellm.scheduler.types import MPQueues, RequestT, ResponseT, WorkerProcessResult
+from guidellm.scheduler.result import (
+    MPQueues,
+    SchedulerRequestInfo,
+    WorkerProcessResult,
+)
+from guidellm.scheduler.types import RequestT, ResponseT
 
 __all__ = [
     "GenerativeRequestsWorker",