vLLM AsyncLLMEngine 与 OpenAI API 源码精读

class AsyncStream:
    """A stream of RequestOutputs for a request that can be
    iterated over asynchronously."""

    def __init__(self, request_id: str) -> None:
        self.request_id = request_id
        self._queue: asyncio.Queue = asyncio.Queue()
        self._finished = False

    def put(self, item: Union[RequestOutput, Exception]) -> None:
        if self._finished:
            return
        self._queue.put_nowait(item)

    def finish(self) -> None:
        self._queue.put_nowait(StopAsyncIteration())
        self._finished = True

    @property
    def finished(self) -> bool:
        return self._finished

    def __aiter__(self):
        return self

    async def __anext__(self) -> RequestOutput:
        result = await self._queue.get()
        if isinstance(result, Exception):
            raise result
        return result

class RequestTracker:
    """Synchronous abstraction for tracking requests."""

    def __init__(self) -> None:
        self._request_streams: Dict[str, AsyncStream] = {}
        self._finished_requests: asyncio.Queue[str] = asyncio.Queue()
        self._new_requests: asyncio.Queue[Tuple[AsyncStream,
                                                dict]] = asyncio.Queue()
        self.new_requests_event = asyncio.Event()

def add_request(self, request_id: str,
                **engine_add_request_kwargs) -> AsyncStream:
    """Add a request to be sent to the engine on the next background
    loop iteration."""
    if request_id in self._request_streams:
        raise KeyError(f"Request {request_id} already exists.")

    stream = AsyncStream(request_id)
    self._new_requests.put_nowait((stream, {
        "request_id": request_id,
        **engine_add_request_kwargs
    }))
    self.new_requests_event.set()
    return stream

def process_request_output(self, request_output: RequestOutput,
                           *, verbose: bool = False) -> None:
    """Process a request output from the engine."""
    request_id = request_output.request_id
    self._request_streams[request_id].put(request_output)
    if request_output.finished:
        self.abort_request(request_id)

class _AsyncLLMEngine(LLMEngine):
    """Extension of LLMEngine to add async methods."""

    async def step_async(self) -> List[RequestOutput]:
        seq_group_metadata_list, scheduler_outputs = self.scheduler.schedule()

        if not scheduler_outputs.is_empty():
            execute_model_req = ExecuteModelRequest(
                seq_group_metadata_list=seq_group_metadata_list,
                blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in,
                blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out,
                blocks_to_copy=scheduler_outputs.blocks_to_copy,
                num_lookahead_slots=scheduler_outputs.num_lookahead_slots,
                running_queue_size=scheduler_outputs.running_queue_size,
            )
            output = await self.model_executor.execute_model_async(
                execute_model_req)
        else:
            output = []

        request_outputs = self._process_model_outputs(
            output, scheduler_outputs.scheduled_seq_groups,
            scheduler_outputs.ignored_seq_groups, seq_group_metadata_list)
        self.do_log_stats(scheduler_outputs, output)
        return request_outputs

async def add_request_async(
    self, request_id: str, prompt: Optional[str],
    sampling_params: SamplingParams,
    prompt_token_ids: Optional[List[int]] = None,
    arrival_time: Optional[float] = None,
    lora_request: Optional[LoRARequest] = None,
    multi_modal_data: Optional[MultiModalData] = None,
) -> None:
    if arrival_time is None:
        arrival_time = time.time()
    prompt_token_ids = await self.encode_request_async(
        request_id=request_id, prompt=prompt,
        prompt_token_ids=prompt_token_ids, lora_request=lora_request)

    return self.add_request(request_id, prompt=prompt,
                            prompt_token_ids=prompt_token_ids,
                            sampling_params=sampling_params,
                            arrival_time=arrival_time,
                            lora_request=lora_request,
                            multi_modal_data=multi_modal_data)

class AsyncLLMEngine:
    _engine_class: Type[_AsyncLLMEngine] = _AsyncLLMEngine

    def __init__(self, worker_use_ray: bool, engine_use_ray: bool,
                 *args, log_requests: bool = True,
                 max_log_len: Optional[int] = None,
                 start_engine_loop: bool = True, **kwargs) -> None:
        self.worker_use_ray = worker_use_ray
        self.engine_use_ray = engine_use_ray
        self.log_requests = log_requests
        self.max_log_len = max_log_len
        self.engine = self._init_engine(*args, **kwargs)

        self.background_loop: Optional[asyncio.Future] = None
        self._background_loop_unshielded: Optional[asyncio.Task] = None
        self.start_engine_loop = start_engine_loop
        self._errored_with: Optional[BaseException] = None

        # Lazy initialized fields
        self._request_tracker: RequestTracker

