diff --git a/python/sglang/multimodal_gen/apps/webui/main.py b/python/sglang/multimodal_gen/apps/webui/main.py index be60dcfb9..1245a2683 100644 --- a/python/sglang/multimodal_gen/apps/webui/main.py +++ b/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): diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 620382fed..81a442a90 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 88784a58c..23d05e5dd 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/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,9 +383,11 @@ class DiffGenerator: lora_path=lora_path, lora_nickname=lora_nickname, merge_lora=merge_lora ) return self.generate( - prompt=prompt, - sampling_params=sampling_params, - **kwargs, + sampling_params_kwargs=dict( + prompt=prompt, + sampling_params=sampling_params, + **kwargs, + ) ) def shutdown(self): diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index d84e53e82..cfa80f9f1 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/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.") diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py index 2f1501672..676520aa3 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index f748120de..ab3af9988 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index 2c1e26e8f..08d658300 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index edca67dc6..d467d0d45 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/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())}, diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index a580fbeef..a3c414040 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index e2fe164fb..7fb9a2e81 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/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: - self.receiver.send_pyobj(output_batch) + if self.receiver is not None and identity is not None: + 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,35 +192,43 @@ class Scheduler: continue # 2: execute, make sure a reply is always sent - try: - first_req = reqs[0] if reqs else None + while self.waiting_queue: + items = self.get_next_batch_to_run() + if not items: + break - handler = self.request_handlers.get(type(first_req)) - if handler: - output_batch = handler(reqs) - else: - output_batch = { - "status": "error", - "message": f"Unknown request type: {type(first_req)}", - } - except Exception as e: - logger.error( - f"Error executing request in scheduler event loop: {e}", - exc_info=True, - ) - # Determine appropriate error response format - output_batch = ( - OutputBatch(error=str(e)) - if reqs and isinstance(reqs[0], Req) - else {"status": "error", "message": str(e)} - ) + identities = [item[0] for item in items] + reqs = [item[1] for item in items] - try: - self.return_result(output_batch) - except zmq.ZMQError as e: - # Reply failed; log and keep loop alive to accept future requests - logger.error(f"ZMQ error sending reply: {e}") - continue + try: + first_req = reqs[0] + handler = self.request_handlers.get(type(first_req)) + if handler: + output_batch = handler(reqs) + else: + output_batch = { + "status": "error", + "message": f"Unknown request type: {type(first_req)}", + } + except Exception as e: + logger.error( + f"Error executing request in scheduler event loop: {e}", + exc_info=True, + ) + # Determine appropriate error response format + output_batch = ( + OutputBatch(error=str(e)) + if reqs and isinstance(reqs[0], Req) + else {"status": "error", "message": str(e)} + ) + + try: + # 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}") + continue logger.info("Scheduler event loop terminated.") if self.receiver is not None: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 71658ebee..4375138a7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 6783a9c76..df878ace4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/scheduler_client.py b/python/sglang/multimodal_gen/runtime/scheduler_client.py index c0c0a00a1..5bfc62533 100644 --- a/python/sglang/multimodal_gen/runtime/scheduler_client.py +++ b/python/sglang/multimodal_gen/runtime/scheduler_client.py @@ -1,6 +1,4 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -import asyncio +import pickle 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() diff --git a/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py b/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py deleted file mode 100644 index 3d32576fb..000000000 --- a/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py +++ /dev/null @@ -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() diff --git a/python/sglang/multimodal_gen/runtime/utils/common.py b/python/sglang/multimodal_gen/runtime/utils/common.py index 70bef8dd6..156800491 100644 --- a/python/sglang/multimodal_gen/runtime/utils/common.py +++ b/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: - set_send_opt() - set_recv_opt() - elif socket_type == zmq.REQ: - set_send_opt() - set_recv_opt() - elif socket_type == zmq.REP: + elif socket_type in [zmq.DEALER, zmq.REQ, zmq.REP, zmq.ROUTER]: set_send_opt() set_recv_opt() else: diff --git a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py index 30de35fdf..2afa203fc 100644 --- a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py +++ b/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: