Repository: jina-ai/jina
File: jina/serve/streamer.py (Python)

Code before:
        os.environ['JINA_STREAMER_ARGS'] = json.dumps(kwargs)


class _ExecutorStreamer:
    def __init__(self, connection_pool: GrpcConnectionPool, executor_name: str) -> None:
        self._connection_pool: GrpcConnectionPool = connection_pool
        self.executor_name = executor_name

    async def post(
        self,
        inputs: DocumentArray,
        request_size: int = 100,
        on: Optional[str] = None,
        parameters: Optional[Dict] = None,
        **kwargs,
    ):
        reqs = []
        for docs_batch in inputs.batch(batch_size=request_size, shuffle=False):
            req = DataRequest()
            req.header.exec_endpoint = on
            req.header.target_executor = self.executor_name
            req.parameters = parameters
            req.data.docs = docs_batch
            reqs.append(req)

        resp, _ = await self._connection_pool.send_requests_once(
            requests=reqs,
            deployment=self.executor_name,
            head=True,
            endpoint=on
        )
        
        return resp.docs
        

Review comment (suggestion):
do tests where u send more than one request, I believe this interface is not intended to be used like this.

Every `req` should lead to a `send_request` call and u should collect docs afterwards.

Code after:
        else:
            raise OSError('JINA_STREAMER_ARGS environment variable is not set')

    @staticmethod
    def _set_env_streamer_args(**kwargs):
        os.environ['JINA_STREAMER_ARGS'] = json.dumps(kwargs)

    async def warmup(self, stop_event: threading.Event):
        '''Executes warmup task on each deployment. This forces the gateway to establish connection and open a
        gRPC channel to each executor so that the first request doesn't need to experience the penalty of
        eastablishing a brand new gRPC channel.
        :param stop_event: signal to indicate if an early termination of the task is required for graceful teardown.
        '''
        self.logger.debug(f'Running GatewayRuntime warmup')
        deployments = {key for key in self._executor_addresses.keys()}

        try:
            deployment_warmup_tasks = []
            for deployment in deployments:
                deployment_warmup_tasks.append(
                    asyncio.create_task(
                        self._connection_pool.warmup(
                            deployment=deployment, stop_event=stop_event
                        )
                    )
                )

            await asyncio.gather(*deployment_warmup_tasks, return_exceptions=True)
        except Exception as ex:
            self.logger.error(f'error with GatewayRuntime warmup up task: {ex}')
            return


class _ExecutorStreamer:
    def __init__(self, connection_pool: GrpcConnectionPool, executor_name: str) -> None:
        self._connection_pool: GrpcConnectionPool = connection_pool
        self.executor_name = executor_name

    async def post(
        self,
        inputs: DocumentArray,
        request_size: int = 100,
        on: Optional[str] = None,
        parameters: Optional[Dict] = None,
        **kwargs,
    ):
        reqs = []
        for docs_batch in inputs.batch(batch_size=request_size, shuffle=False):
            req = DataRequest()
            req.header.exec_endpoint = on
            req.header.target_executor = self.executor_name