723 lines
27 KiB
Python
723 lines
27 KiB
Python
# Copyright 2023-2024 SGLang Team
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""A controller that dispatches requests to multiple data parallel workers."""
|
|
|
|
import faulthandler
|
|
import logging
|
|
import multiprocessing as mp
|
|
import signal
|
|
import threading
|
|
import time
|
|
from enum import Enum, auto
|
|
from typing import Callable, List, Optional
|
|
|
|
import psutil
|
|
import setproctitle
|
|
import zmq
|
|
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
|
|
from sglang.srt.managers.io_struct import (
|
|
ActiveRanksOutput,
|
|
BatchTokenizedEmbeddingReqInput,
|
|
BatchTokenizedGenerateReqInput,
|
|
BlockReqInput,
|
|
TokenizedEmbeddingReqInput,
|
|
TokenizedGenerateReqInput,
|
|
WatchLoadUpdateReq,
|
|
DpRequestInfoqOutput,
|
|
)
|
|
from sglang.srt.managers.schedule_batch import Req
|
|
from sglang.srt.managers.scheduler import run_scheduler_process
|
|
from sglang.srt.observability.cpu_monitor import start_cpu_monitor_thread
|
|
from sglang.srt.observability.req_time_stats import DPControllerReqTimeStats
|
|
from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
|
|
from sglang.srt.server_args import (
|
|
DP_ATTENTION_HANDSHAKE_PORT_DELTA,
|
|
PortArgs,
|
|
ServerArgs,
|
|
)
|
|
from sglang.srt.utils import numa_utils
|
|
from sglang.srt.utils.common import (
|
|
configure_logger,
|
|
kill_itself_when_parent_died,
|
|
maybe_reindex_device_id,
|
|
)
|
|
from sglang.srt.utils.network import NetworkAddress, bind_port, get_zmq_socket
|
|
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
|
|
from sglang.srt.utils.watchdog import Watchdog
|
|
from sglang.utils import TypeBasedDispatcher, get_exception_traceback
|
|
from sglang.srt.disaggregation.utils import (
|
|
DisaggregationMode
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class LoadBalanceMethod(Enum):
|
|
"""Load balance method."""
|
|
|
|
ROUND_ROBIN = auto()
|
|
FOLLOW_BOOTSTRAP_ROOM = auto()
|
|
TOTAL_REQUESTS = auto()
|
|
TOTAL_TOKENS = auto()
|
|
|
|
@classmethod
|
|
def from_str(cls, method: str):
|
|
method = method.upper()
|
|
try:
|
|
return cls[method]
|
|
except KeyError as exc:
|
|
raise ValueError(f"Invalid load balance method: {method}") from exc
|
|
|
|
import random
|
|
class DPBudget:
|
|
def __init__(self, server_args: ServerArgs):
|
|
dp_size = server_args.dp_size
|
|
self.dp_size = dp_size
|
|
self.total_requests = [0] * dp_size
|
|
self.total_tokens = [0] * dp_size
|
|
self.ts_tic = [0.0] * dp_size
|
|
self.pool_tokens = [0] * dp_size
|
|
self.req_sum_tokens = [0] * dp_size
|
|
self.consume_token_num_per_rank = 4000
|
|
self.disaggregation_mode = DisaggregationMode(
|
|
server_args.disaggregation_mode
|
|
)
|
|
|
|
# 每个 rank 一个账本，按 timestamp 排序
|
|
self.reqs_info = [list() for _ in range(dp_size)]
|
|
|
|
def _sort_rank_ledger(self, rank: int):
|
|
self.reqs_info[rank].sort(key=lambda x: x.timestamp)
|
|
|
|
def _append_local_req(self, rank: int, req: Req):
|
|
if req is None:
|
|
return
|
|
|
|
rid = getattr(req, "rid", None)
|
|
if rid is None:
|
|
return
|
|
|
|
num_tokens = len(req.input_ids) if getattr(req, "input_ids", None) is not None else 0
|
|
completion_tokens = (
|
|
len(req.output_ids) if getattr(req, "output_ids", None) is not None else 0
|
|
)
|
|
cached_tokens = getattr(req, "cached_tokens", 0)
|
|
|
|
info = DpRequestInfoqOutput(
|
|
rid=rid,
|
|
timestamp=time.perf_counter(),
|
|
dp_rank=rank,
|
|
num_tokens=num_tokens,
|
|
cached_tokens=cached_tokens,
|
|
completion_tokens=completion_tokens,
|
|
)
|
|
self.reqs_info[rank].append(info)
|
|
|
|
def _merge_rank_reqs(self, rank: int, incoming_reqs: List[DpRequestInfoqOutput]):
|
|
if incoming_reqs is None:
|
|
return
|
|
|
|
if len(incoming_reqs) == 0:
|
|
self.reqs_info[rank] = []
|
|
return
|
|
|
|
incoming_reqs = sorted(incoming_reqs, key=lambda x: x.timestamp)
|
|
earliest_ts = incoming_reqs[0].timestamp
|
|
|
|
# 删除本地账本中比 incoming_reqs 中最早 timestamp 更旧的请求
|
|
local_reqs = [x for x in self.reqs_info[rank] if x.timestamp >= earliest_ts]
|
|
|
|
# 按 rid 合并
|
|
merged = {x.rid: x for x in local_reqs}
|
|
for x in incoming_reqs:
|
|
merged[x.rid] = x
|
|
|
|
self.reqs_info[rank] = sorted(merged.values(), key=lambda x: x.timestamp)
|
|
|
|
def update_budget(self, load_update: WatchLoadUpdateReq):
|
|
for load in load_update.loads:
|
|
if load is None:
|
|
continue
|
|
|
|
rank = load.dp_rank
|
|
if rank is None or rank >= self.dp_size:
|
|
continue
|
|
|
|
self.total_requests[rank] = load.num_reqs
|
|
self.total_tokens[rank] = load.num_tokens
|
|
self.ts_tic[rank] = load.ts_tic
|
|
self.pool_tokens[rank] = load.pool_tokens
|
|
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
|
self._merge_rank_reqs(rank, load.reqs_list)
|
|
else:
|
|
self._merge_rank_reqs(rank, load.reqs_list)
|
|
self.req_sum_tokens[rank] = sum(x.num_tokens for x in self.reqs_info[rank])+self.pool_tokens[rank]
|
|
|
|
def dispatch(self, method: LoadBalanceMethod, req: Req = None):
|
|
if method == LoadBalanceMethod.TOTAL_REQUESTS:
|
|
min_req = min(self.total_requests)
|
|
candidates = [i for i, x in enumerate(self.total_requests) if x == min_req]
|
|
target_rank = random.choice(candidates)
|
|
|
|
elif method == LoadBalanceMethod.TOTAL_TOKENS:
|
|
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
|
pairs = [
|
|
(self.total_tokens[i], self.total_requests[i])
|
|
for i in range(self.dp_size)
|
|
]
|
|
min_pair = min(pairs)
|
|
candidates = [
|
|
i for i in range(self.dp_size)
|
|
if (self.total_tokens[i], self.total_requests[i]) == min_pair
|
|
]
|
|
target_rank = random.choice(candidates)
|
|
min_tokens = self.total_tokens[target_rank]
|
|
else:
|
|
pairs = [
|
|
(self.total_tokens[i], self.total_requests[i])
|
|
for i in range(self.dp_size)
|
|
]
|
|
min_pair = min(pairs)
|
|
candidates = [
|
|
i for i in range(self.dp_size)
|
|
if (self.total_tokens[i], self.total_requests[i]) == min_pair
|
|
]
|
|
target_rank = random.choice(candidates)
|
|
min_tokens = self.total_tokens[target_rank]
|
|
|
|
logger.info(
|
|
f"Dispatching to DP rank {target_rank} with rough_num_tokens={min_tokens} and req_sum_tokens={self.req_sum_tokens[target_rank]}"
|
|
)
|
|
|
|
else:
|
|
return None
|
|
|
|
self.total_requests[target_rank] += 1
|
|
if req is not None:
|
|
self.total_tokens[target_rank] += len(req.input_ids)
|
|
self.req_sum_tokens[target_rank] += len(req.input_ids)
|
|
self._append_local_req(target_rank, req)
|
|
|
|
return target_rank
|
|
|
|
class DataParallelController:
|
|
"""A controller that dispatches requests to multiple data parallel workers."""
|
|
|
|
def __init__(
|
|
self,
|
|
server_args: ServerArgs,
|
|
port_args: PortArgs,
|
|
run_scheduler_process_func: Callable,
|
|
) -> None:
|
|
# Parse args
|
|
self.server_args = server_args
|
|
self.port_args = port_args
|
|
self.load_balance_method = LoadBalanceMethod.from_str(
|
|
server_args.load_balance_method
|
|
)
|
|
self.run_scheduler_process_func = run_scheduler_process_func
|
|
|
|
# For DP balance
|
|
self.global_balance_id = 0
|
|
|
|
# Init inter-process communication
|
|
self.context = zmq.Context(1 + server_args.dp_size)
|
|
if server_args.node_rank == 0:
|
|
self.recv_from_tokenizer = get_zmq_socket(
|
|
self.context, zmq.PULL, port_args.scheduler_input_ipc_name, False
|
|
)
|
|
|
|
# Dispatch method
|
|
self.round_robin_counter = 0
|
|
dispatch_lookup = {
|
|
LoadBalanceMethod.ROUND_ROBIN: self.round_robin_scheduler,
|
|
LoadBalanceMethod.FOLLOW_BOOTSTRAP_ROOM: self.follow_bootstrap_room_scheduler,
|
|
LoadBalanceMethod.TOTAL_REQUESTS: self.total_requests_scheduler,
|
|
LoadBalanceMethod.TOTAL_TOKENS: self.total_tokens_scheduler,
|
|
}
|
|
self.dispatching = dispatch_lookup[self.load_balance_method]
|
|
|
|
# Load balance budget
|
|
self.dp_budget = DPBudget(server_args)
|
|
|
|
# To protect changing env vars to set CUDA_VISIBLE_DEVICES.
|
|
self.env_lock = threading.Lock()
|
|
|
|
# Launch data parallel workers
|
|
self.scheduler_procs = []
|
|
self.workers: List[zmq.Socket] = [None] * server_args.dp_size
|
|
self.status: List[bool] = [True] * server_args.dp_size
|
|
|
|
if server_args.enable_dp_attention:
|
|
self.launch_dp_attention_schedulers(server_args, port_args)
|
|
self.control_message_step = server_args.tp_size
|
|
else:
|
|
self.launch_dp_schedulers(server_args, port_args)
|
|
self.control_message_step = 1
|
|
|
|
self.init_dispatcher()
|
|
|
|
self.soft_watchdog = Watchdog.create(
|
|
debug_name="DataParallelController",
|
|
watchdog_timeout=server_args.soft_watchdog_timeout,
|
|
soft=True,
|
|
test_stuck_time=envs.SGLANG_TEST_STUCK_DP_CONTROLLER.get(),
|
|
)
|
|
|
|
if server_args.enable_metrics:
|
|
start_cpu_monitor_thread("data_parallel_controller")
|
|
|
|
def send_to_all_workers(self, obj):
|
|
for i, worker in enumerate(self.workers):
|
|
if self.status[i]:
|
|
worker.send_pyobj(obj)
|
|
|
|
def send_control_message(self, obj):
|
|
# Send control messages to first worker of tp group
|
|
for worker in self.workers[:: self.control_message_step]:
|
|
worker.send_pyobj(obj)
|
|
|
|
def handle_load_update_req(self, obj):
|
|
self.dp_budget.update_budget(obj)
|
|
|
|
def update_active_ranks(self, ranks: ActiveRanksOutput):
|
|
self.status = ranks.status
|
|
|
|
def dispatching_with_trace(self, req: Req):
|
|
req.time_stats = DPControllerReqTimeStats.new_from_obj(req.time_stats)
|
|
|
|
req.time_stats.set_dp_dispatch_time()
|
|
self.dispatching(req)
|
|
req.time_stats.set_dp_dispatch_finish_time()
|
|
|
|
def dispatch_batch_generate(self, batch_req: BatchTokenizedGenerateReqInput):
|
|
for req in batch_req:
|
|
self.dispatching_with_trace(req)
|
|
|
|
def dispatch_batch_embedding(self, batch_req: BatchTokenizedEmbeddingReqInput):
|
|
for req in batch_req:
|
|
self.dispatching_with_trace(req)
|
|
|
|
def init_dispatcher(self):
|
|
self._request_dispatcher = TypeBasedDispatcher(
|
|
[
|
|
(TokenizedGenerateReqInput, self.dispatching_with_trace),
|
|
(TokenizedEmbeddingReqInput, self.dispatching_with_trace),
|
|
(BatchTokenizedGenerateReqInput, self.dispatch_batch_generate),
|
|
(BatchTokenizedEmbeddingReqInput, self.dispatch_batch_embedding),
|
|
(BlockReqInput, self.send_to_all_workers),
|
|
(WatchLoadUpdateReq, self.handle_load_update_req),
|
|
(ActiveRanksOutput, self.update_active_ranks),
|
|
]
|
|
)
|
|
self._request_dispatcher.add_fallback_fn(self.send_control_message)
|
|
|
|
def launch_dp_schedulers(self, server_args, port_args):
|
|
base_gpu_id = 0
|
|
|
|
threads = []
|
|
sockets = []
|
|
ready_events = []
|
|
for dp_rank in range(server_args.dp_size):
|
|
tmp_port_args = PortArgs.init_new(server_args)
|
|
tmp_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name
|
|
tmp_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name
|
|
|
|
# This port is checked free in PortArgs.init_new.
|
|
# We hold it first so that the next dp worker gets a different port
|
|
sockets.append(bind_port(tmp_port_args.nccl_port))
|
|
|
|
ready_event = threading.Event()
|
|
ready_events.append(ready_event)
|
|
|
|
# Create a thread for each worker
|
|
thread = threading.Thread(
|
|
target=self.launch_tensor_parallel_group_thread,
|
|
args=(server_args, tmp_port_args, base_gpu_id, dp_rank, ready_event),
|
|
)
|
|
threads.append(thread)
|
|
base_gpu_id += (
|
|
server_args.tp_size * server_args.pp_size * server_args.gpu_id_step
|
|
)
|
|
|
|
if server_args.node_rank == 0:
|
|
self.workers[dp_rank] = get_zmq_socket(
|
|
self.context,
|
|
zmq.PUSH,
|
|
tmp_port_args.scheduler_input_ipc_name,
|
|
True,
|
|
)
|
|
|
|
# Free all sockets before starting the threads to launch TP workers
|
|
for sock in sockets:
|
|
sock.close()
|
|
|
|
# Start all threads
|
|
for thread in threads:
|
|
thread.start()
|
|
for event in ready_events:
|
|
event.wait()
|
|
|
|
def launch_tensor_parallel_group_thread(
|
|
self,
|
|
server_args: ServerArgs,
|
|
port_args: PortArgs,
|
|
base_gpu_id: int,
|
|
dp_rank: int,
|
|
ready_event: threading.Event,
|
|
):
|
|
self.launch_tensor_parallel_group(server_args, port_args, base_gpu_id, dp_rank)
|
|
ready_event.set()
|
|
|
|
# This thread cannot be closed because otherwise the `kill_itself_when_parent_died`
|
|
# function in scheduler.py will kill the scheduler.
|
|
while True:
|
|
time.sleep(30 * 24 * 3600)
|
|
|
|
def _broadcast_worker_ports(
|
|
self, server_args: ServerArgs, worker_ports: Optional[List[int]] = None
|
|
) -> List[int]:
|
|
"""Broadcast worker ports from node 0 to all other nodes.
|
|
|
|
Node 0 acts as the server, waiting for all other nodes to connect and
|
|
sending them the pre-allocated worker ports. Other nodes act as clients,
|
|
connecting to node 0 to receive their copy of the worker ports.
|
|
|
|
Args:
|
|
server_args: Server arguments containing node configuration.
|
|
worker_ports: Pre-allocated worker ports to broadcast.
|
|
|
|
Returns:
|
|
List of worker ports (same on all nodes after broadcast).
|
|
"""
|
|
# Determine the endpoint for inter-node communication
|
|
if server_args.dist_init_addr is None:
|
|
na = NetworkAddress(
|
|
server_args.host or "127.0.0.1",
|
|
server_args.port + DP_ATTENTION_HANDSHAKE_PORT_DELTA,
|
|
)
|
|
else:
|
|
na = NetworkAddress.parse(server_args.dist_init_addr)
|
|
na = NetworkAddress(na.host, na.port + DP_ATTENTION_HANDSHAKE_PORT_DELTA)
|
|
endpoint = na.to_tcp()
|
|
|
|
if server_args.node_rank == 0:
|
|
# Node 0: Broadcast worker ports to all other nodes
|
|
return self._broadcast_ports_as_server(
|
|
endpoint, server_args.nnodes - 1, worker_ports
|
|
)
|
|
else:
|
|
# Other nodes: Receive worker ports from node 0
|
|
return self._receive_ports_as_client(endpoint, server_args.node_rank)
|
|
|
|
def _broadcast_ports_as_server(
|
|
self, endpoint: str, expected_clients: int, worker_ports: List[int]
|
|
) -> List[int]:
|
|
"""Broadcast worker ports to all client nodes."""
|
|
logger.debug(f"Broadcasting worker ports to {expected_clients} client nodes")
|
|
logger.debug(f"Worker ports: {worker_ports}")
|
|
|
|
rep_socket = get_zmq_socket(self.context, zmq.REP, endpoint, True)
|
|
|
|
try:
|
|
connected_clients = 0
|
|
while connected_clients < expected_clients:
|
|
# Wait for client handshake
|
|
client_rank = rep_socket.recv().decode()
|
|
logger.debug(f"Received handshake from node {client_rank}")
|
|
|
|
# Send worker ports to client
|
|
rep_socket.send_pyobj(worker_ports)
|
|
connected_clients += 1
|
|
logger.debug(
|
|
f"Sent worker ports to {connected_clients}/{expected_clients} nodes"
|
|
)
|
|
|
|
logger.debug("Worker port broadcast completed")
|
|
return worker_ports
|
|
finally:
|
|
rep_socket.close()
|
|
|
|
def _receive_ports_as_client(self, endpoint: str, node_rank: int) -> List[int]:
|
|
"""Receive worker ports from the server node."""
|
|
logger.debug(f"Connecting to node 0 to receive worker ports")
|
|
|
|
req_socket = get_zmq_socket(self.context, zmq.REQ, endpoint, False)
|
|
req_socket.setsockopt(zmq.RCVTIMEO, 600 * 1000) # 10 minute timeout
|
|
req_socket.setsockopt(zmq.SNDTIMEO, 600 * 1000)
|
|
|
|
try:
|
|
# Send handshake with our node rank
|
|
req_socket.send(str(node_rank).encode())
|
|
|
|
# Receive worker ports
|
|
worker_ports = req_socket.recv_pyobj()
|
|
logger.debug(f"Received {len(worker_ports)} worker ports from node 0")
|
|
return worker_ports
|
|
except zmq.Again:
|
|
logger.error("Timeout waiting for worker ports from node 0")
|
|
raise RuntimeError(
|
|
"Failed to receive worker ports from node 0 within timeout"
|
|
)
|
|
finally:
|
|
req_socket.close()
|
|
|
|
def launch_dp_attention_schedulers(
|
|
self, server_args: ServerArgs, port_args: PortArgs
|
|
):
|
|
# Pre-allocate worker ports on node 0 to avoid conflicts
|
|
worker_ports = []
|
|
if server_args.node_rank == 0:
|
|
for dp_rank in range(server_args.dp_size):
|
|
port_and_socket = get_zmq_socket(self.context, zmq.PUSH)
|
|
worker_ports.append(port_and_socket[0])
|
|
self.workers[dp_rank] = port_and_socket[1]
|
|
logger.debug(f"Assigned port {port_and_socket[0]} to worker {dp_rank}")
|
|
|
|
broadcasted_ports = self._broadcast_worker_ports(
|
|
server_args, worker_ports if worker_ports else None
|
|
)
|
|
self.launch_tensor_parallel_group(
|
|
server_args, port_args, 0, None, broadcasted_ports
|
|
)
|
|
|
|
def launch_tensor_parallel_group(
|
|
self,
|
|
server_args: ServerArgs,
|
|
port_args: PortArgs,
|
|
base_gpu_id: int,
|
|
dp_rank: Optional[int],
|
|
worker_ports: Optional[List[int]] = None,
|
|
):
|
|
if not server_args.enable_dp_attention:
|
|
logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.")
|
|
|
|
memory_saver_adapter = TorchMemorySaverAdapter.create(
|
|
enable=server_args.enable_memory_saver
|
|
)
|
|
|
|
scheduler_pipe_readers = []
|
|
|
|
pp_size_per_node = max(server_args.pp_size // server_args.nnodes, 1)
|
|
nnodes_per_pp_rank = max(server_args.nnodes // server_args.pp_size, 1)
|
|
pp_rank_range = range(
|
|
pp_size_per_node * (server_args.node_rank // nnodes_per_pp_rank),
|
|
pp_size_per_node * (server_args.node_rank // nnodes_per_pp_rank + 1),
|
|
)
|
|
|
|
nnodes_per_tp_group = nnodes_per_pp_rank
|
|
tp_size_per_node = server_args.tp_size // nnodes_per_tp_group
|
|
tp_rank_range = range(
|
|
tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group),
|
|
tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group + 1),
|
|
)
|
|
|
|
attn_cp_rank = 0
|
|
moe_dp_rank = 0
|
|
for pp_rank in pp_rank_range:
|
|
for tp_rank in tp_rank_range:
|
|
rank_port_args = port_args
|
|
|
|
if server_args.enable_dp_attention:
|
|
# dp attention has different sharding logic
|
|
_, _, dp_rank = compute_dp_attention_world_info(
|
|
server_args.enable_dp_attention,
|
|
tp_rank,
|
|
server_args.tp_size,
|
|
server_args.dp_size,
|
|
server_args.attn_cp_size,
|
|
)
|
|
# compute zmq ports for this dp rank
|
|
rank_port_args = PortArgs.init_new(
|
|
server_args, dp_rank, worker_ports
|
|
)
|
|
# Data parallelism reuses the tensor parallelism group,
|
|
# so all dp ranks should use the same nccl port.
|
|
rank_port_args.nccl_port = port_args.nccl_port
|
|
|
|
reader, writer = mp.Pipe(duplex=False)
|
|
gpu_id = (
|
|
server_args.base_gpu_id
|
|
+ base_gpu_id
|
|
+ ((pp_rank % pp_size_per_node) * tp_size_per_node)
|
|
+ (tp_rank % tp_size_per_node) * server_args.gpu_id_step
|
|
)
|
|
attn_dp_size = (
|
|
server_args.dp_size if server_args.enable_dp_attention else 1
|
|
)
|
|
|
|
# Parallelism hierarchy (outermost to innermost):
|
|
# - Attention: Global(TP) -> DP -> ATTN_CP -> ATTN_TP (innermost)
|
|
# - MoE: Global(TP) -> MOE_DP -> EP -> MOE_TP (innermost)
|
|
attn_tp_size = (
|
|
server_args.tp_size // attn_dp_size // server_args.attn_cp_size
|
|
)
|
|
attn_cp_rank = (tp_rank // attn_tp_size) % server_args.attn_cp_size
|
|
moe_dp_rank = tp_rank // (
|
|
server_args.tp_size // server_args.moe_dp_size
|
|
)
|
|
moe_ep_rank = (
|
|
tp_rank
|
|
% (server_args.tp_size // server_args.moe_dp_size)
|
|
// (
|
|
server_args.tp_size
|
|
// server_args.moe_dp_size
|
|
// server_args.ep_size
|
|
)
|
|
)
|
|
|
|
with self.env_lock, maybe_reindex_device_id(gpu_id) as gpu_id:
|
|
proc = mp.Process(
|
|
target=self.run_scheduler_process_func,
|
|
args=(
|
|
server_args,
|
|
rank_port_args,
|
|
gpu_id,
|
|
tp_rank,
|
|
attn_cp_rank,
|
|
moe_dp_rank,
|
|
moe_ep_rank,
|
|
pp_rank,
|
|
dp_rank,
|
|
writer,
|
|
),
|
|
)
|
|
with memory_saver_adapter.configure_subprocess(), numa_utils.configure_subprocess(
|
|
server_args, gpu_id
|
|
):
|
|
proc.start()
|
|
self.scheduler_procs.append(proc)
|
|
scheduler_pipe_readers.append(reader)
|
|
|
|
# Wait for model to finish loading
|
|
scheduler_info = []
|
|
for i in range(len(scheduler_pipe_readers)):
|
|
scheduler_info.append(scheduler_pipe_readers[i].recv())
|
|
|
|
self.max_total_num_tokens = scheduler_info[0]["max_total_num_tokens"]
|
|
self.max_req_input_len = scheduler_info[0]["max_req_input_len"]
|
|
|
|
def maybe_external_dp_rank_routing(self, req: Req):
|
|
if req.routed_dp_rank is not None:
|
|
logger.debug(f"Direct routing to DP rank {req.routed_dp_rank}")
|
|
self.workers[req.routed_dp_rank].send_pyobj(req)
|
|
return True
|
|
return False
|
|
|
|
def round_robin_scheduler(self, req: Req):
|
|
if self.maybe_external_dp_rank_routing(req):
|
|
return
|
|
|
|
while True:
|
|
if self.status[self.round_robin_counter]:
|
|
logger.debug(f"Choose worker {self.round_robin_counter}")
|
|
self.workers[self.round_robin_counter].send_pyobj(req)
|
|
self.round_robin_counter = (self.round_robin_counter + 1) % len(
|
|
self.workers
|
|
)
|
|
break
|
|
self.round_robin_counter = (self.round_robin_counter + 1) % len(
|
|
self.workers
|
|
)
|
|
|
|
def follow_bootstrap_room_scheduler(self, req: Req):
|
|
if self.maybe_external_dp_rank_routing(req):
|
|
return
|
|
|
|
# Set default bootstrap_room if in FAKE auto mode and room is None
|
|
if (
|
|
req.bootstrap_room is None
|
|
and self.server_args.disaggregation_transfer_backend == "fake"
|
|
):
|
|
req.bootstrap_room = self.round_robin_counter
|
|
self.round_robin_counter = (self.round_robin_counter + 1) % len(
|
|
self.workers
|
|
)
|
|
|
|
assert req.bootstrap_room is not None, (
|
|
"req.bootstrap_room should not be None. Do not send requests directly to "
|
|
"prefill or decode instances; send to the router instead."
|
|
)
|
|
target_rank = req.bootstrap_room % len(self.workers)
|
|
self.workers[target_rank].send_pyobj(req)
|
|
|
|
def total_requests_scheduler(self, req: Req):
|
|
if self.maybe_external_dp_rank_routing(req):
|
|
return
|
|
target_worker = self.dp_budget.dispatch(LoadBalanceMethod.TOTAL_REQUESTS)
|
|
self.workers[target_worker].send_pyobj(req)
|
|
|
|
def total_tokens_scheduler(self, req: Req):
|
|
if self.maybe_external_dp_rank_routing(req):
|
|
return
|
|
target_worker = self.dp_budget.dispatch(LoadBalanceMethod.TOTAL_TOKENS,req=req)
|
|
self.workers[target_worker].send_pyobj(req)
|
|
|
|
def event_loop(self):
|
|
while True:
|
|
while True:
|
|
self.soft_watchdog.feed()
|
|
try:
|
|
recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
|
|
except zmq.ZMQError:
|
|
break
|
|
self._request_dispatcher(recv_req)
|
|
|
|
|
|
def run_data_parallel_controller_process(
|
|
server_args: ServerArgs,
|
|
port_args: PortArgs,
|
|
pipe_writer,
|
|
run_scheduler_process_func: Callable = run_scheduler_process,
|
|
):
|
|
setproctitle.setproctitle("sglang::data_parallel_controller")
|
|
faulthandler.enable()
|
|
kill_itself_when_parent_died()
|
|
parent_process = psutil.Process().parent()
|
|
|
|
configure_logger(server_args)
|
|
if server_args.enable_trace:
|
|
process_tracing_init(server_args.otlp_traces_endpoint, "sglang")
|
|
thread_label = "DP Controller"
|
|
if server_args.disaggregation_mode == "prefill":
|
|
thread_label = "Prefill DP Controller"
|
|
elif server_args.disaggregation_mode == "decode":
|
|
thread_label = "Decode DP Controller"
|
|
trace_set_thread_info(thread_label)
|
|
|
|
try:
|
|
controller = DataParallelController(
|
|
server_args, port_args, run_scheduler_process_func
|
|
)
|
|
pipe_writer.send(
|
|
{
|
|
"status": "ready",
|
|
"max_total_num_tokens": controller.max_total_num_tokens,
|
|
"max_req_input_len": controller.max_req_input_len,
|
|
}
|
|
)
|
|
if server_args.node_rank == 0:
|
|
controller.event_loop()
|
|
for proc in controller.scheduler_procs:
|
|
proc.join()
|
|
logger.error(
|
|
f"Scheduler or DataParallelController {proc.pid} terminated with {proc.exitcode}"
|
|
)
|
|
except Exception:
|
|
traceback = get_exception_traceback()
|
|
logger.error(f"DataParallelController hit an exception: {traceback}")
|
|
parent_process.send_signal(signal.SIGQUIT)
|