52 lines
2.1 KiB
Python
52 lines
2.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.sched.base_scheduler import _BaseScheduler
|
|
from vllm_omni.diffusion.sched.interface import (
|
|
DiffusionRequestStatus,
|
|
DiffusionSchedulerOutput,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.diffusion.worker.utils import RunnerOutput
|
|
|
|
|
|
class RequestScheduler(_BaseScheduler):
|
|
"""Diffusion scheduler with vLLM-style waiting/running queues."""
|
|
|
|
def add_request(self, request: OmniDiffusionRequest) -> str:
|
|
return super().add_request(request)
|
|
|
|
def schedule(self) -> DiffusionSchedulerOutput:
|
|
return super().schedule()
|
|
|
|
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
|
|
scheduled_request_ids = sched_output.scheduled_request_ids
|
|
if not scheduled_request_ids:
|
|
return set()
|
|
|
|
terminal_statuses: dict[str, DiffusionRequestStatus] = {}
|
|
terminal_errors: dict[str, str | None] = {}
|
|
for request_id in scheduled_request_ids:
|
|
state = self._request_states.get(request_id)
|
|
if state is None or state.is_finished():
|
|
continue
|
|
req_output = output.get_request_output(request_id)
|
|
result = req_output.result if req_output is not None else None
|
|
if result is None:
|
|
terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
|
|
terminal_errors[request_id] = "No output result"
|
|
elif result.error:
|
|
terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
|
|
terminal_errors[request_id] = result.error
|
|
else:
|
|
terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_COMPLETED
|
|
terminal_errors[request_id] = None
|
|
|
|
return self._finalize_update_from_output(sched_output, terminal_statuses, terminal_errors)
|