Simplify async streaming using run_in_executor instead of manual thread

parent 8af5cf90
......@@ -1047,14 +1047,9 @@ class VulkanBackend(ModelBackend):
total_content = ""
chunk_count = 0
queue = asyncio.Queue()
stream_done = asyncio.Event()
stream_error = None
def run_stream():
"""Run the synchronous stream in a thread."""
nonlocal stream_error
try:
def sync_generator():
"""Synchronous generator that runs in executor."""
print(f"DEBUG: generate_chat_stream: Calling create_chat_completion with tools={tools}")
stream = self.model.create_chat_completion(
messages=messages,
......@@ -1067,25 +1062,17 @@ class VulkanBackend(ModelBackend):
)
print(f"DEBUG: generate_chat_stream: Got stream object: {type(stream)}")
for chunk in stream:
# Put chunk in queue for async consumer
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)
yield chunk
# Get the event loop
try:
# Run synchronous generator in thread pool
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:
chunk = await queue.get()
if chunk is None: # Sentinel
try:
chunk = await loop.run_in_executor(None, lambda: next(gen, None))
if chunk is None:
break
chunk_count += 1
......@@ -1098,6 +1085,8 @@ class VulkanBackend(ModelBackend):
if content:
total_content += content
yield content
except StopIteration:
break
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():
......@@ -1115,8 +1104,6 @@ class VulkanBackend(ModelBackend):
yield chunk
else:
print(f"DEBUG: Stream completed with {chunk_count} chunks")
finally:
thread.join(timeout=5)
def _manual_format_messages(self, messages: List[Dict]) -> str:
"""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