[diffusion] feat: data-parallel serving (--dp-size) (#33725)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -137,6 +137,10 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
):
|
||||
self.local_rank = local_rank
|
||||
self.rank = rank
|
||||
# the rank that materializes output and replies to the client: the
|
||||
# first rank of this DP replica, which is global rank 0 only at dp=1
|
||||
gpus_per_replica = max(1, server_args.num_gpus // (server_args.dp_size or 1))
|
||||
self.is_output_rank = rank % gpus_per_replica == 0
|
||||
self.master_port = master_port
|
||||
# FIXME: should we use tcp as distribute init method?
|
||||
self.server_args = server_args
|
||||
@@ -460,7 +464,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
output_batch = None
|
||||
forward_failed = False
|
||||
try:
|
||||
if self.rank == 0 and not current_platform.is_cpu():
|
||||
if self.is_output_rank and not current_platform.is_cpu():
|
||||
torch.get_device_module().reset_peak_memory_stats()
|
||||
|
||||
start_time = (
|
||||
@@ -474,7 +478,11 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
request_metrics = [
|
||||
item.metrics for item in log_reqs if item.metrics is not None
|
||||
]
|
||||
if self.rank == 0 and request_metrics and not current_platform.is_cpu():
|
||||
if (
|
||||
self.is_output_rank
|
||||
and request_metrics
|
||||
and not current_platform.is_cpu()
|
||||
):
|
||||
baseline_snapshot = capture_memory_snapshot()
|
||||
for metrics in request_metrics:
|
||||
metrics.record_memory_snapshot("before_forward", baseline_snapshot)
|
||||
@@ -501,13 +509,13 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
self._record_output_peak_memory(output_batch)
|
||||
|
||||
output_metrics = self._iter_output_metrics(output_batch)
|
||||
if self.rank == 0 and output_metrics and not current_platform.is_cpu():
|
||||
if self.is_output_rank and output_metrics and not current_platform.is_cpu():
|
||||
peak_snapshot = capture_memory_snapshot()
|
||||
for metrics in output_metrics:
|
||||
metrics.record_memory_snapshot("after_forward", peak_snapshot)
|
||||
|
||||
if (
|
||||
self.rank == 0
|
||||
self.is_output_rank
|
||||
and not req.suppress_logs
|
||||
and not current_platform.is_cpu()
|
||||
and logger.isEnabledFor(logging.DEBUG)
|
||||
@@ -581,7 +589,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
def _materialize_raw_frame_transport(
|
||||
self, output_batch: OutputBatch, req: Req
|
||||
) -> None:
|
||||
if self.rank != 0:
|
||||
if not self.is_output_rank:
|
||||
return
|
||||
if output_batch.output is not None:
|
||||
output_batch.raw_frame_content_type = RAW_RGB_CONTENT_TYPE
|
||||
@@ -603,7 +611,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
output_batch: OutputBatch,
|
||||
save_output_paths: Callable[[OutputBatch], None],
|
||||
) -> None:
|
||||
if self.rank == 0:
|
||||
if self.is_output_rank:
|
||||
save_output_paths(output_batch)
|
||||
output_batch.output = None
|
||||
output_batch.audio = None
|
||||
@@ -614,7 +622,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
) -> None:
|
||||
"""materialize the output from tensor to numpy frames for faster serialization"""
|
||||
if (
|
||||
self.rank != 0
|
||||
not self.is_output_rank
|
||||
or output_batch.output is None
|
||||
or not getattr(req, "return_frames", False)
|
||||
):
|
||||
@@ -678,7 +686,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
return np.asarray(materialized.frames)
|
||||
|
||||
def _record_output_peak_memory(self, output_batch: OutputBatch) -> None:
|
||||
if self.rank != 0 or current_platform.is_cpu():
|
||||
if not self.is_output_rank or current_platform.is_cpu():
|
||||
return
|
||||
peak_reserved_bytes = torch.get_device_module().max_memory_reserved()
|
||||
output_batch.peak_memory_mb = peak_reserved_bytes / (1024**2)
|
||||
@@ -691,7 +699,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
|
||||
def _save_output_paths(self, req: Req, output_batch: OutputBatch) -> None:
|
||||
"""save outputs to files"""
|
||||
if self.rank != 0 or output_batch.output is None:
|
||||
if not self.is_output_rank or output_batch.output is None:
|
||||
return
|
||||
|
||||
dynamic_output_paths = None
|
||||
@@ -742,7 +750,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
reqs: list[Req],
|
||||
output_batch: OutputBatch,
|
||||
) -> None:
|
||||
if self.rank != 0 or output_batch.output is None:
|
||||
if not self.is_output_rank or output_batch.output is None:
|
||||
return
|
||||
if len(output_batch.output) != len(reqs):
|
||||
raise RuntimeError(
|
||||
|
||||
@@ -104,15 +104,23 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag
|
||||
|
||||
set_global_server_args(server_args=server_args)
|
||||
|
||||
# Inter-process Communication
|
||||
# Each DP replica is a contiguous rank block (dp is the outermost
|
||||
# layout axis); its first rank binds the replica's ingress, and the
|
||||
# sp/cfg/tp broadcast relay in recv_reqs -- replica-internal by
|
||||
# construction -- fans requests out within the replica only.
|
||||
gpus_per_replica = max(1, server_args.num_gpus // server_args.dp_size)
|
||||
self.dp_replica = gpu_id // gpus_per_replica
|
||||
self.context = zmq.Context(io_threads=2)
|
||||
endpoint = server_args.scheduler_endpoint
|
||||
if gpu_id == 0:
|
||||
if gpu_id % gpus_per_replica == 0:
|
||||
endpoint = server_args.scheduler_endpoint_for(self.dp_replica)
|
||||
# router allocates identify (envelope) for each connection
|
||||
self.receiver, actual_endpoint = get_zmq_socket(
|
||||
self.context, zmq.ROUTER, endpoint, True
|
||||
)
|
||||
logger.info(f"Scheduler bind at endpoint: {actual_endpoint}")
|
||||
logger.info(
|
||||
f"Scheduler (dp replica {self.dp_replica}) bind at endpoint: "
|
||||
f"{actual_endpoint}"
|
||||
)
|
||||
else:
|
||||
self.receiver = None
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
@@ -1172,9 +1180,8 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag
|
||||
self._disagg_event_loop()
|
||||
return
|
||||
|
||||
logger.debug(
|
||||
f"Rank 0 scheduler listening on tcp://*:{self.server_args.scheduler_port}"
|
||||
)
|
||||
if self.receiver is not None:
|
||||
logger.debug("Driver scheduler of dp replica %d listening", self.dp_replica)
|
||||
|
||||
while self._running:
|
||||
# Update queue depth for metrics
|
||||
|
||||
@@ -1,11 +1,29 @@
|
||||
import itertools
|
||||
import pickle
|
||||
import time
|
||||
import zlib
|
||||
from typing import Any, Optional
|
||||
|
||||
import zmq
|
||||
import zmq.asyncio
|
||||
|
||||
from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import (
|
||||
GetWeightsChecksumReqInput,
|
||||
ReleaseMemoryOccupationReqInput,
|
||||
ResumeMemoryOccupationReqInput,
|
||||
UpdateWeightFromDiskReqInput,
|
||||
UpdateWeightFromTensorCheckerReqInput,
|
||||
UpdateWeightFromTensorReqInput,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.entrypoints.utils import (
|
||||
ListLorasReq,
|
||||
MergeLoraWeightsReq,
|
||||
SetLoraReq,
|
||||
ShutdownReq,
|
||||
UnmergeLoraWeightsReq,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.ipc_array import materialize_file_refs
|
||||
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.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
@@ -15,6 +33,22 @@ from sglang.multimodal_gen.runtime.utils.request_logger import (
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Control ops mutate replica state (weights, LoRA, memory, shutdown), so with
|
||||
# DP they must reach every replica rather than one.
|
||||
_CONTROL_REQ_TYPES = (
|
||||
SetLoraReq,
|
||||
MergeLoraWeightsReq,
|
||||
UnmergeLoraWeightsReq,
|
||||
ListLorasReq,
|
||||
ShutdownReq,
|
||||
UpdateWeightFromDiskReqInput,
|
||||
UpdateWeightFromTensorReqInput,
|
||||
UpdateWeightFromTensorCheckerReqInput,
|
||||
GetWeightsChecksumReqInput,
|
||||
ReleaseMemoryOccupationReqInput,
|
||||
ResumeMemoryOccupationReqInput,
|
||||
)
|
||||
|
||||
|
||||
async def run_zeromq_broker(server_args: ServerArgs):
|
||||
"""
|
||||
@@ -49,6 +83,34 @@ async def run_zeromq_broker(server_args: ServerArgs):
|
||||
pass
|
||||
|
||||
|
||||
def _session_key(batch: Any) -> str | None:
|
||||
"""Realtime sessions hold GPU state on one replica, so every request of a
|
||||
session must land on the same one."""
|
||||
reqs = batch if isinstance(batch, list) else [batch]
|
||||
for req in reqs:
|
||||
if isinstance(req, Req) and req.realtime_session_id:
|
||||
return req.realtime_session_id
|
||||
return None
|
||||
|
||||
|
||||
def _select_replica(batch: Any, dp_size: int, counter: "itertools.count") -> int:
|
||||
if dp_size <= 1:
|
||||
return 0
|
||||
session = _session_key(batch)
|
||||
if session is not None:
|
||||
return zlib.crc32(session.encode()) % dp_size
|
||||
return next(counter) % dp_size
|
||||
|
||||
|
||||
def _merge_fanout_results(results: list[Any]) -> Any:
|
||||
"""One reply for a control op sent to every replica: the first error wins,
|
||||
because "succeeded" must mean succeeded everywhere."""
|
||||
for result in results:
|
||||
if isinstance(result, OutputBatch) and result.error:
|
||||
return result
|
||||
return results[0]
|
||||
|
||||
|
||||
class SchedulerClient:
|
||||
"""
|
||||
A synchronous, singleton client for communicating with the Scheduler service.
|
||||
@@ -57,9 +119,9 @@ class SchedulerClient:
|
||||
|
||||
def __init__(self):
|
||||
self.context = None
|
||||
self.scheduler_socket = None
|
||||
self.server_args = None
|
||||
self.request_logger: Optional[DiffusionRequestLogger] = None
|
||||
self._replica_counter = itertools.count()
|
||||
|
||||
def initialize(self, server_args: ServerArgs):
|
||||
if self.context is not None and not self.context.closed:
|
||||
@@ -69,39 +131,38 @@ class SchedulerClient:
|
||||
self.server_args = server_args
|
||||
self.request_logger = DiffusionRequestLogger.from_server_args(server_args)
|
||||
self.context = zmq.Context()
|
||||
self.scheduler_socket = self.context.socket(zmq.REQ)
|
||||
|
||||
# Set socket options for the main communication socket
|
||||
self.scheduler_socket.setsockopt(zmq.LINGER, 0)
|
||||
|
||||
# 100 minute timeout for generation
|
||||
self.scheduler_socket.setsockopt(zmq.RCVTIMEO, 6000000)
|
||||
|
||||
scheduler_endpoint = self.server_args.scheduler_endpoint
|
||||
self.scheduler_socket.connect(scheduler_endpoint)
|
||||
logger.debug(
|
||||
f"SchedulerClient connected to backend scheduler at {scheduler_endpoint}"
|
||||
)
|
||||
|
||||
def forward(self, batch: Any, timeout_ms: int | None = None) -> Any:
|
||||
"""Sends a batch or request to the scheduler and waits for the response."""
|
||||
self.request_logger.log_received_request(batch)
|
||||
previous_timeout_ms = None
|
||||
if timeout_ms is not None:
|
||||
previous_timeout_ms = self.scheduler_socket.getsockopt(zmq.RCVTIMEO)
|
||||
self.scheduler_socket.setsockopt(zmq.RCVTIMEO, timeout_ms)
|
||||
return self._forward_routed(batch, timeout_ms)
|
||||
|
||||
def _forward_one(self, endpoint: str, batch: Any, timeout_ms: int | None) -> Any:
|
||||
socket = self.context.socket(zmq.REQ)
|
||||
socket.setsockopt(zmq.LINGER, 0)
|
||||
socket.setsockopt(zmq.RCVTIMEO, timeout_ms if timeout_ms else 6000000)
|
||||
try:
|
||||
self.scheduler_socket.send_pyobj(batch)
|
||||
output_batch = self.scheduler_socket.recv_pyobj()
|
||||
socket.connect(endpoint)
|
||||
socket.send_pyobj(batch)
|
||||
output_batch = socket.recv_pyobj()
|
||||
_materialize_output_batch_file_refs(output_batch)
|
||||
self.request_logger.log_finished_request(batch, output_batch)
|
||||
return output_batch
|
||||
except zmq.error.Again:
|
||||
logger.error("Timeout waiting for response from scheduler.")
|
||||
logger.error("Timeout waiting for response from %s.", endpoint)
|
||||
raise TimeoutError("Scheduler did not respond in time.")
|
||||
finally:
|
||||
if previous_timeout_ms is not None and self.scheduler_socket is not None:
|
||||
self.scheduler_socket.setsockopt(zmq.RCVTIMEO, previous_timeout_ms)
|
||||
socket.close()
|
||||
|
||||
def _forward_routed(self, batch: Any, timeout_ms: int | None) -> Any:
|
||||
self.request_logger.log_received_request(batch)
|
||||
endpoints = self.server_args.scheduler_endpoints
|
||||
if isinstance(batch, _CONTROL_REQ_TYPES):
|
||||
results = [self._forward_one(ep, batch, timeout_ms) for ep in endpoints]
|
||||
output_batch = _merge_fanout_results(results)
|
||||
else:
|
||||
replica = _select_replica(batch, len(endpoints), self._replica_counter)
|
||||
output_batch = self._forward_one(endpoints[replica], batch, timeout_ms)
|
||||
self.request_logger.log_finished_request(batch, output_batch)
|
||||
return output_batch
|
||||
|
||||
def ping(self) -> bool:
|
||||
"""
|
||||
@@ -111,27 +172,22 @@ class SchedulerClient:
|
||||
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()
|
||||
for endpoint in self.server_args.scheduler_endpoints:
|
||||
ping_socket = self.context.socket(zmq.REQ)
|
||||
ping_socket.setsockopt(zmq.LINGER, 0)
|
||||
ping_socket.setsockopt(zmq.RCVTIMEO, 2000) # 2-second timeout for pings
|
||||
try:
|
||||
ping_socket.connect(endpoint)
|
||||
ping_socket.send_pyobj({"method": "ping"})
|
||||
ping_socket.recv_pyobj()
|
||||
except zmq.error.Again:
|
||||
return False
|
||||
finally:
|
||||
ping_socket.close()
|
||||
return True
|
||||
|
||||
def close(self):
|
||||
"""Closes the socket and terminates the context."""
|
||||
if self.scheduler_socket:
|
||||
self.scheduler_socket.close()
|
||||
self.scheduler_socket = None
|
||||
"""Terminates the context."""
|
||||
if self.context:
|
||||
self.context.term()
|
||||
self.context = None
|
||||
@@ -150,6 +206,7 @@ class AsyncSchedulerClient:
|
||||
self.context = None
|
||||
self.server_args = None
|
||||
self.request_logger: Optional[DiffusionRequestLogger] = None
|
||||
self._replica_counter = itertools.count()
|
||||
|
||||
def initialize(self, server_args: ServerArgs):
|
||||
if self.context is not None and not self.context.closed:
|
||||
@@ -171,24 +228,33 @@ class AsyncSchedulerClient:
|
||||
"AsyncSchedulerClient is not initialized. Call initialize() first."
|
||||
)
|
||||
|
||||
# Create a temporary REQ socket for this request to allow concurrency
|
||||
endpoints = self.server_args.scheduler_endpoints
|
||||
if isinstance(batch, _CONTROL_REQ_TYPES):
|
||||
# replica state (weights, LoRA, memory) must change everywhere
|
||||
results = [await self._forward_one(ep, batch) for ep in endpoints]
|
||||
output_batch = _merge_fanout_results(results)
|
||||
else:
|
||||
replica = _select_replica(batch, len(endpoints), self._replica_counter)
|
||||
output_batch = await self._forward_one(endpoints[replica], batch)
|
||||
self.request_logger.log_finished_request(batch, output_batch)
|
||||
return output_batch
|
||||
|
||||
async def _forward_one(self, endpoint: str, batch: Any) -> Any:
|
||||
# a temporary REQ socket per request keeps concurrent requests from
|
||||
# interleaving on one socket's strict send/recv alternation
|
||||
socket = self.context.socket(zmq.REQ)
|
||||
socket.setsockopt(zmq.LINGER, 0)
|
||||
# 100 minute timeout
|
||||
socket.setsockopt(zmq.RCVTIMEO, 6000000)
|
||||
|
||||
endpoint = self.server_args.scheduler_endpoint
|
||||
socket.connect(endpoint)
|
||||
|
||||
try:
|
||||
await socket.send(pickle.dumps(batch))
|
||||
payload = await socket.recv()
|
||||
output_batch = pickle.loads(payload)
|
||||
_materialize_output_batch_file_refs(output_batch)
|
||||
self.request_logger.log_finished_request(batch, output_batch)
|
||||
return output_batch
|
||||
except zmq.error.Again:
|
||||
logger.error("Timeout waiting for response from scheduler.")
|
||||
logger.error("Timeout waiting for response from %s.", endpoint)
|
||||
raise TimeoutError("Scheduler did not respond in time.")
|
||||
finally:
|
||||
socket.close()
|
||||
@@ -201,21 +267,19 @@ class AsyncSchedulerClient:
|
||||
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()
|
||||
for endpoint in self.server_args.scheduler_endpoints:
|
||||
ping_socket = self.context.socket(zmq.REQ)
|
||||
ping_socket.setsockopt(zmq.LINGER, 0)
|
||||
ping_socket.setsockopt(zmq.RCVTIMEO, 2000)
|
||||
try:
|
||||
ping_socket.connect(endpoint)
|
||||
await ping_socket.send(pickle.dumps({"method": "ping"}))
|
||||
await ping_socket.recv()
|
||||
except zmq.error.Again:
|
||||
return False
|
||||
finally:
|
||||
ping_socket.close()
|
||||
return True
|
||||
|
||||
def close(self):
|
||||
"""Closes the socket and terminates the context."""
|
||||
|
||||
@@ -222,7 +222,6 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
# number of data parallelism groups
|
||||
dp_size: int = 1
|
||||
# number of gpu in a dp group
|
||||
dp_degree: int = 1
|
||||
# cfg parallel (None = auto-decide based on num_gpus)
|
||||
enable_cfg_parallel: Optional[bool] = None
|
||||
# number of GPUs in each CFG parallel group (None = auto, 1 = disabled, N > 1 = enabled)
|
||||
@@ -356,6 +355,8 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
webui_port: int | None = 12312
|
||||
|
||||
scheduler_port: int = 5555
|
||||
# settled ingress ports, one per DP replica; None until ports are settled
|
||||
scheduler_ports: list[int] | None = None
|
||||
batching_mode: str = "dynamic"
|
||||
batching_max_size: int = 1
|
||||
batching_delay_ms: float = 0.0
|
||||
@@ -974,7 +975,10 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
requested_ports = []
|
||||
if needs_http:
|
||||
requested_ports.append((self.port, "HTTP"))
|
||||
requested_ports.append((self.scheduler_port, "Scheduler"))
|
||||
for replica in range(self.dp_size or 1):
|
||||
requested_ports.append(
|
||||
(self.scheduler_port + replica, f"Scheduler[{replica}]")
|
||||
)
|
||||
if self.master_port is not None:
|
||||
requested_ports.append((self.master_port, "Master"))
|
||||
seen_ports: dict[int, str] = {}
|
||||
@@ -998,6 +1002,13 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
initial_scheduler_port, avoid=settled_ports
|
||||
)
|
||||
settled_ports.add(self.scheduler_port)
|
||||
self.scheduler_ports = [self.scheduler_port]
|
||||
for _ in range((self.dp_size or 1) - 1):
|
||||
port = self.settle_port(
|
||||
self.scheduler_ports[-1] + 1, avoid=settled_ports
|
||||
)
|
||||
settled_ports.add(port)
|
||||
self.scheduler_ports.append(port)
|
||||
if self.master_port is not None:
|
||||
self.master_port = self.settle_port(
|
||||
self.master_port, 37, avoid=settled_ports
|
||||
@@ -2047,10 +2058,22 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
Internal endpoint for scheduler.
|
||||
Prefers the configured host but normalizes localhost -> 127.0.0.1 to avoid ZMQ issues.
|
||||
"""
|
||||
return self.scheduler_endpoint_for(0)
|
||||
|
||||
def scheduler_endpoint_for(self, replica: int) -> str:
|
||||
"""Ingress endpoint of one DP replica's driver rank."""
|
||||
scheduler_host = self.host
|
||||
if scheduler_host is None or scheduler_host == "localhost":
|
||||
scheduler_host = "127.0.0.1"
|
||||
return f"tcp://{scheduler_host}:{self.scheduler_port}"
|
||||
if self.scheduler_ports is not None:
|
||||
port = self.scheduler_ports[replica]
|
||||
else:
|
||||
port = self.scheduler_port + replica
|
||||
return f"tcp://{scheduler_host}:{port}"
|
||||
|
||||
@property
|
||||
def scheduler_endpoints(self) -> list[str]:
|
||||
return [self.scheduler_endpoint_for(r) for r in range(self.dp_size or 1)]
|
||||
|
||||
def settle_port(
|
||||
self,
|
||||
@@ -2459,8 +2482,11 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
if self.dp_size < 1:
|
||||
raise ValueError("--dp-size must be a natural number")
|
||||
|
||||
if self.dp_size > 1:
|
||||
raise ValueError("DP is not yet supported")
|
||||
if self.dp_size > 1 and self.disagg_role != RoleType.MONOLITHIC:
|
||||
raise ValueError(
|
||||
"--dp-size > 1 is only supported for monolithic serving; "
|
||||
"disaggregated roles scale by adding role instances instead"
|
||||
)
|
||||
|
||||
num_gpus_per_group = self.dp_size * self.tp_size
|
||||
if self.enable_cfg_parallel:
|
||||
|
||||
@@ -1144,6 +1144,7 @@ STANDALONE_FILES = {
|
||||
"../single_test_file/test_disagg_server.py",
|
||||
"../single_test_file/test_ar_models.py",
|
||||
"../single_test_file/test_ipc_a2a_2_gpu.py",
|
||||
"../single_test_file/test_dp_serving_2_gpu.py",
|
||||
],
|
||||
}
|
||||
|
||||
@@ -1178,6 +1179,8 @@ STANDALONE_FILE_EST_TIMES = {
|
||||
"../single_test_file/test_ar_models.py": 600.0,
|
||||
# no model load; the cost is the one-time JIT build of the sync kernels
|
||||
"../single_test_file/test_ipc_a2a_2_gpu.py": 240.0,
|
||||
# zimage server startup dominates; six short requests after warmup
|
||||
"../single_test_file/test_dp_serving_2_gpu.py": 900.0,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,150 @@
|
||||
"""Data-parallel serving must route, replicate, and agree.
|
||||
|
||||
Launches one server with --dp-size 2 (one GPU per replica) and checks the three
|
||||
properties that make DP real rather than a parsed flag: both replica drivers
|
||||
bind their own ingress, both replicas serve traffic (round-robin means two
|
||||
sequential requests land on different replicas), and a fixed seed produces the
|
||||
same image bytes from either replica.
|
||||
|
||||
pytest -v python/sglang/multimodal_gen/test/single_test_file/test_dp_serving_2_gpu.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import unittest
|
||||
import urllib.request
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
_MODEL = "Tongyi-MAI/Z-Image-Turbo"
|
||||
_PORT = 30811
|
||||
_STARTUP_TIMEOUT_S = 1200
|
||||
|
||||
|
||||
def _post_generation(prompt: str, seed: int) -> dict:
|
||||
payload = json.dumps(
|
||||
{
|
||||
"prompt": prompt,
|
||||
"size": "512x512",
|
||||
"seed": seed,
|
||||
"num_inference_steps": 20,
|
||||
"response_format": "b64_json",
|
||||
}
|
||||
).encode()
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{_PORT}/v1/images/generations",
|
||||
data=payload,
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=600) as resp:
|
||||
return json.loads(resp.read())
|
||||
|
||||
|
||||
def _image_md5(response: dict) -> str:
|
||||
b64 = response["data"][0]["b64_json"]
|
||||
return hashlib.md5(base64.b64decode(b64)).hexdigest()
|
||||
|
||||
|
||||
class TestDpServingTwoGpu(CustomTestCase):
|
||||
def test_two_replicas_serve_and_agree(self):
|
||||
if not current_platform.is_cuda():
|
||||
self.skipTest("DP e2e is exercised on CUDA")
|
||||
if torch.cuda.device_count() < 2:
|
||||
self.skipTest("needs 2 GPUs")
|
||||
|
||||
log_path = Path("/tmp/dp_serving_test.log")
|
||||
fh = open(log_path, "w")
|
||||
proc = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"sglang.multimodal_gen.runtime.entrypoints.cli.main",
|
||||
"serve",
|
||||
"--model-path",
|
||||
_MODEL,
|
||||
"--num-gpus",
|
||||
"2",
|
||||
"--dp-size",
|
||||
"2",
|
||||
"--enable-cfg-parallel",
|
||||
"false",
|
||||
"--port",
|
||||
str(_PORT),
|
||||
],
|
||||
stdout=fh,
|
||||
stderr=subprocess.STDOUT,
|
||||
preexec_fn=os.setsid,
|
||||
env=os.environ.copy(),
|
||||
)
|
||||
try:
|
||||
deadline = time.monotonic() + _STARTUP_TIMEOUT_S
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
with urllib.request.urlopen(
|
||||
f"http://127.0.0.1:{_PORT}/health", timeout=5
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
break
|
||||
except Exception:
|
||||
time.sleep(5)
|
||||
else:
|
||||
self.fail(f"server not healthy in {_STARTUP_TIMEOUT_S}s")
|
||||
|
||||
# warm each replica once so first-request costs stay out of timing
|
||||
for _ in range(2):
|
||||
_post_generation("warm pear", seed=7)
|
||||
|
||||
# same seed through two round-robined replicas -> identical bytes
|
||||
first = _post_generation("a pear on a table", seed=42)
|
||||
second = _post_generation("a pear on a table", seed=42)
|
||||
self.assertEqual(_image_md5(first), _image_md5(second))
|
||||
|
||||
log = log_path.read_text(errors="ignore")
|
||||
self.assertIn("dp replica 0) bind", log)
|
||||
self.assertIn("dp replica 1) bind", log)
|
||||
|
||||
# distribution, not just agreement: a concurrent pair must run in
|
||||
# about one request's wall time. Two requests serialized on a
|
||||
# single replica would take ~2x the single-request time, so the
|
||||
# 1.6x bound separates the behaviors with margin for jitter.
|
||||
t0 = time.monotonic()
|
||||
single = _post_generation("a pear on a table", seed=42)
|
||||
single_s = time.monotonic() - t0
|
||||
self.assertTrue(single["data"])
|
||||
|
||||
t0 = time.monotonic()
|
||||
with ThreadPoolExecutor(max_workers=2) as pool:
|
||||
futures = [
|
||||
pool.submit(_post_generation, "a pear on a table", 42)
|
||||
for _ in range(2)
|
||||
]
|
||||
results = [f.result() for f in futures]
|
||||
pair_s = time.monotonic() - t0
|
||||
for r in results:
|
||||
self.assertTrue(r["data"])
|
||||
self.assertLess(
|
||||
pair_s,
|
||||
1.6 * single_s,
|
||||
f"concurrent pair took {pair_s:.2f}s vs single {single_s:.2f}s; "
|
||||
"requests are serializing on one replica",
|
||||
)
|
||||
finally:
|
||||
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
|
||||
fh.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,71 @@
|
||||
"""Routing invariants for DP serving.
|
||||
|
||||
Generation requests go to exactly one replica, control ops reach every replica,
|
||||
and a realtime session always lands on the same replica it started on.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
|
||||
from sglang.multimodal_gen.runtime.entrypoints.utils import SetLoraReq, ShutdownReq
|
||||
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 (
|
||||
_CONTROL_REQ_TYPES,
|
||||
_merge_fanout_results,
|
||||
_select_replica,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
|
||||
|
||||
def _req(**kwargs) -> Req:
|
||||
return Req(prompt="a pear", **kwargs)
|
||||
|
||||
|
||||
def test_round_robin_covers_every_replica():
|
||||
counter = itertools.count()
|
||||
picks = [_select_replica([_req()], 3, counter) for _ in range(6)]
|
||||
assert picks == [0, 1, 2, 0, 1, 2]
|
||||
|
||||
|
||||
def test_dp1_always_selects_replica_zero():
|
||||
counter = itertools.count()
|
||||
assert {_select_replica([_req()], 1, counter) for _ in range(5)} == {0}
|
||||
|
||||
|
||||
def test_session_requests_stick_to_one_replica():
|
||||
counter = itertools.count()
|
||||
session_req = _req(realtime_session_id="session-abc")
|
||||
picks = {_select_replica([session_req], 4, counter) for _ in range(8)}
|
||||
assert len(picks) == 1
|
||||
# and the counter was never consumed by session traffic
|
||||
assert _select_replica([_req()], 4, counter) == 0
|
||||
|
||||
|
||||
def test_control_reqs_are_recognized():
|
||||
assert isinstance(SetLoraReq(lora_nickname="x", lora_path="y"), _CONTROL_REQ_TYPES)
|
||||
assert isinstance(ShutdownReq(), _CONTROL_REQ_TYPES)
|
||||
assert not isinstance([_req()], _CONTROL_REQ_TYPES)
|
||||
assert not isinstance(_req(), _CONTROL_REQ_TYPES)
|
||||
|
||||
|
||||
def test_fanout_merge_surfaces_the_failing_replica():
|
||||
ok = OutputBatch(output=None)
|
||||
bad = OutputBatch(error="replica 1 exploded")
|
||||
assert _merge_fanout_results([ok, bad]) is bad
|
||||
assert _merge_fanout_results([ok, OutputBatch(output=None)]) is ok
|
||||
|
||||
|
||||
def test_scheduler_endpoints_one_per_replica():
|
||||
args = ServerArgs.__new__(ServerArgs)
|
||||
args.host = "localhost"
|
||||
args.dp_size = 3
|
||||
args.scheduler_port = 6000
|
||||
args.scheduler_ports = None
|
||||
assert ServerArgs.scheduler_endpoints.fget(args) == [
|
||||
"tcp://127.0.0.1:6000",
|
||||
"tcp://127.0.0.1:6001",
|
||||
"tcp://127.0.0.1:6002",
|
||||
]
|
||||
# settled ports need not be consecutive; the endpoint list follows them
|
||||
args.scheduler_ports = [6000, 7005, 7100]
|
||||
assert ServerArgs.scheduler_endpoint_for(args, 2) == "tcp://127.0.0.1:7100"
|
||||
Reference in New Issue
Block a user