@classmethod
def from_engine_args(cls, engine_args: AsyncEngineArgs,
                     start_engine_loop: bool = True,
                     usage_context: UsageContext = ...) -> "AsyncLLMEngine":
    engine_config = engine_args.create_engine_config()

    if engine_config.device_config.device_type == "neuron":
        executor_class = NeuronExecutorAsync
    elif engine_config.device_config.device_type == "cpu":
        executor_class = CPUExecutorAsync
    elif engine_config.parallel_config.worker_use_ray:
        initialize_ray_cluster(engine_config.parallel_config)
        executor_class = RayGPUExecutorAsync
    else:
        assert engine_config.parallel_config.world_size == 1
        executor_class = GPUExecutorAsync

    return cls(worker_use_ray, engine_use_ray,
               **engine_config.to_dict(),
               executor_class=executor_class, ...)

async def run_engine_loop(self):
    has_requests_in_progress = False
    while True:
        if not has_requests_in_progress:
            logger.debug("Waiting for new requests...")
            await self._request_tracker.wait_for_new_requests()
            logger.debug("Got new requests!")

        # Abort if iteration takes too long due to unrecoverable errors
        try:
            has_requests_in_progress = await asyncio.wait_for(
                self.engine_step(), ENGINE_ITERATION_TIMEOUT_S)
        except asyncio.TimeoutError as exc:
            logger.error(
                "Engine iteration timed out. This should never happen!")
            self.set_errored(exc)
            raise
        await asyncio.sleep(0)

async def engine_step(self) -> bool:
    """Kick the engine to process the waiting requests.
    Returns True if there are in-progress requests."""

    new_requests, finished_requests = (
        self._request_tracker.get_new_and_finished_requests())

    for new_request in new_requests:
        try:
            if self.engine_use_ray:
                await self.engine.add_request.remote(**new_request)
            else:
                await self.engine.add_request_async(**new_request)
        except ValueError as e:
            self._request_tracker.process_exception(
                new_request["request_id"], e, verbose=self.log_requests)

    if finished_requests:
        await self._engine_abort(finished_requests)

    if self.engine_use_ray:
        request_outputs = await self.engine.step.remote()
    else:
        request_outputs = await self.engine.step_async()

    # Put the outputs into the corresponding streams.
    for request_output in request_outputs:
        self._request_tracker.process_request_output(
            request_output, verbose=self.log_requests)

    return len(request_outputs) > 0

async def generate(
    self, prompt: Optional[str],
    sampling_params: SamplingParams,
    request_id: str,
    prompt_token_ids: Optional[List[int]] = None,
    lora_request: Optional[LoRARequest] = None,
    multi_modal_data: Optional[MultiModalData] = None
) -> AsyncIterator[RequestOutput]:
    """Generate outputs for a request.
    ...
    Details:
        - If the engine is not running, start the background loop,
          which iteratively invokes engine_step to process requests.
        - Add the request to the engine's RequestTracker.
        - Wait for the request outputs from AsyncStream and yield them.
    """
    arrival_time = time.time()

    try:
        stream = await self.add_request(
            request_id, prompt, sampling_params,
            prompt_token_ids=prompt_token_ids,
            arrival_time=arrival_time,
            lora_request=lora_request,
            multi_modal_data=multi_modal_data,
        )

        async for request_output in stream:
            yield request_output
    except (Exception, asyncio.CancelledError) as e:
        self._abort(request_id)
        raise e

@app.get("/health")
async def health() -> Response:
    await openai_serving_chat.engine.check_health()
    return Response(status_code=200)

@app.get("/v1/models")
async def show_available_models():
    models = await openai_serving_chat.show_available_models()
    return JSONResponse(content=models.model_dump())

@app.post("/v1/chat/completions")
async def create_chat_completion(request: ChatCompletionRequest,
                                 raw_request: Request):
    generator = await openai_serving_chat.create_chat_completion(
        request, raw_request)
    if isinstance(generator, ErrorResponse):
        return JSONResponse(content=generator.model_dump(),
                            status_code=generator.code)
    if request.stream:
        return StreamingResponse(content=generator,
                                 media_type="text/event-stream")
    else:
        return JSONResponse(content=generator.model_dump())

