chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2025-12-21 00:39:54 +08:00
committed by GitHub
parent bee8ac5b88
commit bc18cb8631
16 changed files with 310 additions and 305 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/multimodal_gen/apps/webui/main.py
View File

@@ -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):

34
python/sglang/multimodal_gen/configs/sample/sampling_params.py
View File

@@ -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:

10
python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
View File

@@ -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):

10
python/sglang/multimodal_gen/runtime/entrypoints/http_server.py
View File

@@ -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.")

4
python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py
View File

@@ -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:

9
python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py
View File

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

8
python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py
View File

@@ -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(

4
python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py
View File

@@ -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())},

3
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
View File

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

103
python/sglang/multimodal_gen/runtime/managers/scheduler.py
View File

@@ -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:

3
python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
View File

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

38
python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py
View File

@@ -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:

285
python/sglang/multimodal_gen/runtime/scheduler_client.py
View File

@@ -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()

92
python/sglang/multimodal_gen/runtime/sync_scheduler_client.py
View File

@@ -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()

8
python/sglang/multimodal_gen/runtime/utils/common.py
View File

@@ -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:

2
python/sglang/multimodal_gen/runtime/utils/perf_logger.py
View File

@@ -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: