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

[PD] Tiny code cleanup for prefill info registering (#19414)

Browse Source 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
2026-02-27 01:35:26 +08:00
committed by GitHub
parent e4b708d3e9
commit 288300aafd
2 changed files with 18 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/disaggregation/base/conn.py
View File

@@ -64,7 +64,7 @@ class BaseKVManager(ABC):
@abstractmethod
def register_to_bootstrap(self):
"""Register to the bootstrap server."""
"""Register prefill server info to the bootstrap server."""
...

36
python/sglang/srt/disaggregation/common/conn.py
View File

@@ -226,7 +226,7 @@ class CommonKVManager(BaseKVManager):
return None
def register_to_bootstrap(self):
"""Register KVSender to bootstrap server via HTTP POST."""
"""Register prefill server info to bootstrap server via HTTP POST."""
if self.dist_init_addr:
# Multi-node case: bootstrap server's host is dist_init_addr
if self.dist_init_addr.startswith("["): # [ipv6]:port or [ipv6]
@@ -244,7 +244,6 @@ class CommonKVManager(BaseKVManager):
bootstrap_server_url = f"{host}:{self.bootstrap_port}"
url = f"http://{bootstrap_server_url}/route"
payload = {
"role": "Prefill",
"attn_tp_size": self.attn_tp_size,
"attn_tp_rank": self.attn_tp_rank,
"attn_dp_size": self.attn_dp_size,
@@ -666,7 +665,6 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
async def _handle_route_put(self, request: web.Request):
data = await request.json()
role = data["role"]
attn_tp_size = data["attn_tp_size"]
attn_tp_rank = data["attn_tp_rank"]
attn_dp_size = data["attn_dp_size"]
@@ -701,29 +699,29 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
)
self.follow_bootstrap_room = load_balance_method == "follow_bootstrap_room"
if role == "Prefill":
if system_dp_size == 1:
dp_group = attn_dp_rank
else:
dp_group = system_dp_rank
if system_dp_size == 1:
dp_group = attn_dp_rank
else:
dp_group = system_dp_rank
# Add lock to make sure thread-safe
async with self.lock:
if dp_group not in self.prefill_port_table:
self.prefill_port_table[dp_group] = {}
if attn_tp_rank not in self.prefill_port_table[dp_group]:
self.prefill_port_table[dp_group][attn_tp_rank] = {}
# Add lock to make sure thread-safe
async with self.lock:
if dp_group not in self.prefill_port_table:
self.prefill_port_table[dp_group] = {}
if attn_tp_rank not in self.prefill_port_table[dp_group]:
self.prefill_port_table[dp_group][attn_tp_rank] = {}
self.prefill_port_table[dp_group][attn_tp_rank][pp_rank] = {
"rank_ip": rank_ip,
"rank_port": rank_port,
}
self._registered_count += 1
expected = self.dp_size * self.attn_tp_size * self.pp_size
logger.debug(
f"Register prefill bootstrap: DP{dp_group} TP{attn_tp_rank} PP{pp_rank} with rank_ip: {rank_ip} and rank_port: {rank_port}"
f" ({self._registered_count}/{expected} registered)"
)
expected = self.dp_size * self.attn_tp_size * self.pp_size
logger.debug(
f"Register prefill bootstrap: DP{dp_group} TP{attn_tp_rank} PP{pp_rank} with rank_ip: {rank_ip} and rank_port: {rank_port}"
f" ({self._registered_count}/{expected} registered)"
)
return web.Response(text="OK", status=200)