Simplify async streaming using run_in_executor instead of manual thread

parent 8af5cf90
...@@ -1047,14 +1047,9 @@ class VulkanBackend(ModelBackend): ...@@ -1047,14 +1047,9 @@ class VulkanBackend(ModelBackend):
total_content = "" total_content = ""
chunk_count = 0 chunk_count = 0
queue = asyncio.Queue()
stream_done = asyncio.Event()
stream_error = None
def run_stream(): def sync_generator():
"""Run the synchronous stream in a thread.""" """Synchronous generator that runs in executor."""
nonlocal stream_error
try:
print(f"DEBUG: generate_chat_stream: Calling create_chat_completion with tools={tools}") print(f"DEBUG: generate_chat_stream: Calling create_chat_completion with tools={tools}")
stream = self.model.create_chat_completion( stream = self.model.create_chat_completion(
messages=messages, messages=messages,
...@@ -1067,25 +1062,17 @@ class VulkanBackend(ModelBackend): ...@@ -1067,25 +1062,17 @@ class VulkanBackend(ModelBackend):
) )
print(f"DEBUG: generate_chat_stream: Got stream object: {type(stream)}") print(f"DEBUG: generate_chat_stream: Got stream object: {type(stream)}")
for chunk in stream: for chunk in stream:
# Put chunk in queue for async consumer yield chunk
asyncio.run_coroutine_threadsafe(queue.put(chunk), loop)
asyncio.run_coroutine_threadsafe(queue.put(None), loop) # Sentinel
except Exception as e:
stream_error = e
asyncio.run_coroutine_threadsafe(queue.put(None), loop)
# Get the event loop try:
# Run synchronous generator in thread pool
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
gen = sync_generator()
# Start stream in background thread
import threading
thread = threading.Thread(target=run_stream)
thread.start()
try:
while True: while True:
chunk = await queue.get() try:
if chunk is None: # Sentinel chunk = await loop.run_in_executor(None, lambda: next(gen, None))
if chunk is None:
break break
chunk_count += 1 chunk_count += 1
...@@ -1098,6 +1085,8 @@ class VulkanBackend(ModelBackend): ...@@ -1098,6 +1085,8 @@ class VulkanBackend(ModelBackend):
if content: if content:
total_content += content total_content += content
yield content yield content
except StopIteration:
break
print(f"DEBUG: generate_chat_stream yielded {chunk_count} chunks, total content length: {len(total_content)}") print(f"DEBUG: generate_chat_stream yielded {chunk_count} chunks, total content length: {len(total_content)}")
if chunk_count == 0 or not total_content.strip(): if chunk_count == 0 or not total_content.strip():
...@@ -1115,8 +1104,6 @@ class VulkanBackend(ModelBackend): ...@@ -1115,8 +1104,6 @@ class VulkanBackend(ModelBackend):
yield chunk yield chunk
else: else:
print(f"DEBUG: Stream completed with {chunk_count} chunks") print(f"DEBUG: Stream completed with {chunk_count} chunks")
finally:
thread.join(timeout=5)
def _manual_format_messages(self, messages: List[Dict]) -> str: def _manual_format_messages(self, messages: List[Dict]) -> str:
"""Manual fallback for formatting messages when create_chat_completion fails.""" """Manual fallback for formatting messages when create_chat_completion fails."""
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment