guidellm 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

guidellm/scheduler/worker.py

@@ -0,0 +1,511 @@
+import asyncio
+import math
+import multiprocessing
+import multiprocessing.queues
+import time
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
+from dataclasses import dataclass
+from typing import (
+    Any,
+    Generic,
+    Literal,
+    Optional,
+    Union,
+)
+
+from loguru import logger
+from pydantic import Field
+
+from guidellm.backend import (
+    Backend,
+    BackendType,
+    RequestArgs,
+    ResponseSummary,
+    StreamingTextResponse,
+)
+from guidellm.objects import StandardBaseModel
+from guidellm.request import GenerationRequest
+from guidellm.scheduler.result import SchedulerRequestInfo
+from guidellm.scheduler.types import RequestT, ResponseT
+
+__all__ = [
+    "WorkerProcessRequest",
+    "WorkerProcessResult",
+    "ResolveStatus",
+    "WorkerDescription",
+    "RequestsWorker",
+    "GenerativeRequestsWorkerDescription",
+    "GenerativeRequestsWorker",
+]
+
+
+@dataclass
+class WorkerProcessRequest(Generic[RequestT]):
+    request: RequestT
+    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 ResolveStatus:
+    requested: bool
+    completed: bool
+    errored: bool
+    canceled: bool
+
+    request_start: float
+    request_end: float
+
+
+class WorkerDescription(StandardBaseModel):
+    type_: Literal["worker"] = "worker"
+
+
+class RequestsWorker(ABC, Generic[RequestT, ResponseT]):
+    """
+    An abstract base class for a worker that processes requests.
+    This class defines the interface for a worker that can resolve requests
+    asynchronously or synchronously within the Scheduler class.
+    Subclasses must implement the `resolve` method,
+    which takes a request directly given from the load generator,
+    along with the desired start_time for the request and a timeout_time.
+    The `resolve` method should return the response from the backend.
+    """
+
+    @property
+    @abstractmethod
+    def description(self) -> WorkerDescription:
+        """
+        An abstract property that must be implemented by subclasses.
+        This property should return a Serializable class representing the information
+        about the worker instance.
+        """
+        ...
+
+    @abstractmethod
+    async def prepare_multiprocessing(self):
+        """
+        An abstract method that must be implemented by subclasses.
+        This is useful for workers that have instance state that can not
+        be shared across processes and should be cleared out and re-initialized
+        for each new process.
+        """
+        ...
+
+    @abstractmethod
+    async def resolve(
+        self,
+        request: RequestT,
+        timeout_time: float,
+    ) -> tuple[ResolveStatus, ResponseT]:
+        """
+        An abstract method that must be implemented by subclasses.
+        This method should handle the resolution of a request through asyncio,
+        including any necessary backend processing and response handling.
+
+        :param request: The request to be resolved generated by the load generator.
+        :param timeout_time: The timeout time for the request, if there is no timeout
+            given, then this will be math.inf.
+        :return: The response from the worker.
+        """
+        ...
+
+    async def get_request(
+        self, requests_queue: multiprocessing.Queue
+    ) -> Optional[WorkerProcessRequest[RequestT]]:
+        return await asyncio.to_thread(requests_queue.get)  # type: ignore[attr-defined]
+
+    async def send_result(
+        self,
+        results_queue: multiprocessing.Queue,
+        result: WorkerProcessResult[RequestT, ResponseT],
+    ):
+        await asyncio.to_thread(results_queue.put, result)  # type: ignore[attr-defined]
+
+    async def resolve_scheduler_request(
+        self,
+        request: Any,
+        queued_time: float,
+        dequeued_time: float,
+        start_time: float,
+        timeout_time: float,
+        results_queue: multiprocessing.Queue,
+        process_id: int,
+    ):
+        info = SchedulerRequestInfo(
+            targeted_start_time=start_time,
+            queued_time=queued_time,
+            dequeued_time=dequeued_time,
+            scheduled_time=time.time(),
+            process_id=process_id,
+        )
+        result: WorkerProcessResult[RequestT, ResponseT] = WorkerProcessResult(
+            type_="request_scheduled",
+            request=request,
+            response=None,
+            info=info,
+        )
+        asyncio.create_task(self.send_result(results_queue, result))
+
+        if (wait_time := start_time - time.time()) > 0:
+            await asyncio.sleep(wait_time)
+
+        info.worker_start = time.time()
+        result = WorkerProcessResult(
+            type_="request_start",
+            request=request,
+            response=None,
+            info=info,
+        )
+        asyncio.create_task(self.send_result(results_queue, result))
+
+        status, response = await self.resolve(request, timeout_time)
+        info.worker_end = time.time()
+        info.requested = status.requested
+        info.completed = status.completed
+        info.errored = status.errored
+        info.canceled = status.canceled
+        info.request_start = status.request_start
+        info.request_end = status.request_end
+        result = WorkerProcessResult(
+            type_="request_complete",
+            request=request,
+            response=response,
+            info=info,
+        )
+        asyncio.create_task(self.send_result(results_queue, result))
+
+    def process_loop_synchronous(
+        self,
+        requests_queue: multiprocessing.Queue,
+        results_queue: multiprocessing.Queue,
+        process_id: int,
+    ):
+        async def _process_runner():
+            while (
+                process_request := await self.get_request(requests_queue)
+            ) is not None:
+                dequeued_time = time.time()
+
+                await self.resolve_scheduler_request(
+                    request=process_request.request,
+                    queued_time=process_request.queued_time,
+                    dequeued_time=dequeued_time,
+                    start_time=process_request.start_time,
+                    timeout_time=process_request.timeout_time,
+                    results_queue=results_queue,
+                    process_id=process_id,
+                )
+
+        try:
+            asyncio.run(_process_runner())
+        except Exception as exc:  # noqa: BLE001
+            logger.error(
+                f"Error in worker process {process_id}: {exc}",
+                exc_info=True,
+                stack_info=True,
+            )
+
+    def process_loop_asynchronous(
+        self,
+        requests_queue: multiprocessing.Queue,
+        results_queue: multiprocessing.Queue,
+        max_concurrency: int,
+        process_id: int,
+    ):
+        async def _process_runner():
+            pending = asyncio.Semaphore(max_concurrency)
+
+            if pending.locked():
+                raise ValueError("Async worker called with max_concurrency < 1")
+
+            while (
+                process_request := await self.get_request(requests_queue)
+            ) is not None:
+                dequeued_time = time.time()
+
+                await pending.acquire()
+
+                def _task_done(_: asyncio.Task):
+                    nonlocal pending
+                    pending.release()
+
+                task = asyncio.create_task(
+                    self.resolve_scheduler_request(
+                        request=process_request.request,
+                        queued_time=process_request.queued_time,
+                        dequeued_time=dequeued_time,
+                        start_time=process_request.start_time,
+                        timeout_time=process_request.timeout_time,
+                        results_queue=results_queue,
+                        process_id=process_id,
+                    )
+                )
+                task.add_done_callback(_task_done)
+                await asyncio.sleep(0)  # enable start task immediately
+
+        try:
+            asyncio.run(_process_runner())
+        except Exception as exc:  # noqa: BLE001
+            logger.error(
+                f"Error in worker process {process_id}: {exc}",
+                exc_info=True,
+                stack_info=True,
+            )
+
+
+class GenerativeRequestsWorkerDescription(WorkerDescription):
+    type_: Literal["generative_requests_worker"] = "generative_requests_worker"  # type: ignore[assignment]
+    backend_type: BackendType
+    backend_target: str
+    backend_model: str
+    backend_info: dict[str, Any] = Field(
+        default_factory=dict,
+    )
+
+
+class GenerativeRequestsWorker(RequestsWorker[GenerationRequest, ResponseSummary]):
+    """
+    A class that handles the execution of requests using a backend.
+    This class is responsible for sending requests to the backend,
+    handling responses, and managing errors.
+
+    :param backend: The backend to use for handling requests.
+        This should be an instance of Backend such as an OpenAIHTTPBackend.
+    """
+
+    def __init__(self, backend: Backend):
+        self.backend = backend
+
+    @property
+    def description(self) -> GenerativeRequestsWorkerDescription:
+        """
+        Get the description of the worker.
+        :return: The description of the worker.
+        """
+        return GenerativeRequestsWorkerDescription(
+            backend_type=self.backend.type_,
+            backend_target=self.backend.target,
+            backend_model=self.backend.model or "None",
+            backend_info=self.backend.info,
+        )
+
+    async def prepare_multiprocessing(self):
+        """
+        Prepare the worker for multiprocessing.
+        This is useful for workers that have instance state that can not
+        be shared across processes and should be cleared out and re-initialized
+        for each new process.
+        """
+        await self.backend.prepare_multiprocessing()
+
+    def process_loop_synchronous(
+        self,
+        requests_queue: multiprocessing.Queue,
+        results_queue: multiprocessing.Queue,
+        process_id: int,
+    ):
+        asyncio.run(self.backend.validate())
+        super().process_loop_synchronous(
+            requests_queue=requests_queue,
+            results_queue=results_queue,
+            process_id=process_id,
+        )
+
+    def process_loop_asynchronous(
+        self,
+        requests_queue: multiprocessing.Queue,
+        results_queue: multiprocessing.Queue,
+        max_concurrency: int,
+        process_id: int,
+    ):
+        asyncio.run(self.backend.validate())
+        super().process_loop_asynchronous(
+            requests_queue=requests_queue,
+            results_queue=results_queue,
+            max_concurrency=max_concurrency,
+            process_id=process_id,
+        )
+
+    async def resolve(
+        self,
+        request: GenerationRequest,
+        timeout_time: float,
+    ) -> tuple[ResolveStatus, ResponseSummary]:
+        """
+        Resolve a request by sending it to the backend and handling the response.
+        This method sends the request to the backend, waits for a response,
+        and handles any errors that may occur during the process.
+
+        :param request: The request to resolve.
+        :param timeout_time: The time to wait for a response before timing out.
+            If timeout_time is math.inf, the request will not timeout.
+        :return: A ResponseSummary object containing the response from the backend.
+            If an error occurs, the ResponseSummary will contain the error message.
+        """
+        resolve_start_time = time.time()
+        response = None
+        error: Optional[str] = None
+        status = ResolveStatus(
+            requested=False,
+            completed=False,
+            errored=False,
+            canceled=False,
+            request_start=-1,
+            request_end=-1,
+        )
+
+        try:
+            if timeout_time < time.time():
+                raise asyncio.TimeoutError(
+                    "The timeout time has already passed."
+                )  # exit early
+
+            status.requested = True
+            request_func, request_kwargs = self._create_request_func_kwargs(request)
+
+            async def _runner():
+                # wrap function so we can enforce timeout and
+                # still return the latest state from the backend
+                async for resp in request_func(**request_kwargs):  # type: ignore[operator]
+                    nonlocal response
+                    response = resp
+
+            await asyncio.wait_for(
+                _runner(),
+                timeout=timeout_time - time.time() if timeout_time < math.inf else None,
+            )
+
+            if not response:
+                raise ValueError(
+                    f"No response received for request: {request} "
+                    f"and backend: {self.backend}"
+                )
+            if not isinstance(response, ResponseSummary):
+                raise ValueError(
+                    f"Received no ResponseSummary for request: {request} "
+                    f"and backend: {self.backend}, received: {response}"
+                )
+
+            status.completed = True
+        except asyncio.TimeoutError:
+            error = "TimeoutError: The request timed out before completing."
+            status.errored = True
+            status.canceled = True
+        except Exception as exc:  # noqa: BLE001
+            error = str(exc)
+            status.errored = True
+
+        return self._handle_response(
+            status=status,
+            request=request,
+            response=response,
+            error=error,
+            resolve_start_time=resolve_start_time,
+        )
+
+    def _create_request_func_kwargs(
+        self,
+        request: GenerationRequest,
+    ) -> tuple[
+        AsyncGenerator[Union[StreamingTextResponse, ResponseSummary], None],
+        dict[str, Any],
+    ]:
+        request_func: AsyncGenerator[
+            Union[StreamingTextResponse, ResponseSummary], None
+        ]
+        request_kwargs: dict[str, Any]
+
+        if request.request_type == "text_completions":
+            request_func = self.backend.text_completions  # type: ignore[assignment]
+            request_kwargs = {
+                "prompt": request.content,
+                "request_id": request.request_id,
+                "prompt_token_count": request.stats.get("prompt_tokens", None),
+                "output_token_count": request.constraints.get("output_tokens", None),
+                **request.params,
+            }
+        elif request.request_type == "chat_completions":
+            request_func = self.backend.chat_completions  # type: ignore[assignment]
+            request_kwargs = {
+                "content": request.content,
+                "request_id": request.request_id,
+                "prompt_token_count": request.stats.get("prompt_tokens", None),
+                "output_token_count": request.constraints.get("output_tokens", None),
+                **request.params,
+            }
+        else:
+            raise ValueError(
+                f"Invalid request type: {request.request_type} for {request}"
+            )
+
+        return request_func, request_kwargs
+
+    def _handle_response(
+        self,
+        status: ResolveStatus,
+        request: GenerationRequest,
+        response: Any,
+        error: Optional[str],
+        resolve_start_time: float,
+    ) -> tuple[ResolveStatus, ResponseSummary]:
+        if response is None or not isinstance(
+            response, (ResponseSummary, StreamingTextResponse)
+        ):
+            # nothing received or invalid response, fill in defaults for error
+            if response:
+                error = str(
+                    ValueError(
+                        f"Invalid response: {type(response)} for request: {request}; "
+                    )
+                ) + (error or "")
+
+            response = ResponseSummary(
+                value="",
+                request_args=RequestArgs(
+                    target=self.backend.target,
+                    headers={},
+                    payload={},
+                ),
+                start_time=resolve_start_time,
+                end_time=status.request_end,
+                first_iter_time=None,
+                last_iter_time=None,
+                request_id=request.request_id,
+                error=error or "Unknown error",
+            )
+        elif isinstance(response, StreamingTextResponse):
+            response = ResponseSummary(
+                value=response.value,
+                request_args=RequestArgs(
+                    target=self.backend.target,
+                    headers={},
+                    payload={},
+                ),
+                start_time=response.start_time,
+                end_time=time.time(),
+                first_iter_time=response.first_iter_time,
+                last_iter_time=response.time if response.iter_count > 0 else None,
+                request_prompt_tokens=request.stats.get("prompt_tokens", None),
+                request_output_tokens=request.constraints.get("output_tokens", None),
+                response_prompt_tokens=None,
+                response_output_tokens=response.iter_count,
+                request_id=request.request_id,
+                error=error or "Unknown error",
+            )
+
+        response.error = error
+        status.request_start = response.start_time
+        status.request_end = response.end_time
+
+        return status, response

