add streaming support (#66)

Author: HaiHui886

llmserve/backend/llm/engines/_base.py

@@ -12,6 +12,7 @@
 from llmserve.backend.logger import get_logger
 from llmserve.backend.server.models import Args, Prompt
 import asyncio
+from typing import AsyncGenerator, Generator
 
 logger = get_logger(__name__)
 
@@ -63,4 +64,8 @@ async def predict(
 
     @abstractmethod
     async def check_health(self):
+        pass
+    
+    @abstractmethod
+    def stream_generate_texts(self, prompt: str) -> Generator[str, None, None]:
         pass

llmserve/backend/llm/engines/generic.py

@@ -31,6 +31,12 @@
 from llmserve.backend.server.models import Args, LLMConfig, Prompt, Response
 from llmserve.backend.server.utils import render_gradio_params
 from ._base import LLMEngine
+
+from typing import AsyncGenerator, Generator
+from queue import Empty
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
+
 logger = get_logger(__name__)
 
 @timeit
@@ -278,6 +284,12 @@ def ping(self) -> bool:
         """Ping the worker."""
         return True
 
+    async def worker_stream_generate_texts(self, prompt: str, **kwargs) -> Generator[str, None, None]: # type: ignore
+        logger.info(f"Call PredictionWorker.worker_stream_generate_texts with kwargs: {kwargs}")
+        for s in self.generator.streamGenerate(prompt, **kwargs):
+            logger.info(f"PredictionWorker.worker_stream_generate_texts -> yield ->{s}")
+            yield s
+    
 class GenericEngine(LLMEngine):
     base_worker_group = None
 
@@ -416,4 +428,15 @@ async def check_health(self):
                 raise RuntimeError(
                     f"At least one prediction worker is dead. Dead workers: {dead_actors}. "
                     "Reinitializing worker group."
-                )
+                )
+    
+    def stream_generate_texts(self, prompt: str) -> Generator[str, None, None]: # type: ignore
+        logger.info(f"GenericEngine.stream_generate_texts -> worker.length: {len(self.base_worker_group)}")
+        worker0 = self.base_worker_group[0]
+        for strHandle in worker0.worker_stream_generate_texts.remote(
+            prompt,
+            **self.args.model_config.generation.all_generate_kwargs if self.args.model_config.generation else {}
+        ):
+            val = ray.get(strHandle)
+            logger.info(f"GenericEngine.stream_generate_texts -> yield -> {val}")
+            yield val

llmserve/backend/llm/pipelines/_base.py

@@ -22,6 +22,8 @@
 from .processors import StopOnTokens
 from .utils import tokenize_stopping_sequences_where_needed
 
+from typing import AsyncGenerator, Generator
+
 if TYPE_CHECKING:
     from ..initializers._base import LLMInitializer
 
@@ -338,6 +340,9 @@ def _sanitize_parameters(
 
         return preprocess_params, forward_params, postprocess_params
 
+    @abstractmethod
+    def streamGenerate(self, prompt: str, **generate_kwargs) -> Generator[str, None, None]:
+        pass
 
 class StreamingPipeline(BasePipeline):
     def stream(

llmserve/backend/llm/pipelines/default_pipeline.py

@@ -12,6 +12,12 @@
 from .processors import StopOnTokens
 from .utils import construct_prompts, truncate_to_first_stop_token
 
+from typing import AsyncGenerator, Generator
+import asyncio
+from transformers import TextIteratorStreamer
+from threading import Thread
+from queue import Empty
+
 logger = get_logger(__name__)
 
 
@@ -160,3 +166,31 @@ def postprocess(self, model_outputs, **postprocess_kwargs) -> List[Response]:
             response.generation_time = model_outputs["generation_time"]
             response.postprocessing_time = et
         return decoded
+
+    def streamGenerate(self, prompt: str, **generate_kwargs) -> Generator[str, None, None]:
+        logger.info(f"DefaultPipeline.streamGenerate with generate_kwargs: {generate_kwargs}")
+        streamer = TextIteratorStreamer(self.tokenizer, timeout=0, skip_prompt=True, skip_special_tokens=True)
+        input_ids = self.tokenizer([prompt], return_tensors="pt")
+        # generation_kwargs = dict(input_ids, streamer=streamer, max_new_tokens=20)
+        max_new_tokens = 256
+        if generate_kwargs["max_new_tokens"]:
+            max_new_tokens = generate_kwargs["max_new_tokens"]
+        generation_kwargs = dict(input_ids, streamer=streamer, max_new_tokens=max_new_tokens)
+        thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
+        thread.start()
+        while True:
+            try:
+                for token in streamer:
+                    logger.info(f'DefaultPipeline.streamGenerate -> Yield -> "{token}" -> "{type(token)}"')
+                    yield token
+                break
+            except Empty:
+                asyncio.sleep(0.001)
+        
+        # start = 0
+        # while True:
+        #     val = prompt + str(start)
+        #     logger.info(f"PredictionWorker.worker_stream_generate_texts -> yield -> {val}")
+        #     yield val
+        #     start += 1
+        #     asyncio.sleep(1)

llmserve/backend/llm/predictor.py

@@ -14,6 +14,8 @@
 from llmserve.backend.logger import get_logger
 from llmserve.backend.server.models import Args, Prompt
 
+from typing import AsyncGenerator, Generator
+
 initialize_node_remote = ray.remote(initialize_node)
 logger = get_logger(__name__)
 
@@ -59,7 +61,6 @@ async def rollover(self, scaling_config: ScalingConfig, pg_timeout_s: float = 60
                 args = self.args
             )
 
-
             self.new_worker_group = await self._create_worker_group(
                 scaling_config, pg_timeout_s=pg_timeout_s
             )
@@ -178,4 +179,10 @@ async def _predict_async(
 
     # Called by Serve to check the replica's health.
     async def check_health(self):
-        self.engine.check_health()
+        self.engine.check_health()
+        
+    async def stream_generate_texts(self, prompt: str) -> Generator[str, None, None]: # type: ignore
+        logger.info(f"call LLMPredictor.stream_generate_texts")
+        for s in self.engine.stream_generate_texts(prompt):
+            logger.info(f"LLMPredictor.stream_generate_texts -> yield ->{s}")
+            yield s

llmserve/backend/server/app.py

@@ -46,6 +46,10 @@
 from llmserve.api import sdk
 from llmserve.common.utils import _replace_prefix, _reverse_prefix
 
+from starlette.responses import StreamingResponse
+from typing import AsyncGenerator, Generator
+from ray.serve.handle import DeploymentHandle, DeploymentResponseGenerator
+
 # logger = get_logger(__name__)
 logger = get_logger("ray.serve")
 
@@ -303,7 +307,6 @@ async def generate_text_batch(
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}:{self.args.model_config.model_id}"
 
-
 @serve.deployment(
     # TODO make this configurable in llmserve run
     autoscaling_config={
@@ -315,12 +318,16 @@ def __repr__(self) -> str:
 )
 @serve.ingress(app)
 class RouterDeployment:
-    def __init__(
-        self, models: Dict[str, ClassNode], model_configurations: Dict[str, Args]
-    ) -> None:
+    def __init__(self, models: Dict[str, DeploymentHandle], model_configurations: Dict[str, Args]) -> None:
         self._models = models
         # TODO: Remove this once it is possible to reconfigure models on the fly
         self._model_configurations = model_configurations
+        logger.info(f"init: _models.keys: {self._models.keys()}")
+        # logger.info(f"init model_configurations: {model_configurations}")
+        for modelkey in self._models.keys():
+            if self._model_configurations[modelkey].model_config.stream:
+                logger.info(f"Set stream=true for {modelkey}")
+                self._models[modelkey] = self._models[modelkey].options(stream=True)
 
     @app.post("/{model}/run/predict")
     async def predict(self, model: str, prompt: Union[Prompt, List[Prompt]]) -> Union[Dict[str, Any], List[Dict[str, Any]], List[Any]]:
@@ -364,6 +371,30 @@ async def metadata(self, model: str) -> Dict[str, Dict[str, Any]]:
     async def models(self) -> List[str]:
         return list(self._models.keys())
 
+    @app.post("/run/stream")
+    def streamer(self, data: dict) -> StreamingResponse:
+        logger.info(f"data: {data}")
+        logger.info(f'Got stream -> body: {data}, keys: {self._models.keys()}')
+        prompt = data.get("prompt")
+        model = data.get("model")
+        modelKeys = list(self._models.keys())
+        modelID = model
+        for item in modelKeys:
+            logger.info(f"_reverse_prefix(item): {_reverse_prefix(item)}")
+            if _reverse_prefix(item) == model:
+                modelID = item
+                logger.info(f"set stream model id: {item}")
+        logger.info(f"search stream model key: {modelID}")
+        return StreamingResponse(self.streamer_generate_text(modelID, prompt), media_type="text/plain")
+
+    async def streamer_generate_text(self, modelID: str, prompt: str) -> AsyncGenerator[str, None]:
+        logger.info(f'streamer_generate_text: {modelID}, prompt: "{prompt}"')
+        r: DeploymentResponseGenerator = self._models[modelID].stream_generate_texts.remote(prompt)
+        async for i in r:
+            # logger.info(f"RouterDeployment.streamer_generate_text -> yield -> {type(i)}->{i}")
+            if not isinstance(i, str):
+                continue
+            yield i
 
 @serve.deployment(
     # TODO make this configurable in llmserve run

llmserve/backend/server/models.py

@@ -377,6 +377,7 @@ def all_generate_kwargs(self) -> Dict[str, Any]:
 
 
 class LLMConfig(BaseModelExtended):
+    stream: bool = False # enable steaming api
     warmup: bool    # need warmup?
     model_task: str    # need verification, TODO
     model_id: str

models/text-generation--facebook--opt-125m.yaml

@@ -12,6 +12,7 @@ deployment_config:
   ray_actor_options:
     num_cpus: 0.1    # for a model deployment, we have 3 actor created, 1 and 2 will cost 0.1 cpu, and the model infrence will cost 6(see the setting in the end of the file)
 model_config:
+  stream: False
   warmup: True
   model_task: text-generation
   model_id: facebook/opt-125m
@@ -20,7 +21,7 @@ model_config:
     # s3_mirror_config:
       # endpoint_url: http://39.107.108.170:9000 # Optinal for custom S3 storage endpoint url 
       # bucket_uri: s3://opt-125m/facemodel/  # Must include hash file with commit id in repo
-      # bucket_uri: /tmp/hub/opt-125m/ # Local path of model with hash file
+      # bucket_uri: /Users/hhwang/models/opt-125m/ # Local path of model with hash file
       # git_uri: https://portal.opencsg.com/models/opt-125m.git # git address for git clone
     initializer:
       type: SingleDevice