[diffusion] refactor: change zmq socket type to router for scheduler (#15479)

python/sglang/multimodal_gen/apps/webui/main.py

@@ -9,8 +9,8 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
     post_process_sample,
     prepare_request,
 )
+from sglang.multimodal_gen.runtime.scheduler_client import sync_scheduler_client
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.sync_scheduler_client import sync_scheduler_client
 
 
 def add_webui_args(parser: argparse.ArgumentParser):

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -17,7 +17,7 @@ from enum import Enum, auto
 from typing import Any
 
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-from sglang.multimodal_gen.utils import StoreBoolean, align_to
+from sglang.multimodal_gen.utils import StoreBoolean
 
 logger = init_logger(__name__)
 
@@ -340,7 +340,6 @@ class SamplingParams:
                 self.num_frames = new_num_frames
 
         self._set_output_file_name()
-        self.log(pipeline_config=server_args.pipeline_config)
 
     @classmethod
     def from_pretrained(cls, model_path: str, **kwargs) -> "SamplingParams":
@@ -680,37 +679,6 @@ class SamplingParams:
     def output_file_path(self):
         return os.path.join(self.output_path, self.output_file_name)
 
-    def log(self, pipeline_config):
-        # TODO: in some cases (e.g., TI2I), height and weight might be undecided at this moment
-        if self.height:
-            target_height = align_to(self.height, 16)
-        else:
-            target_height = -1
-        if self.width:
-            target_width = align_to(self.width, 16)
-        else:
-            target_width = -1
-
-        # Log sampling parameters
-        debug_str = f"""Sampling params:
-                       width: {target_width}
-                      height: {target_height}
-                  num_frames: {self.num_frames}
-                      prompt: {self.prompt}
-                  neg_prompt: {self.negative_prompt}
-                        seed: {self.seed}
-                 infer_steps: {self.num_inference_steps}
-      num_outputs_per_prompt: {self.num_outputs_per_prompt}
-              guidance_scale: {self.guidance_scale}
-     embedded_guidance_scale: {pipeline_config.embedded_cfg_scale}
-                    n_tokens: {self.n_tokens}
-                  flow_shift: {pipeline_config.flow_shift}
-                  image_path: {self.image_path}
-                 save_output: {self.save_output}
-            output_file_path: {self.output_file_path()}
-        """  # type: ignore[attr-defined]
-        logger.info(debug_str)
-
 
 @dataclass
 class CacheParams:

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -28,8 +28,8 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
 from sglang.multimodal_gen.runtime.launch_server import launch_server
 from sglang.multimodal_gen.runtime.pipelines_core import Req
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
+from sglang.multimodal_gen.runtime.scheduler_client import sync_scheduler_client
 from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