guidellm/utils/__init__.py

@@ -1,40 +1,27 @@
-from .injector import create_report, inject_data
-from .progress import BenchmarkReportProgress
+from .colors import Colors
+from .hf_transformers import (
+    check_load_processor,
+)
+from .random import IntegerRangeSampler
 from .text import (
+    EndlessTextCreator,
     clean_text,
     filter_text,
-    is_path,
-    is_path_like,
-    is_url,
+    is_puncutation,
     load_text,
-    load_text_lines,
-    parse_text_objects,
-    split_lines_by_punctuation,
     split_text,
-)
-from .transformers import (
-    load_transformers_dataset,
-    resolve_transformers_dataset,
-    resolve_transformers_dataset_column,
-    resolve_transformers_dataset_split,
+    split_text_list_by_length,
 )
 
 __all__ = [
-    "BenchmarkReportProgress",
-    "clean_text",
-    "create_report",
+    "IntegerRangeSampler",
+    "Colors",
+    "check_load_processor",
     "filter_text",
-    "inject_data",
-    "is_path",
-    "is_path_like",
-    "is_url",
-    "load_text",
-    "load_text_lines",
-    "load_transformers_dataset",
-    "parse_text_objects",
-    "resolve_transformers_dataset",
-    "resolve_transformers_dataset_column",
-    "resolve_transformers_dataset_split",
-    "split_lines_by_punctuation",
+    "clean_text",
     "split_text",
+    "load_text",
+    "is_puncutation",
+    "EndlessTextCreator",
+    "split_text_list_by_length",
 ]

guidellm/utils/colors.py

@@ -0,0 +1,8 @@
+__all__ = ["Colors"]
+
+
+class Colors:
+    INFO: str = "light_steel_blue"
+    PROGRESS: str = "dark_slate_gray1"
+    SUCCESS: str = "chartreuse1"
+    ERROR: str = "orange_red1"

guidellm/utils/hf_transformers.py

@@ -0,0 +1,35 @@
+from pathlib import Path
+from typing import Any, Optional, Union
+
+from transformers import AutoTokenizer, PreTrainedTokenizerBase  # type: ignore[import]
+
+__all__ = [
+    "check_load_processor",
+]
+
+
+def check_load_processor(
+    processor: Optional[Union[str, Path, PreTrainedTokenizerBase]],
+    processor_args: Optional[dict[str, Any]],
+    error_msg: str,
+) -> PreTrainedTokenizerBase:
+    if processor is None:
+        raise ValueError(f"Processor/Tokenizer is required for {error_msg}.")
+
+    try:
+        if isinstance(processor, (str, Path)):
+            loaded = AutoTokenizer.from_pretrained(
+                processor,
+                **(processor_args or {}),
+            )
+        else:
+            loaded = processor
+    except Exception as err:
+        raise ValueError(
+            f"Failed to load processor/Tokenizer for {error_msg}."
+        ) from err
+
+    if not isinstance(loaded, PreTrainedTokenizerBase):
+        raise ValueError(f"Invalid processor/Tokenizer for {error_msg}.")
+
+    return loaded

