Method _mainloop

tensorrt_llm/executor/postproc_worker.py:175–216

The loop for handle_response and keep producing outputs.

(self)

            await self._push_pipe.put_async(batch)

    async def _mainloop(self):
        ''' The loop for handle_response and keep producing outputs. '''

        async def handle_single_input(inp: PostprocWorker.Input,
                                      batch: List[PostprocWorker.Output]):
            assert isinstance(
                inp, PostprocWorker.Input
            ), f"Expect PostprocWorker.Input, got {type(inp)}."
            client_id = inp.rsp.client_id
            is_final = inp.rsp.result.is_final if is_llm_response(
                inp.rsp) else True
            res, metrics, perf_metrics, disaggregated_params = await self._handle_input(
                inp)
            batch.append(
                PostprocWorker.Output(
                    client_id=client_id,
                    res=res,
                    is_final=is_final,
                    metrics=metrics,
                    request_perf_metrics=perf_metrics,
                    disaggregated_params=disaggregated_params,
                ))
            if is_final:
                self._records.pop(client_id)

        while not self._to_stop.is_set():
            batch = []
            inputs: Optional[List[PostprocWorker.Input]
                             | PostprocWorker.
                             Input] = await self._pull_pipe.get_async()

            if not isinstance(inputs, list):
                inputs = [inputs]

            for inp in inputs:
                if inp is None:
                    self._to_stop.set()
                    yield None
                    break
                await handle_single_input(inp, batch)

            yield batch

    def start(self):
        ''' Start the workflow in the current thread. '''