-from sglang.multimodal_gen.runtime.sync_scheduler_client import sync_scheduler_client
 from sglang.multimodal_gen.runtime.utils.logging_utils import (
     init_logger,
     log_batch_completion,
@@ -383,10 +383,12 @@ class DiffGenerator:
             lora_path=lora_path, lora_nickname=lora_nickname, merge_lora=merge_lora
         )
         return self.generate(
+            sampling_params_kwargs=dict(
                 prompt=prompt,
                 sampling_params=sampling_params,
                 **kwargs,
             )
+        )
 
     def shutdown(self):
         """

python/sglang/multimodal_gen/runtime/entrypoints/http_server.py

@@ -16,7 +16,7 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
     post_process_sample,
     prepare_request,
 )
-from sglang.multimodal_gen.runtime.scheduler_client import scheduler_client
+from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
 from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args
 from sglang.srt.managers.io_struct import VertexGenerateReqInput
 
@@ -27,13 +27,13 @@ VERTEX_ROUTE = os.environ.get("AIP_PREDICT_ROUTE", "/vertex_generate")
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     from sglang.multimodal_gen.runtime.scheduler_client import (
+        async_scheduler_client,
         run_zeromq_broker,
-        scheduler_client,
     )
 
     # 1. Initialize the singleton client that connects to the backend Scheduler
     server_args = app.state.server_args
-    scheduler_client.initialize(server_args)
+    async_scheduler_client.initialize(server_args)
 
     # 2. Start the ZMQ Broker in the background to handle offline requests
     broker_task = asyncio.create_task(run_zeromq_broker(server_args))
@@ -43,7 +43,7 @@ async def lifespan(app: FastAPI):
     # On shutdown
     print("FastAPI app is shutting down...")
     broker_task.cancel()
-    scheduler_client.close()
+    async_scheduler_client.close()
 
 
 # Health router
@@ -105,7 +105,7 @@ def encode_video_to_base64(file_path: str):
 async def forward_to_scheduler(req_obj, sp):
     """Forwards request to scheduler and processes the result."""
     try:
-        response = await scheduler_client.forward(req_obj)
+        response = await async_scheduler_client.forward(req_obj)
         if response.output is None:
             raise RuntimeError("Model generation returned no output.")
 

python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py

@@ -7,7 +7,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
     SetLoraReq,
     UnmergeLoraWeightsReq,
 )
-from sglang.multimodal_gen.runtime.scheduler_client import scheduler_client
+from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 router = APIRouter(prefix="/v1")
@@ -16,7 +16,7 @@ logger = init_logger(__name__)
 
 async def _handle_lora_request(req: Any, success_msg: str, failure_msg: str):
     try:
-        response = await scheduler_client.forward(req)
+        response = await async_scheduler_client.forward(req)
         if isinstance(response, dict) and response.get("status") == "ok":
             return {"status": "ok", "message": success_msg}
         else:

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

@@ -27,7 +27,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
 )
 from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
-from sglang.multimodal_gen.runtime.scheduler_client import scheduler_client
+from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
 from sglang.multimodal_gen.runtime.server_args import get_global_server_args
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import shallow_asdict
@@ -55,7 +55,7 @@ def _build_sampling_params_from_request(
     size: Optional[str],
     output_format: Optional[str],
     background: Optional[str],
-    image_path: Optional[str] = None,
+    image_path: Optional[list[str]] = None,
     seed: Optional[int] = None,
     generator_device: Optional[str] = None,
     negative_prompt: Optional[str] = None,
@@ -120,8 +120,9 @@ async def generations(
         server_args=get_global_server_args(),
         sampling_params=sampling,
     )
+
     # Run synchronously for images and save to disk
-    save_file_path = await process_generation_batch(scheduler_client, batch)
+    save_file_path = await process_generation_batch(async_scheduler_client, batch)
 
     await IMAGE_STORE.upsert(
         request_id,
@@ -218,7 +219,7 @@ async def edits(
     )
     batch = _build_req_from_sampling(sampling)
 
-    save_file_path = await process_generation_batch(scheduler_client, batch)
+    save_file_path = await process_generation_batch(async_scheduler_client, batch)
 
     await IMAGE_STORE.upsert(
         request_id,

python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py

@@ -10,6 +10,7 @@ import httpx
 from fastapi import UploadFile
 
 from sglang.multimodal_gen.runtime.entrypoints.utils import post_process_sample
+from sglang.multimodal_gen.runtime.scheduler_client import AsyncSchedulerClient
 from sglang.multimodal_gen.runtime.utils.logging_utils import (
     init_logger,
     log_batch_completion,
@@ -156,7 +157,7 @@ async def _save_base64_image_to_path(base64_data: str, target_path: str) -> str:
 
 
 async def process_generation_batch(
-    scheduler_client,
+    scheduler_client: AsyncSchedulerClient,
     batch,
 ):
     total_start_time = time.perf_counter()
@@ -164,7 +165,10 @@ async def process_generation_batch(
         result = await scheduler_client.forward([batch])
 
         if result.output is None:
-            raise RuntimeError("Model generation returned no output.")
+            error_msg = getattr(result, "error", "Unknown error")
+            raise RuntimeError(
+                f"Model generation returned no output. Error from scheduler: {error_msg}"
+            )
 
         save_file_path = str(os.path.join(batch.output_path, batch.output_file_name))
         post_process_sample(

python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py

@@ -115,10 +115,10 @@ def _video_job_from_sampling(
 
 
 async def _dispatch_job_async(job_id: str, batch: Req) -> None:
-    from sglang.multimodal_gen.runtime.scheduler_client import scheduler_client
+    from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
 
     try:
-        await process_generation_batch(scheduler_client, batch)
+        await process_generation_batch(async_scheduler_client, batch)
         await VIDEO_STORE.update_fields(
             job_id,
             {"status": "completed", "progress": 100, "completed_at": int(time.time())},

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -108,6 +108,9 @@ class GPUWorker:
                 output_batch.timings.total_duration_ms = duration_ms
                 PerformanceLogger.log_request_summary(timings=output_batch.timings)
         except Exception as e:
+            logger.error(
+                f"Error executing request {req.request_id}: {e}", exc_info=True
+            )
             if output_batch is None:
                 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
                     OutputBatch,

python/sglang/multimodal_gen/runtime/managers/scheduler.py

@@ -1,6 +1,8 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
 
 # SPDX-License-Identifier: Apache-2.0
+import pickle
+from collections import deque
 from typing import Any, List
 
 import zmq
@@ -49,8 +51,9 @@ class Scheduler:
         self.context = zmq.Context(io_threads=2)
         endpoint = server_args.scheduler_endpoint()
         if gpu_id == 0:
+            # router allocates identify (envelope) for each connection
             self.receiver, actual_endpoint = get_zmq_socket(
-                self.context, zmq.REP, endpoint, True
+                self.context, zmq.ROUTER, endpoint, True
             )
             logger.info(f"Scheduler bind at endpoint: {actual_endpoint}")
         else:
@@ -73,8 +76,12 @@ class Scheduler:
             MergeLoraWeightsReq: self._handle_merge_lora,
             UnmergeLoraWeightsReq: self._handle_unmerge_lora,
             Req: self._handle_generation,
+            List[Req]: self._handle_generation,
         }
 
+        # FIFO, new reqs are appended
+        self.waiting_queue: deque[tuple[bytes, Req]] = deque()
+
     def _handle_set_lora(self, reqs: List[Any]):
         # TODO: return set status
         req = reqs[0]
@@ -94,20 +101,31 @@ class Scheduler:
     def _handle_generation(self, reqs: List[Req]):
         return self.worker.execute_forward(reqs)
 
-    def return_result(self, output_batch: OutputBatch):
+    def return_result(self, output_batch: OutputBatch, identity: bytes | None = None):
         """
         replies to client, only on rank 0
         """
-        if self.receiver is not None:
+        if self.receiver is not None and identity is not None:
-            self.receiver.send_pyobj(output_batch)
+            self.receiver.send_multipart([identity, b"", pickle.dumps(output_batch)])
 
-    def recv_reqs(self) -> List[Any]:
+    def get_next_batch_to_run(self) -> list[tuple[bytes, Req]] | None:
+        """pull a req from waiting_queue"""
+        if not self.waiting_queue:
+            return None
+
+        # pop the first (earliest)
+        item = self.waiting_queue.popleft()
+
+        return [item]
+
+    def recv_reqs(self) -> List[tuple[bytes, Any]]:
         """
         For non-main schedulers, reqs are broadcasted from main using broadcast_pyobj
         """
         if self.receiver is not None:
             try:
-                recv_reqs = self.receiver.recv_pyobj()
+                identity, _, payload = self.receiver.recv_multipart()
+                recv_reqs = pickle.loads(payload)
             except zmq.ZMQError:
                 # re-raise or handle appropriately to let the outer loop continue
                 raise
@@ -115,6 +133,9 @@ class Scheduler:
             # Ensure recv_reqs is a list
             if not isinstance(recv_reqs, list):
                 recv_reqs = [recv_reqs]
+
+            # Pack with identity for rank 0
+            recv_reqs = [(identity, req) for req in recv_reqs]
         else:
             recv_reqs = None
 
@@ -147,22 +168,22 @@ class Scheduler:
 
         return recv_reqs
 
-    # TODO: queueing, cancellation
     def event_loop(self) -> None:
         """
         The main event loop that listens for ZMQ requests.
         Handles abortion
         """
 
-        logger.info(
+        logger.debug(
             f"Rank 0 scheduler listening on tcp://*:{self.server_args.scheduler_port}"
         )
 
         while self._running:
-            reqs = None
             # 1: receive requests
             try:
-                reqs = self.recv_reqs()
+                new_reqs = self.recv_reqs()
+                # after processing input reqs
+                self.waiting_queue.extend(new_reqs)
             except Exception as e:
                 logger.error(
                     f"Error receiving requests in scheduler event loop: {e}",
@@ -171,9 +192,16 @@ class Scheduler:
                 continue
 
             # 2: execute, make sure a reply is always sent
-            try:
+            while self.waiting_queue:
-                first_req = reqs[0] if reqs else None
+                items = self.get_next_batch_to_run()
+                if not items:
+                    break
 
+                identities = [item[0] for item in items]
+                reqs = [item[1] for item in items]
+
+                try:
+                    first_req = reqs[0]
                     handler = self.request_handlers.get(type(first_req))
                     if handler:
                         output_batch = handler(reqs)
@@ -195,7 +223,8 @@ class Scheduler:
                     )
 
                 try:
-                self.return_result(output_batch)
+                    # TODO: Support sending back to multiple identities if batched
+                    self.return_result(output_batch, identities[0])
                 except zmq.ZMQError as e:
                     # Reply failed; log and keep loop alive to accept future requests
                     logger.error(f"ZMQ error sending reply: {e}")

python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py

@@ -363,10 +363,13 @@ class ComposedPipelineBase(ABC):
                 "LoRA adapter is set, but not effective. Please make sure the LoRA weights are merged"
             )
 
+        batch.log(server_args=server_args)
+
         # Execute each stage
         logger.info(
             "Running pipeline stages: %s",
             list(self._stage_name_mapping.keys()),
             main_process_only=True,
         )
+
         return self.executor.execute(self.stages, batch, server_args)

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -25,12 +25,17 @@ from sglang.multimodal_gen.configs.sample.teacache import (
     WanTeaCacheParams,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import align_to
 
 if TYPE_CHECKING:
 
     from sglang.multimodal_gen.runtime.utils.perf_logger import RequestTimings
 
 
+logger = init_logger(__name__)
+
+
 @dataclass
 class Req:
     """
@@ -202,7 +207,7 @@ class Req:
         batch_size *= self.num_outputs_per_prompt
         return batch_size
 
-    def output_file_path(self, num_outputs, output_idx):
+    def output_file_path(self, num_outputs=1, output_idx=None):
         output_file_name = self.output_file_name
         if num_outputs > 1 and output_file_name:
             base, ext = os.path.splitext(output_file_name)
@@ -230,6 +235,37 @@ class Req:
     def __str__(self):
         return pprint.pformat(asdict(self), indent=2, width=120)
 
+    def log(self, server_args: ServerArgs):
+        # TODO: in some cases (e.g., TI2I), height and weight might be undecided at this moment
+        if self.height:
+            target_height = align_to(self.height, 16)
+        else:
+            target_height = -1
+        if self.width:
+            target_width = align_to(self.width, 16)
+        else:
+            target_width = -1
+
+        # Log sampling parameters
+        debug_str = f"""Sampling params:
+                       width: {target_width}
+                      height: {target_height}
+                  num_frames: {self.num_frames}
+                      prompt: {self.prompt}
+                  neg_prompt: {self.negative_prompt}
+                        seed: {self.seed}
+                 infer_steps: {self.num_inference_steps}
+      num_outputs_per_prompt: {self.num_outputs_per_prompt}
+              guidance_scale: {self.guidance_scale}
+     embedded_guidance_scale: {server_args.pipeline_config.embedded_cfg_scale}
+                    n_tokens: {self.n_tokens}
+                  flow_shift: {server_args.pipeline_config.flow_shift}
+                  image_path: {self.image_path}
+                 save_output: {self.save_output}
+            output_file_path: {self.output_file_path()}
+        """  # type: ignore[attr-defined]
+        logger.info(debug_str)
+
 
 @dataclass
 class OutputBatch:

