249 lines
11 KiB
Python
249 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Callable
|
|
from concurrent.futures import Future
|
|
from dataclasses import replace
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from vllm.utils.import_utils import resolve_obj_by_qualname
|
|
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
|
|
from vllm.v1.outputs import KVConnectorOutput
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.sched.interface import CachedRequestData
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
|
|
from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
|
|
|
|
|
|
class DiffusionExecutor(ABC):
|
|
"""Abstract base class for Diffusion executors."""
|
|
|
|
uses_multiproc: bool = False
|
|
|
|
@staticmethod
|
|
def get_class(od_config: OmniDiffusionConfig) -> type[DiffusionExecutor]:
|
|
executor_class: type[DiffusionExecutor]
|
|
distributed_executor_backend = od_config.distributed_executor_backend
|
|
# Mirror vLLM's `world_size == 1 -> "uni"` default
|
|
# (vllm/config/parallel.py). A single-GPU diffusion deployment has
|
|
# nothing to distribute: spawning a worker only adds MessageQueues,
|
|
# /dev/shm output segments, and a second model load. Explicit "mp"
|
|
# keeps process isolation and RPC timeouts. Multi-GPU still defaults
|
|
# to "mp".
|
|
if distributed_executor_backend is None:
|
|
num_gpus = od_config.num_gpus or 1
|
|
distributed_executor_backend = "uni" if num_gpus == 1 else "mp"
|
|
|
|
if isinstance(distributed_executor_backend, type):
|
|
if not issubclass(distributed_executor_backend, DiffusionExecutor):
|
|
raise TypeError(
|
|
"distributed_executor_backend must be a subclass of "
|
|
f"DiffusionExecutor. Got {distributed_executor_backend}."
|
|
)
|
|
executor_class = distributed_executor_backend
|
|
elif distributed_executor_backend == "ray":
|
|
raise NotImplementedError("ray backend is not yet supported.")
|
|
elif distributed_executor_backend == "mp":
|
|
from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
|
|
|
|
executor_class = MultiprocDiffusionExecutor
|
|
elif distributed_executor_backend == "uni":
|
|
from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor
|
|
|
|
executor_class = UniProcDiffusionExecutor
|
|
elif distributed_executor_backend == "external_launcher":
|
|
raise NotImplementedError("external_launcher backend is not yet supported.")
|
|
elif isinstance(distributed_executor_backend, str):
|
|
try:
|
|
executor_class = resolve_obj_by_qualname(distributed_executor_backend)
|
|
except (ImportError, ValueError) as e:
|
|
raise ValueError(
|
|
f"Failed to load executor backend '{distributed_executor_backend}'. "
|
|
f"Ensure it is a valid python path. Error: {e}"
|
|
) from e
|
|
|
|
if not issubclass(executor_class, DiffusionExecutor):
|
|
raise TypeError(
|
|
f"distributed_executor_backend must be a subclass of DiffusionExecutor. Got {executor_class}."
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown distributed executor backend: {distributed_executor_backend}")
|
|
return executor_class
|
|
|
|
def __init__(self, od_config: OmniDiffusionConfig) -> None:
|
|
self.od_config = od_config
|
|
self._init_executor()
|
|
|
|
@abstractmethod
|
|
def _init_executor(self) -> None:
|
|
"""Initialize the executor (e.g., launch workers, setup IPC)."""
|
|
pass
|
|
|
|
@property
|
|
@abstractmethod
|
|
def is_dead(self) -> bool:
|
|
"""Whether the executor is shut down or has failed fatally."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
|
|
"""Execute request-mode work from a scheduler output."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def execute_batch(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
|
|
"""Execute request-mode work through the request-batch path."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
|
|
"""Execute step-mode work from a scheduler output."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def collective_rpc(
|
|
self,
|
|
method: str,
|
|
timeout: float | None = None,
|
|
args: tuple = (),
|
|
kwargs: dict | None = None,
|
|
unique_reply_rank: int | None = None,
|
|
exec_all_ranks: bool = False,
|
|
) -> Any:
|
|
"""Execute a method on workers."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def check_health(self) -> None:
|
|
"""Check if the executor and workers are healthy."""
|
|
pass
|
|
|
|
def register_failure_callback(self, callback: Callable[[], None]) -> None:
|
|
"""Register a callback invoked when the executor fatally fails.
|
|
|
|
Executors without a background failure monitor can keep the default
|
|
no-op implementation.
|
|
"""
|
|
return None
|
|
|
|
def wait_output_ready(self, async_output_id: str) -> Future[DiffusionOutput]:
|
|
"""Resolve deferred output; only asynchronous executors implement this."""
|
|
raise NotImplementedError(f"{type(self).__name__} does not support asynchronous output")
|
|
|
|
def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
|
|
"""Collect rank-local native specs after every Worker loads its model."""
|
|
|
|
result = self.collective_rpc(
|
|
"get_kv_cache_specs",
|
|
unique_reply_rank=0,
|
|
exec_all_ranks=True,
|
|
)
|
|
if not isinstance(result, list):
|
|
raise TypeError(f"get_kv_cache_specs returned {type(result).__name__}, expected list")
|
|
return result
|
|
|
|
def determine_available_kv_memory(self, profile_requests: list[OmniDiffusionRequest]) -> list[int]:
|
|
"""Profile and collect the KV memory budget on every Worker rank."""
|
|
|
|
result = self.collective_rpc(
|
|
"determine_available_kv_memory",
|
|
args=(profile_requests,),
|
|
unique_reply_rank=0,
|
|
exec_all_ranks=True,
|
|
)
|
|
if not isinstance(result, list) or not all(isinstance(value, int) for value in result):
|
|
raise TypeError("determine_available_kv_memory must return list[int]")
|
|
return result
|
|
|
|
def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig], resolved_max_model_len: int) -> None:
|
|
"""Send rank-local configs and the resolved model length to all Workers."""
|
|
|
|
# The default control-plane RPC mode executes on every rank and has
|
|
# rank 0 return the gathered rank statuses, so failures on nonzero
|
|
# ranks are not silently dropped.
|
|
self.collective_rpc("set_kv_cache_configs", args=(kv_cache_configs, resolved_max_model_len))
|
|
|
|
def remove_diffusion_kv_requests(self, request_ids: list[str | tuple[str, int]]) -> None:
|
|
"""Clear request rows on every Worker after Scheduler retirement."""
|
|
|
|
unique_request_ids = list(dict.fromkeys(request_ids))
|
|
if not unique_request_ids:
|
|
return
|
|
self.collective_rpc(
|
|
"remove_diffusion_kv_requests",
|
|
args=(unique_request_ids,),
|
|
)
|
|
|
|
def prepare_kv_for_forward(self, scheduler_output: DiffusionSchedulerOutput) -> KVConnectorOutput | None:
|
|
if scheduler_output.kv_prefetch_connector_metadata is not None:
|
|
current = replace(
|
|
scheduler_output,
|
|
kv_transfer_request_ids=scheduler_output.kv_transfer_request_ids
|
|
- scheduler_output.kv_prefetch_request_ids,
|
|
kv_prefetch_connector_metadata=None,
|
|
kv_prefetch_request_ids=set(),
|
|
)
|
|
# First complete the current request on every rank. Only then
|
|
# submit B, preventing Mooncake from coalescing its bytes with A.
|
|
self.prepare_kv_for_forward(current)
|
|
return self.prepare_kv_for_forward(
|
|
replace(
|
|
current,
|
|
kv_connector_metadata=scheduler_output.kv_prefetch_connector_metadata,
|
|
kv_transfer_request_ids=scheduler_output.kv_prefetch_request_ids,
|
|
kv_required_request_ids=set(),
|
|
kv_finished_request_ids=set(),
|
|
)
|
|
)
|
|
if (
|
|
scheduler_output.kv_connector_metadata is None
|
|
and not scheduler_output.kv_required_request_ids
|
|
and not scheduler_output.kv_poll_only
|
|
):
|
|
return None
|
|
transfer_output = replace(
|
|
scheduler_output,
|
|
scheduled_new_reqs=[],
|
|
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
|
kv_prefetch_job=None,
|
|
)
|
|
outputs: list[KVConnectorOutput] = self.collective_rpc(
|
|
"prepare_kv_for_forward",
|
|
args=(transfer_output,),
|
|
unique_reply_rank=0,
|
|
exec_all_ranks=True,
|
|
)
|
|
if len(outputs) != self.od_config.num_gpus or any(output.invalid_block_ids for output in outputs):
|
|
# Missing ranks / invalid pages cannot establish safe ownership.
|
|
raise RuntimeError("Diffusion KV receive failed on one or more ranks")
|
|
if scheduler_output.kv_required_request_ids is not None:
|
|
# Prefetch workers retain completion events until retirement.
|
|
# Their cumulative snapshots can be intersected directly.
|
|
finished = set.intersection(*(set(output.finished_recving or ()) for output in outputs))
|
|
if not scheduler_output.kv_required_request_ids.issubset(finished):
|
|
raise RuntimeError("Required diffusion KV receive did not complete on every rank")
|
|
outputs[0].finished_recving = finished
|
|
return outputs[0]
|
|
completed = getattr(self, "_kv_receive_completed_ranks", {})
|
|
for rank, output in enumerate(outputs):
|
|
for request_id in output.finished_recving or ():
|
|
completed.setdefault(request_id, set()).add(rank)
|
|
finished = {request_id for request_id, ranks in completed.items() if len(ranks) == len(outputs)}
|
|
for request_id in finished:
|
|
del completed[request_id]
|
|
self._kv_receive_completed_ranks = completed
|
|
outputs[0].finished_recving = finished
|
|
return outputs[0]
|
|
|
|
@abstractmethod
|
|
def shutdown(self) -> None:
|
|
"""Shutdown the executor and release resources."""
|
|
pass
|