Fix streaming bug (#820)

sgl-project · Jul 30, 2024 · daf593a · daf593a
1 parent bece265
commit daf593a
Showing 1 changed file with 17 additions and 2 deletions.
diff --git a/python/sglang/lang/interpreter.py b/python/sglang/lang/interpreter.py
@@ -553,6 +553,7 @@ def _execute_select(self, expr: SglSelect):
                 "output_token_logprobs": output_token_logprobs,
             }
             self.variable_event[name].set()
+            self.stream_var_event[name].set()
         self.text_ += decision
 
     def _execute_variable(self, expr: SglVariable):
@@ -778,7 +779,14 @@ def text_iter(self, var_name: Optional[str] = None):
                     if self.stream_executor.is_finished:
                         break
             else:
-                event = self.stream_executor.stream_var_event[var_name]
+                event = None
+                while not event:
+                    if var_name in self.stream_executor.stream_var_event:
+                        event = self.stream_executor.stream_var_event[var_name]
+                    if self.stream_executor.is_finished:
+                        yield ""
+                        return
+
                 while True:
                     event.wait()
                     event.clear()
@@ -813,7 +821,14 @@ async def text_async_iter(
                     if self.stream_executor.is_finished:
                         break
             else:
-                event = self.stream_executor.stream_var_event[var_name]
+                event = None
+                while not event:
+                    if var_name in self.stream_executor.stream_var_event:
+                        event = self.stream_executor.stream_var_event[var_name]
+                    if self.stream_executor.is_finished:
+                        yield ""
+                        return
+
                 while True:
                     await loop.run_in_executor(None, event.wait)
                     event.clear()