python/sglang/multimodal_gen/runtime/scheduler_client.py

@@ -1,6 +1,4 @@
-# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+import pickle
-
-import asyncio
 from typing import Any
 
 import zmq
@@ -12,113 +10,6 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 logger = init_logger(__name__)
 
 
-# Using a singleton pattern to hold the ZMQ context and the socket connected to the scheduler
-class SchedulerClient:
-    """
-    A gateway for Scheduler, forwarding the ForwardBatch from http endpoints (or somewhere else) to background scheduler, with TCP socket
-    """
-
-    _instance = None
-
-    def __new__(cls, *args, **kwargs):
-        if not cls._instance:
-            cls._instance = super(SchedulerClient, cls).__new__(cls)
-        return cls._instance
-
-    def __init__(self, *args, **kwargs):
-        # Ensure the initialization runs only once for the singleton instance
-        if getattr(self, "_init_done", False):
-            return
-        # Queue + worker to strictly serialize ZeroMQ REQ/REP interactions
-        self._request_queue = asyncio.Queue()
-        self._worker_task = None
-        self._closing = False
-        self._init_done = True
-
-    def initialize(self, server_args: ServerArgs):
-        self.server_args = server_args
-        self.context = zmq.asyncio.Context()
-        # This is the REQ socket used to connect to the backend Scheduler
-        self.scheduler_socket = self.context.socket(zmq.REQ)
-        scheduler_endpoint = server_args.scheduler_endpoint()
-        self.scheduler_socket.connect(scheduler_endpoint)
-        logger.info(
-            f"Scheduler client connected to backend scheduler at {scheduler_endpoint}"
-        )
-        # Worker will be lazily started on the first forward call to ensure a running loop exists
-
-    async def forward(self, batch: Any) -> Any:
-        """Enqueue a request to the backend Scheduler and await the reply."""
-        if self._closing:
-            raise RuntimeError(
-                "SchedulerClient is closing; cannot forward new requests"
-            )
-
-        await self._ensure_worker_started()
-
-        loop = asyncio.get_running_loop()
-        future = loop.create_future()
-        await self._request_queue.put((batch, future))
-        return await future
-
-    async def _ensure_worker_started(self):
-        # Start the worker only once and only when an event loop is running
-        if self._worker_task is None or self._worker_task.done():
-            self._worker_task = asyncio.create_task(self._worker_loop())
-
-    async def _worker_loop(self):
-        while True:
-            try:
-                item = await self._request_queue.get()
-                try:
-                    batch, future = item
-                except Exception:
-                    # Malformed queue item; skip
-                    self._request_queue.task_done()
-                    continue
-
-                try:
-                    await self.scheduler_socket.send_pyobj(batch)
-                    response = await self.scheduler_socket.recv_pyobj()
-                    if not future.done():
-                        future.set_result(response)
-                except Exception as e:
-                    if not future.done():
-                        future.set_exception(e)
-                finally:
-                    self._request_queue.task_done()
-            except asyncio.CancelledError:
-                # Drain remaining items with cancellation error to avoid hanging waiters
-                while True:
-                    try:
-                        batch, future = self._request_queue.get_nowait()
-                    except asyncio.QueueEmpty:
-                        break
-                    try:
-                        if not future.done():
-                            future.set_exception(asyncio.CancelledError())
-                    finally:
-                        self._request_queue.task_done()
-                raise
-
-    def close(self):
-        self._closing = True
-        # Cancel worker if running
-        if self._worker_task is not None:
-            self._worker_task.cancel()
-        try:
-            self.scheduler_socket.close()
-        finally:
-            try:
-                self.context.term()
-            except Exception:
-                pass
-
-
-# Singleton instance
-scheduler_client = SchedulerClient()
-
-
 async def run_zeromq_broker(server_args: ServerArgs):
     """
     This function runs as a background task in the FastAPI process.
@@ -134,16 +25,182 @@ async def run_zeromq_broker(server_args: ServerArgs):
     while True:
         try:
             # 1. Receive a request from an offline client
-            request_batch = await socket.recv_pyobj()
+            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 scheduler_client.forward(request_batch)
+            response_batch = await async_scheduler_client.forward(request_batch)
 
             # 3. Send the Scheduler's reply back to the offline client
-            await socket.send_pyobj(response_batch)
+            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
-            await socket.send_pyobj({"status": "error", "message": str(e)})
+            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
+
+    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.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) -> Any:
+        """Sends a batch or request to the scheduler and waits for the response."""
+        try:
+            self.scheduler_socket.send_pyobj(batch)
+            output_batch = self.scheduler_socket.recv_pyobj()
+            return output_batch
+        except zmq.error.Again:
+            logger.error("Timeout waiting for response from scheduler.")
+            raise TimeoutError("Scheduler did not respond in time.")
+
+    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
+
+    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.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."""
+        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()
+            return pickle.loads(payload)
+        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()

