import pickle import time from typing import Any, Optional import zmq import zmq.asyncio from sglang.multimodal_gen.runtime.ipc_array import materialize_file_refs from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.request_logger import ( DiffusionRequestLogger, ) logger = init_logger(__name__) async def run_zeromq_broker(server_args: ServerArgs): """ This function runs as a background task in the FastAPI process. It listens for TCP requests from offline clients (e.g., DiffGenerator). """ ctx = zmq.asyncio.Context() socket = ctx.socket(zmq.REP) broker_endpoint = f"tcp://127.0.0.1:{server_args.broker_port}" socket.bind(broker_endpoint) logger.info(f"ZMQ Broker is listening for offline jobs on {broker_endpoint}") while True: try: # 1. Receive a request from an offline client payload = await socket.recv() request_batch = pickle.loads(payload) logger.info("Broker received an offline job from a client.") # 2. Forward the request to the main Scheduler via the shared client response_batch = await async_scheduler_client.forward(request_batch) # 3. Send the Scheduler's reply back to the offline client await socket.send(pickle.dumps(response_batch)) except Exception as e: logger.error(f"Error in ZMQ Broker: {e}", exc_info=True) # A reply must be sent to prevent the client from hanging try: await socket.send(pickle.dumps({"status": "error", "message": str(e)})) except Exception: pass class SchedulerClient: """ A synchronous, singleton client for communicating with the Scheduler service. Designed for use in DiffGenerator, where synchronous usage is preferred """ def __init__(self): self.context = None self.scheduler_socket = None self.server_args = None self.request_logger: Optional[DiffusionRequestLogger] = None def initialize(self, server_args: ServerArgs): if self.context is not None and not self.context.closed: logger.warning("SchedulerClient is already initialized. Re-initializing.") self.close() self.server_args = server_args self.request_logger = DiffusionRequestLogger.from_server_args(server_args) self.context = zmq.Context() self.scheduler_socket = self.context.socket(zmq.REQ) # Set socket options for the main communication socket self.scheduler_socket.setsockopt(zmq.LINGER, 0) # 100 minute timeout for generation self.scheduler_socket.setsockopt(zmq.RCVTIMEO, 6000000) scheduler_endpoint = self.server_args.scheduler_endpoint self.scheduler_socket.connect(scheduler_endpoint) logger.debug( f"SchedulerClient connected to backend scheduler at {scheduler_endpoint}" ) def forward(self, batch: Any, timeout_ms: int | None = None) -> Any: """Sends a batch or request to the scheduler and waits for the response.""" self.request_logger.log_received_request(batch) previous_timeout_ms = None if timeout_ms is not None: previous_timeout_ms = self.scheduler_socket.getsockopt(zmq.RCVTIMEO) self.scheduler_socket.setsockopt(zmq.RCVTIMEO, timeout_ms) try: self.scheduler_socket.send_pyobj(batch) output_batch = self.scheduler_socket.recv_pyobj() _materialize_output_batch_file_refs(output_batch) self.request_logger.log_finished_request(batch, output_batch) return output_batch except zmq.error.Again: logger.error("Timeout waiting for response from scheduler.") raise TimeoutError("Scheduler did not respond in time.") finally: if previous_timeout_ms is not None and self.scheduler_socket is not None: self.scheduler_socket.setsockopt(zmq.RCVTIMEO, previous_timeout_ms) def ping(self) -> bool: """ Checks if the scheduler server is alive using a temporary socket. """ if self.context is None or self.context.closed: logger.error("Cannot ping: client is not initialized.") return False ping_socket = self.context.socket(zmq.REQ) ping_socket.setsockopt(zmq.LINGER, 0) ping_socket.setsockopt(zmq.RCVTIMEO, 2000) # 2-second timeout for pings endpoint = self.server_args.scheduler_endpoint try: ping_socket.connect(endpoint) ping_socket.send_pyobj({"method": "ping"}) ping_socket.recv_pyobj() return True except zmq.error.Again: return False finally: ping_socket.close() def close(self): """Closes the socket and terminates the context.""" if self.scheduler_socket: self.scheduler_socket.close() self.scheduler_socket = None if self.context: self.context.term() self.context = None class AsyncSchedulerClient: """ An asynchronous, singleton client for communicating with the Scheduler service. Designed for use in asynchronous environments like FastAPI entrypoints. To support high concurrency, it creates a new REQ socket for each request rather than sharing a single one (which would cause ZMQ state errors). """ def __init__(self): self.context = None self.server_args = None self.request_logger: Optional[DiffusionRequestLogger] = None def initialize(self, server_args: ServerArgs): if self.context is not None and not self.context.closed: logger.warning( "AsyncSchedulerClient is already initialized. Re-initializing." ) self.close() self.server_args = server_args self.request_logger = DiffusionRequestLogger.from_server_args(server_args) self.context = zmq.asyncio.Context() logger.debug("AsyncSchedulerClient initialized with zmq.asyncio.Context") async def forward(self, batch: Any) -> Any: """Sends a batch or request to the scheduler and waits for the response.""" self.request_logger.log_received_request(batch) if self.context is None: raise RuntimeError( "AsyncSchedulerClient is not initialized. Call initialize() first." ) # Create a temporary REQ socket for this request to allow concurrency socket = self.context.socket(zmq.REQ) socket.setsockopt(zmq.LINGER, 0) # 100 minute timeout socket.setsockopt(zmq.RCVTIMEO, 6000000) endpoint = self.server_args.scheduler_endpoint socket.connect(endpoint) try: await socket.send(pickle.dumps(batch)) payload = await socket.recv() output_batch = pickle.loads(payload) _materialize_output_batch_file_refs(output_batch) self.request_logger.log_finished_request(batch, output_batch) return output_batch except zmq.error.Again: logger.error("Timeout waiting for response from scheduler.") raise TimeoutError("Scheduler did not respond in time.") finally: socket.close() async def ping(self) -> bool: """ Checks if the scheduler server is alive using a temporary socket. """ if self.context is None or self.context.closed: logger.error("Cannot ping: client is not initialized.") return False ping_socket = self.context.socket(zmq.REQ) ping_socket.setsockopt(zmq.LINGER, 0) ping_socket.setsockopt(zmq.RCVTIMEO, 2000) endpoint = self.server_args.scheduler_endpoint try: ping_socket.connect(endpoint) await ping_socket.send(pickle.dumps({"method": "ping"})) await ping_socket.recv() return True except zmq.error.Again: return False finally: ping_socket.close() def close(self): """Closes the socket and terminates the context.""" if self.context: self.context.term() self.context = None # Singleton instances for easy access async_scheduler_client = AsyncSchedulerClient() sync_scheduler_client = SchedulerClient() def _materialize_output_batch_file_refs(output_batch: Any) -> None: if not isinstance(output_batch, OutputBatch): return start_time = time.perf_counter() output_batch.output = materialize_file_refs(output_batch.output) if output_batch.metrics is not None: output_batch.metrics.record_stage( "SchedulerClient.materialize_file_refs", time.perf_counter() - start_time, )