@app.post("/v1/completions")
async def create_completion(request: CompletionRequest, raw_request: Request):
    generator = await openai_serving_completion.create_completion(
        request, raw_request)
    if request.stream:
        return StreamingResponse(content=generator,
                                 media_type="text/event-stream")
    else:
        return JSONResponse(content=generator.model_dump())

if __name__ == "__main__":
    args = parse_args()

    # CORS, authentication middleware ...

    engine_args = AsyncEngineArgs.from_cli_args(args)
    engine = AsyncLLMEngine.from_engine_args(
        engine_args, usage_context=UsageContext.OPENAI_API_SERVER)

    openai_serving_chat = OpenAIServingChat(
        engine, served_model_names, args.response_role,
        args.lora_modules, args.chat_template)
    openai_serving_completion = OpenAIServingCompletion(
        engine, served_model_names, args.lora_modules)

    uvicorn.run(app, host=args.host, port=args.port,
                timeout_keep_alive=TIMEOUT_KEEP_ALIVE, ...)

async def create_chat_completion(
    self, request: ChatCompletionRequest, raw_request: Request
) -> Union[ErrorResponse, AsyncGenerator[str, None],
           ChatCompletionResponse]:
    # 1. 模型检查
    error_check_ret = await self._check_model(request)
    if error_check_ret is not None:
        return error_check_ret

    # 2. 解析消息，应用 chat template
    conversation: List[ConversationMessage] = []
    for m in request.messages:
        messages, _ = self._parse_chat_message_content(
            m["role"], m["content"])
        conversation.extend(messages)
    prompt = self.tokenizer.apply_chat_template(
        conversation=conversation, tokenize=False,
        add_generation_prompt=request.add_generation_prompt)

    # 3. Tokenize 并构造采样参数
    request_id = f"cmpl-{random_uuid()}"
    prompt_ids, prompt_text = self._validate_prompt_and_tokenize(
        request, prompt=prompt)
    sampling_params = request.to_sampling_params()

    # 4. 调用 engine.generate()
    result_generator = self.engine.generate(
        prompt_text, sampling_params, request_id, prompt_ids,
        lora_request)

    # 5. 分流: streaming vs non-streaming
    if request.stream:
        return self.chat_completion_stream_generator(
            request, result_generator, request_id, conversation)
    else:
        return await self.chat_completion_full_generator(
            request, raw_request, result_generator, request_id,
            conversation)

async def chat_completion_stream_generator(
        self, request, result_generator, request_id, conversation
) -> AsyncGenerator[str, None]:
    model_name = self.served_model_names[0]
    created_time = int(time.time())
    previous_texts = [""] * request.n
    previous_num_tokens = [0] * request.n

    async for res in result_generator:
        if first_iteration:
            # 发送 role delta
            for i in range(request.n):
                chunk = ChatCompletionStreamResponse(
                    choices=[ChatCompletionResponseStreamChoice(
                        index=i, delta=DeltaMessage(role=role),
                        finish_reason=None)])
                yield f"data: {chunk.model_dump_json(...)}\n\n"
            first_iteration = False

        for output in res.outputs:
            delta_text = output.text[len(previous_texts[i]):]
            # 发送 content delta
            if output.finish_reason is None:
                chunk = ...  # delta with content
                yield f"data: {chunk.model_dump_json(...)}\n\n"
            else:
                chunk = ...  # final chunk with finish_reason + usage
                yield f"data: {chunk.model_dump_json(...)}\n\n"

    yield "data: [DONE]\n\n"

async def chat_completion_full_generator(
    self, request, raw_request, result_generator, request_id,
    conversation
) -> Union[ErrorResponse, ChatCompletionResponse]:
    final_res: Optional[RequestOutput] = None

    async for res in result_generator:
        if await raw_request.is_disconnected():
            # Abort the request if the client disconnects.
            await self.engine.abort(request_id)
            return self.create_error_response("Client disconnected")
        final_res = res

    assert final_res is not None
    # 构造完整的 ChatCompletionResponse ...
    return ChatCompletionResponse(choices=choices, usage=usage, ...)