python/sglang/multimodal_gen/runtime/sync_scheduler_client.py

@@ -1,92 +0,0 @@
-# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
-from typing import Any
-
-import zmq
-
-from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
-
-
-class SyncSchedulerClient:
-    """
-    A synchronous, singleton client for communicating with the Scheduler service.
-    Designed for use in synchronous environments like the DiffGenerator or standalone scripts.
-    """
-
-    _instance = None
-
-    def __new__(cls, *args, **kwargs):
-        if not cls._instance:
-            cls._instance = super(SyncSchedulerClient, cls).__new__(cls)
-        return cls._instance
-
-    def initialize(self, server_args: ServerArgs):
-        if hasattr(self, "context") and not self.context.closed:
-            logger.warning(
-                "SyncSchedulerClient is already initialized. Re-initializing."
-            )
-            self.close()
-
-        self.server_args = server_args
-        self.context = zmq.Context()  # Standard synchronous context
-        self.scheduler_socket = self.context.socket(zmq.REQ)
-
-        # Set socket options for the main communication socket
-        self.scheduler_socket.setsockopt(zmq.LINGER, 0)
-        self.scheduler_socket.setsockopt(
-            zmq.RCVTIMEO, 6000000
-        )  # 10 minute timeout for generation
-
-        scheduler_endpoint = self.server_args.scheduler_endpoint()
-        self.scheduler_socket.connect(scheduler_endpoint)
-        logger.debug(
-            f"SyncSchedulerClient connected to backend scheduler at {scheduler_endpoint}"
-        )
-
-    def forward(self, batch: Any) -> Any:
-        """Sends a batch or request to the scheduler and waits for the response."""
-        try:
-            self.scheduler_socket.send_pyobj(batch)
-            output_batch = self.scheduler_socket.recv_pyobj()
-            return output_batch
-        except zmq.error.Again:
-            logger.error("Timeout waiting for response from scheduler.")
-            raise TimeoutError("Scheduler did not respond in time.")
-
-    def ping(self) -> bool:
-        """
-        Checks if the scheduler server is alive using a temporary socket.
-        This avoids interfering with the state of the main REQ/REP socket.
-        """
-        if not hasattr(self, "context") 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 hasattr(self, "scheduler_socket"):
-            self.scheduler_socket.close()
-        if hasattr(self, "context"):
-            self.context.term()
-
-
-# Singleton instance for easy access
-sync_scheduler_client = SyncSchedulerClient()

python/sglang/multimodal_gen/runtime/utils/common.py

@@ -170,13 +170,7 @@ def get_zmq_socket(
         set_send_opt()
     elif socket_type == zmq.PULL:
         set_recv_opt()
-    elif socket_type == zmq.DEALER:
+    elif socket_type in [zmq.DEALER, zmq.REQ, zmq.REP, zmq.ROUTER]:
-        set_send_opt()
-        set_recv_opt()
-    elif socket_type == zmq.REQ:
-        set_send_opt()
-        set_recv_opt()
-    elif socket_type == zmq.REP:
         set_send_opt()
         set_recv_opt()
     else:

python/sglang/multimodal_gen/runtime/utils/perf_logger.py

@@ -168,7 +168,7 @@ class StageProfiler:
 
         if self.simple_log:
             self.logger.info(
-                f"[{self.stage_name}] finished in {execution_time_s:.4f} seconds"
+                f"[{self.stage_name}] finished in {execution_time_s:.4f} seconds",
             )
 
         if self.metrics_enabled and self.timings: