sgl-project--sglang/python/sglang/multimodal_gen/runtime/scheduler_client.py

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,
        )