Method stream_with_timeout

tensorrt_llm/executor/rpc/rpc_server.py:538–565

()

            if adjusted_timeout is not None and adjusted_timeout > 0:
                # Create a task for the async generator with timeout
                async def stream_with_timeout():
                    nonlocal chunk_index
                    async for result in func(*req.args, **req.kwargs):
                        if result is None or result == []:
                            # Skip None values or empty list to save bandwidth
                            # TODO[Superjomn]: add a flag to control this behavior
                            continue
                        # Check if shutdown was triggered
                        if self._shutdown_event.is_set():
                            raise RPCCancelled(
                                "Server is shutting down, streaming cancelled")

                        logger_debug(
                            f"[server] RPC Server got data and ready to send result {result}"
                        )
                        response = RPCResponse(req.request_id,
                                               result=result,
                                               error=None,
                                               is_streaming=True,
                                               chunk_index=chunk_index,
                                               stream_status='data')
                        if not await self._send_response(req, response):
                            # Stop streaming after a pickle error
                            return
                        logger_debug(
                            f"[server] Sent response for request {req.request_id}",
                            color="green")
                        chunk_index += 1

                # Use wait_for for timeout handling
                await asyncio.wait_for(stream_with_timeout(),