guidellm.collectors.process

class CollectorProcess:
    """
    Own the child process that runs a benchmark's collectors.

    Example:
        ::

            collectors = CollectorProcess(args)
            await collectors.start()
            try:
                await collectors.benchmark_started()
                ...  # run the benchmark
                await collectors.benchmark_stopped()
                results = await collectors.results(start_time, end_time)
            finally:
                await collectors.shutdown()
    """

    def __init__(self, args: list[CollectorArgs]):
        """
        :param args: Arguments for each collector, in the order results are returned
        """
        self.args = args
        self._process: BaseProcess | None = None
        self._connection: Connection | None = None

    async def start(self) -> None:
        """
        Start the child process and create the collectors in it.

        :raises RuntimeError: If a collector cannot be created
        """
        context = get_context(settings.mp_context_type)
        parent_connection, child_connection = context.Pipe()
        self._process = context.Process(
            target=_run_collectors,
            args=(child_connection, [arg.model_dump() for arg in self.args]),
            name="guidellm-collectors",
            daemon=True,
        )
        self._process.start()
        child_connection.close()
        self._connection = parent_connection
        await self._request("ready")

    async def benchmark_started(self) -> None:
        """Start every collector for a new benchmark."""
        await self._request("start")

    async def benchmark_stopped(self) -> None:
        """Stop every collector once the benchmark's requests have ended."""
        await self._request("stop")

    async def results(self, start_time: float, end_time: float) -> list[dict[str, Any]]:
        """
        Return each collector's result for the latest benchmark.

        :param start_time: Start of the measurement window in seconds since epoch
        :param end_time: End of the measurement window in seconds since epoch
        :return: One result per collector, in configuration order
        """
        return await self._request("result", start_time, end_time)

    async def shutdown(self) -> None:
        """Close the collectors and wait for the child process to exit."""
        if self._process is None or self._connection is None:
            return
        try:
            await self._request("close")
        except (EOFError, OSError, RuntimeError) as err:
            logger.warning("Collector process did not close cleanly: {}", err)
        self._connection.close()
        await asyncio.to_thread(self._process.join, _SHUTDOWN_TIMEOUT)
        if self._process.is_alive():
            self._process.terminate()
        self._process = None
        self._connection = None

    async def _request(self, command: str, *params: Any) -> Any:
        if self._connection is None:
            raise RuntimeError("Collector process has not been started")
        if command != "ready":
            self._connection.send((command, *params))
        status, payload = await asyncio.to_thread(self._connection.recv)
        if status != "ok":
            raise RuntimeError(f"Collector process failed on {command!r}: {payload}")

        return payload