File size: 11,162 Bytes
31dc8dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 | from __future__ import annotations
import asyncio
import uuid
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from functools import partial
from typing import Any, AsyncIterator, Awaitable, Callable
from diffulex.sampling_params import SamplingParams
from diffulex.server.protocol import PromptInput
from diffulex.mixin.async_serving.engine import (
ServingAbort,
ServingCommand,
ServingError,
ServingEvent,
ServingGenerate,
ServingReply,
)
GenerateResult = ServingReply
@dataclass
class QueuedCommand:
command: ServingCommand
def default_engine_factory(model: str, **engine_kwargs):
from diffulex import strategy as _strategy # noqa: F401
from diffulex.engine.engine import DiffulexEngine
return DiffulexEngine(model, **engine_kwargs)
class EngineLoop:
"""Single-owner background loop for online serving.
The FastAPI event loop owns admission futures, while all engine calls run through
a one-worker executor. This keeps DiffulexEngine mutations serialized without
blocking request admission while a model step is running.
"""
def __init__(
self,
model: str,
engine_kwargs: dict[str, Any] | None = None,
*,
engine_factory: Callable[..., Any] | None = None,
idle_sleep_s: float = 0.001,
) -> None:
self.model = model
self.engine_kwargs = dict(engine_kwargs or {})
self.engine_factory = engine_factory or default_engine_factory
self.idle_sleep_s = idle_sleep_s
self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="diffulex-engine")
self._queue: asyncio.Queue[QueuedCommand] | None = None
self._task: asyncio.Task | None = None
self._started: asyncio.Future[None] | None = None
self._stop_event: asyncio.Event | None = None
self._engine = None
self._futures: dict[str, asyncio.Future[GenerateResult]] = {}
self._stream_queues: dict[str, asyncio.Queue[ServingEvent | None]] = {}
self._engine_has_work = False
@property
def model_id(self) -> str:
return self.engine_kwargs.get("model_name") or self.model
@property
def is_running(self) -> bool:
return self._task is not None and not self._task.done()
async def start(self) -> None:
if self.is_running:
if self._started is not None and not self._started.done():
await self._started
return
self._queue = asyncio.Queue()
self._stop_event = asyncio.Event()
self._started = asyncio.get_running_loop().create_future()
self._task = asyncio.create_task(self.run_loop(), name="diffulex-engine-loop")
await self._started
async def stop(self) -> None:
if self._stop_event is not None:
self._stop_event.set()
if self._task is not None:
await self._task
self._executor.shutdown(wait=True)
async def generate(self, prompt: str | list[int], sampling_params: SamplingParams) -> GenerateResult:
if self._queue is None or self._task is None:
raise RuntimeError("EngineLoop.start() must be called before generate()")
if self._task.done():
raise RuntimeError("EngineLoop stopped before generate()")
loop = asyncio.get_running_loop()
future = loop.create_future()
request_id = f"diffulex-{uuid.uuid4().hex}"
self._futures[request_id] = future
queued = QueuedCommand(
command=ServingGenerate(
rid=request_id,
input=PromptInput(prompt),
sampling_params=sampling_params,
),
)
await self._queue.put(queued)
return await future
async def generate_stream(
self,
prompt: str | list[int],
sampling_params: SamplingParams,
*,
stream_mode: str = "denoise",
is_disconnected: Callable[[], Awaitable[bool]] | None = None,
) -> AsyncIterator[ServingEvent]:
if self._queue is None or self._task is None:
raise RuntimeError("EngineLoop.start() must be called before generate_stream()")
if self._task.done():
raise RuntimeError("EngineLoop stopped before generate_stream()")
request_id = f"diffulex-{uuid.uuid4().hex}"
stream_queue: asyncio.Queue[ServingEvent | None] = asyncio.Queue()
self._stream_queues[request_id] = stream_queue
await self._queue.put(
QueuedCommand(
command=ServingGenerate(
rid=request_id,
input=PromptInput(prompt),
sampling_params=sampling_params,
stream=True,
stream_mode=stream_mode,
),
)
)
completed = False
try:
while True:
try:
event = await asyncio.wait_for(stream_queue.get(), timeout=0.5)
except asyncio.TimeoutError:
if is_disconnected is not None and await is_disconnected():
break
continue
if event is None:
break
if isinstance(event, (ServingReply, ServingError)):
completed = True
yield event
if completed:
break
finally:
self._stream_queues.pop(request_id, None)
if not completed and self._queue is not None:
await self._queue.put(QueuedCommand(command=ServingAbort(request_id)))
async def render_chat_prompt(self, messages: list[dict[str, str]]) -> str:
if self._task is not None and self._task.done():
raise RuntimeError("EngineLoop stopped before render_chat_prompt()")
def render() -> str:
if self._engine is None:
raise RuntimeError("EngineLoop.start() must complete before render_chat_prompt()")
return self._engine.render_chat_prompt_for_serving(messages)
return await self.call_engine(render)
async def run_loop(self) -> None:
try:
self._engine = await self.call_engine(lambda: self.engine_factory(self.model, **self.engine_kwargs))
self.mark_started()
await self.serve()
except Exception as exc:
self.mark_started(exc)
raise
finally:
self.fail_pending(RuntimeError("EngineLoop stopped before request completion"))
self.discard_queued_commands()
if self._engine is not None:
await self.call_engine(self._engine.exit)
self._engine = None
def mark_started(self, exc: Exception | None = None) -> None:
if self._started is None or self._started.done():
return
if exc is None:
self._started.set_result(None)
else:
self._started.set_exception(exc)
async def serve(self) -> None:
assert self._queue is not None
assert self._stop_event is not None
while not self._stop_event.is_set():
commands = await self.collect_ready_commands(
wait_for_one=not self.has_waiters() and not self._engine_has_work,
)
if commands or self.has_waiters() or self._engine_has_work:
events = await self.call_engine(self._engine.run_serving_tick, commands)
self._engine_has_work = await self.call_engine(lambda: not self._engine.is_finished())
self.resolve_events(events)
await asyncio.sleep(0)
elif self.idle_sleep_s > 0:
await asyncio.sleep(self.idle_sleep_s)
async def collect_ready_commands(self, *, wait_for_one: bool) -> list[ServingCommand]:
assert self._queue is not None
assert self._stop_event is not None
queued_commands: list[QueuedCommand] = []
if wait_for_one:
get_task = asyncio.create_task(self._queue.get())
stop_task = asyncio.create_task(self._stop_event.wait())
done, pending = await asyncio.wait({get_task, stop_task}, return_when=asyncio.FIRST_COMPLETED)
for task in pending:
task.cancel()
if stop_task in done and self._stop_event.is_set():
if not get_task.done():
return []
if get_task in done:
queued_commands.append(get_task.result())
while True:
try:
queued_commands.append(self._queue.get_nowait())
except asyncio.QueueEmpty:
break
return [queued.command for queued in queued_commands if self.should_admit(queued.command)]
def should_admit(self, command: ServingCommand) -> bool:
if isinstance(command, ServingGenerate):
future = self._futures.get(command.request_id)
stream_queue = self._stream_queues.get(command.request_id)
if future is None and stream_queue is None:
self._futures.pop(command.request_id, None)
return False
if future is not None and future.cancelled():
self._futures.pop(command.request_id, None)
return False
return True
def resolve_events(self, events: list[ServingEvent]) -> None:
for event in events:
stream_queue = self._stream_queues.get(event.rid)
if stream_queue is not None:
stream_queue.put_nowait(event)
if isinstance(event, (ServingReply, ServingError)):
self._stream_queues.pop(event.rid, None)
stream_queue.put_nowait(None)
continue
future = self._futures.pop(event.rid, None)
if future is None or future.done():
continue
if isinstance(event, ServingError):
future.set_exception(RuntimeError(event.message))
elif isinstance(event, ServingReply):
future.set_result(event)
def has_waiters(self) -> bool:
return bool(self._futures or self._stream_queues)
async def call_engine(self, fn: Callable, *args, **kwargs):
loop = asyncio.get_running_loop()
return await loop.run_in_executor(self._executor, partial(fn, *args, **kwargs))
def discard_queued_commands(self) -> None:
if self._queue is None:
return
while True:
try:
self._queue.get_nowait()
except asyncio.QueueEmpty:
return
def fail_pending(self, exc: Exception) -> None:
for future in list(self._futures.values()):
if not future.done():
future.set_exception(exc)
self._futures.clear()
for request_id, stream_queue in list(self._stream_queues.items()):
stream_queue.put_nowait(ServingError(request_id, str(exc)))
stream_queue.put_nowait(None)
self._stream_queues.clear()
|