guidellm/utils/random.py

@@ -0,0 +1,43 @@
+import random
+from collections.abc import Iterator
+from typing import Optional
+
+__all__ = ["IntegerRangeSampler"]
+
+
+class IntegerRangeSampler:
+    def __init__(
+        self,
+        average: int,
+        variance: Optional[int],
+        min_value: Optional[int],
+        max_value: Optional[int],
+        random_seed: int,
+    ):
+        self.average = average
+        self.variance = variance
+        self.min_value = min_value
+        self.max_value = max_value
+        self.seed = random_seed
+        self.rng = random.Random(random_seed)  # noqa: S311
+
+    def __iter__(self) -> Iterator[int]:
+        calc_min = self.min_value
+        if calc_min is None:
+            calc_min = max(
+                1, self.average - 5 * self.variance if self.variance else self.average
+            )
+        calc_max = self.max_value
+        if calc_max is None:
+            calc_max = (
+                self.average + 5 * self.variance if self.variance else self.average
+            )
+
+        while True:
+            if calc_min == calc_max:
+                yield calc_min
+            elif not self.variance:
+                yield self.rng.randint(calc_min, calc_max + 1)
+            else:
+                rand = self.rng.gauss(self.average, self.variance)
+                yield round(max(calc_min, min(calc_max, rand)))