Repository: kserve/kserve

Issue:
Support streaming for v1 predict endpoint
/kind feature

**Describe the solution you'd like**
in open inference protocol we have implemented the streaming for [v2 generate endpoint,](https://github.com/kserve/kserve/blob/ad88ab4ae567803320ea5044ef6bee3dd27fe639/python/kserve/kserve/protocol/rest/v2_endpoints.py#L211) for users still on v1 endpoint we would like to support delivering back streaming response if the handlers returns `AsyncIterator[Any]`.


in [model.py](https://github.com/kserve/kserve/blob/ad88ab4ae567803320ea5044ef6bee3dd27fe639/python/kserve/kserve/model.py#L240) we add `AsyncIterator[Any]` to the `postprocess` return type
```python
async def postprocess(self, result: Union[Dict, InferResponse], headers: Dict[str, str] = None) \
            -> Union[Dict, InferResponse, AsyncIterator[Any]]:
```

in [v1_endpoints.py](https://github.com/kserve/kserve/blob/ad88ab4ae567803320ea5044ef6bee3dd27fe639/python/kserve/kserve/protocol/rest/v1_endpoints.py#L83) we return the `StreamingResponse` instead if handler yields `AsyncIterator`.

```python
        if isinstance(response, AsyncIterator):
            return StreamingResponse(response)
        if not isinstance(response, dict):
            return Response(content=response, headers=response_headers)
            
        return response
```


**Anything else you would like to add:**
[Miscellaneous information that will assist in solving the issue.]


**Links to the design documents:**
[Optional, start with the short-form [RFC template](https://docs.google.com/document/d/1UcBeLfZ_JMGpVrPJmYtEIVH_9Y4U3AEKQdq_IKuOMrU) to outline your ideas and get early feedback.]
[Required, use the longer-form [design doc template](https://docs.google.com/document/d/1Mtoui_PP2a9N59NjYHnsvrdJ8t2iKFwIJAx1zRO_I1c) to specify and discuss your design in more detail]



Discussion:

Resolving patch:
diff --git a/python/kserve/kserve/model.py b/python/kserve/kserve/model.py
--- a/python/kserve/kserve/model.py
+++ b/python/kserve/kserve/model.py
@@ -300,7 +300,7 @@ async def _grpc_predict(self, payload: Union[ModelInferRequest, InferRequest], h
         return async_result
 
     async def predict(self, payload: Union[Dict, InferRequest, ModelInferRequest],
-                      headers: Dict[str, str] = None) -> Union[Dict, InferResponse]:
+                      headers: Dict[str, str] = None) -> Union[Dict, InferResponse, AsyncIterator[Any]]:
         """ The `predict` handler can be overridden for performing the inference.
             By default, the predict handler makes call to predictor for the inference step.
 
diff --git a/python/kserve/kserve/protocol/rest/v1_endpoints.py b/python/kserve/kserve/protocol/rest/v1_endpoints.py
--- a/python/kserve/kserve/protocol/rest/v1_endpoints.py
+++ b/python/kserve/kserve/protocol/rest/v1_endpoints.py
@@ -11,9 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Optional, Union, Dict, List
+from typing import Optional, Union, Dict, List, AsyncIterator
 
 from fastapi import Request, Response
+from starlette.responses import StreamingResponse
 
 from kserve.errors import ModelNotReady
 from ..dataplane import DataPlane
@@ -80,8 +81,10 @@ async def predict(self, model_name: str, request: Request) -> Union[Response, Di
                                                            response=response,
                                                            headers=headers, req_attributes=req_attributes)
 
-        if not isinstance(response, dict):
+        if isinstance(response, bytes) or isinstance(response, str):
             return Response(content=response, headers=response_headers)
+        if isinstance(response, AsyncIterator):
+            return StreamingResponse(content=response)
         return response
 
     async def explain(self, model_name: str, request: Request) -> Union[Response, Dict]:


Test patch:
diff --git a/python/kserve/test/test_server.py b/python/kserve/test/test_server.py
--- a/python/kserve/test/test_server.py
+++ b/python/kserve/test/test_server.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import io
 import json
 import os
@@ -21,6 +22,7 @@
 
 import avro.io
 import avro.schema
+import httpx
 import pytest
 from cloudevents.conversion import to_binary, to_structured
 from cloudevents.http import CloudEvent
@@ -48,6 +50,8 @@
         }
         '''
 
+fake_stream_data = 'some streamed data'
+
 
 def dummy_cloud_event(data, set_contenttype: bool = False, add_extension: bool = False,
                       contenttype: str = "application/json"):
@@ -68,6 +72,51 @@ def dummy_cloud_event(data, set_contenttype: bool = False, add_extension: bool =
     return event
 
 
+async def fake_data_streamer():
+    for i in range(10):
+        yield fake_stream_data.encode()
+        await asyncio.sleep(0.5)  # sleep 1/2 second
+
+
+class DummyStreamModel(Model):
+    def __init__(self, name):
+        super().__init__(name)
+        self.name = name
+        self.ready = False
+
+    def load(self):
+        self.ready = True
+
+    async def predict(self, request, headers=None):
+        return fake_data_streamer()
+
+
+class TestStreamPredict:
+    @pytest.fixture(scope="class")
+    def app(self):  # pylint: disable=no-self-use
+        model = DummyStreamModel("TestModel")
+        model.load()
+        server = ModelServer()
+        server.register_model(model)
+        rest_server = RESTServer(server.dataplane, server.model_repository_extension)
+        return rest_server.create_application()
+
+    @pytest.fixture(scope='class')
+    def http_server_client(self, app):
+        return TestClient(app)
+
+    def test_predict_stream(self, http_server_client):
+        with http_server_client.stream('POST', '/v1/models/TestModel:predict',
+                                       content=b'{"instances":[[1,2]]}') as response:
+            response: httpx.Response
+            all_data = []
+            for value in response.iter_bytes():
+                data = value.decode()
+                assert fake_stream_data in data
+                all_data.append(data)
+        assert all([fake_stream_data in data for data in all_data]), "Unexpected number of streamed responses"
+
+
 class DummyModel(Model):
     def __init__(self, name):
         super().__init__(name)