fix issue: non-support streaming pipeline cannot work when call it as…

… streaming (#84)
OpenCSGs · Apr 10, 2024 · f8be0c5 · f8be0c5
1 parent 4208e4f
commit f8be0c5
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/llmserve/backend/llm/engines/generic.py b/llmserve/backend/llm/engines/generic.py
@@ -376,7 +376,7 @@ async def launch_engine(
         )
 
         self.base_worker_group = worker_group
-        self.can_stream = await asyncio.gather(*[worker_group[0].can_stream.remote()])
+        self.can_stream = ray.get(worker_group[0].can_stream.remote())
         return worker_group
 
     async def predict(
@@ -503,5 +503,5 @@ async def stream(
                 f"Pipeline {self.args.model_config.initialization.pipeline} does not support streaming. Ignoring queue."
             )
             yield await self.predict(
-                prompts, timeout_s=timeout_s, start_timestamp=start_timestamp
+                prompts, timeout_s=timeout_s, start_timestamp=start_timestamp, lock=lock
             )
diff --git a/llmserve/backend/llm/predictor.py b/llmserve/backend/llm/predictor.py
@@ -174,7 +174,7 @@ async def _predict_async(
         Returns:
             A list of generated texts.
         """
-        prediction = await self.engine.predict(prompts, timeout_s=timeout_s, start_timestamp=start_timestamp, lock = self._base_worker_group_lock)
+        prediction = await self.engine.predict(prompts, timeout_s=timeout_s, start_timestamp=start_timestamp, lock=self._base_worker_group_lock)
         return prediction
 
     async def _stream_async(
@@ -197,7 +197,7 @@ async def _stream_async(
         Returns:
             A list of generated texts.
         """
-        async for s in self.engine.stream(prompts, timeout_s=timeout_s, start_timestamp=start_timestamp, lock = self._base_worker_group_lock):
+        async for s in self.engine.stream(prompts, timeout_s=timeout_s, start_timestamp=start_timestamp, lock=self._base_worker_group_lock):
             yield s
 
     # Called by Serve to check the replica's health.