72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import zmq
|
|
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class Scheduler:
|
|
_instance = None
|
|
|
|
def __new__(cls, *args, **kwargs):
|
|
if not cls._instance:
|
|
cls._instance = super().__new__(cls)
|
|
return cls._instance
|
|
|
|
def initialize(self, od_config: OmniDiffusionConfig):
|
|
if hasattr(self, "context") and not self.context.closed:
|
|
logger.warning("SyncSchedulerClient is already initialized. Re-initializing.")
|
|
self.close()
|
|
|
|
self.od_config = od_config
|
|
self.context = zmq.Context() # Standard synchronous context
|
|
|
|
# Initialize MessageQueue for broadcasting requests
|
|
# Assuming all readers are local for now as per current launch_engine implementation
|
|
self.mq = MessageQueue(
|
|
n_reader=od_config.num_gpus,
|
|
n_local_reader=od_config.num_gpus,
|
|
local_reader_ranks=list(range(od_config.num_gpus)),
|
|
)
|
|
|
|
self.result_mq = None
|
|
|
|
def initialize_result_queue(self, handle):
|
|
# Initialize MessageQueue for receiving results
|
|
# We act as rank 0 reader for this queue
|
|
self.result_mq = MessageQueue.create_from_handle(handle, rank=0)
|
|
logger.info("SyncScheduler initialized result MessageQueue")
|
|
|
|
def get_broadcast_handle(self):
|
|
return self.mq.export_handle()
|
|
|
|
def add_req(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
|
|
"""Sends a request to the scheduler and waits for the response."""
|
|
try:
|
|
# Broadcast request to all workers
|
|
self.mq.enqueue(requests)
|
|
# Wait for result from Rank 0 (or whoever sends it)
|
|
if self.result_mq is None:
|
|
raise RuntimeError("Result queue not initialized")
|
|
|
|
output = self.result_mq.dequeue()
|
|
return output
|
|
except zmq.error.Again:
|
|
logger.error("Timeout waiting for response from scheduler.")
|
|
raise TimeoutError("Scheduler did not respond in time.")
|
|
|
|
def close(self):
|
|
"""Closes the socket and terminates the context."""
|
|
if hasattr(self, "context"):
|
|
self.context.term()
|
|
|
|
|
|
# Singleton instance for easy access
|
|
scheduler = Scheduler()
|