[PD] Fix multi-tokenizer disaggregation metrics labels (#30412)
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
This commit is contained in:
@@ -36,7 +36,7 @@ import setproctitle
|
|||||||
import zmq
|
import zmq
|
||||||
import zmq.asyncio
|
import zmq.asyncio
|
||||||
|
|
||||||
from sglang.srt.disaggregation.utils import DisaggregationMode, TransferBackend
|
from sglang.srt.disaggregation.utils import TransferBackend
|
||||||
from sglang.srt.managers.disagg_service import start_disagg_service
|
from sglang.srt.managers.disagg_service import start_disagg_service
|
||||||
from sglang.srt.managers.io_struct import (
|
from sglang.srt.managers.io_struct import (
|
||||||
BaseBatchReq,
|
BaseBatchReq,
|
||||||
@@ -634,24 +634,16 @@ class TokenizerWorker(TokenizerManager):
|
|||||||
import torch
|
import torch
|
||||||
|
|
||||||
torch.set_num_threads(1)
|
torch.set_num_threads(1)
|
||||||
# prevent init prefill bootstrapserver again
|
super().__init__(
|
||||||
disaggregation_mode = server_args.disaggregation_mode
|
server_args,
|
||||||
server_args.override(
|
port_args,
|
||||||
"tokenizer_worker.suppress_bootstrap", disaggregation_mode="null"
|
start_pd_bootstrap_service=False,
|
||||||
)
|
)
|
||||||
super().__init__(server_args, port_args)
|
|
||||||
|
|
||||||
self.worker_id = os.getpid()
|
self.worker_id = os.getpid()
|
||||||
self.tokenizer_ipc_name = port_args.tokenizer_ipc_name
|
self.tokenizer_ipc_name = port_args.tokenizer_ipc_name
|
||||||
|
|
||||||
# For PD disaggregtion
|
# For PD disaggregation
|
||||||
self.server_args.override(
|
|
||||||
"tokenizer_worker.restore_disaggregation_mode",
|
|
||||||
disaggregation_mode=disaggregation_mode,
|
|
||||||
)
|
|
||||||
self.disaggregation_mode = DisaggregationMode(
|
|
||||||
self.server_args.disaggregation_mode
|
|
||||||
)
|
|
||||||
self.disaggregation_transfer_backend = TransferBackend(
|
self.disaggregation_transfer_backend = TransferBackend(
|
||||||
self.server_args.disaggregation_transfer_backend
|
self.server_args.disaggregation_transfer_backend
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -279,6 +279,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
|||||||
self,
|
self,
|
||||||
server_args: ServerArgs,
|
server_args: ServerArgs,
|
||||||
port_args: PortArgs,
|
port_args: PortArgs,
|
||||||
|
*,
|
||||||
|
start_pd_bootstrap_service: bool = True,
|
||||||
):
|
):
|
||||||
# Parse args
|
# Parse args
|
||||||
self.server_args = server_args
|
self.server_args = server_args
|
||||||
@@ -318,7 +320,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
|||||||
self.init_lora()
|
self.init_lora()
|
||||||
|
|
||||||
# Init PD disaggregation and encoder disaggregation
|
# Init PD disaggregation and encoder disaggregation
|
||||||
self.init_disaggregation()
|
self.init_disaggregation(start_pd_bootstrap_service=start_pd_bootstrap_service)
|
||||||
|
|
||||||
# Init metric collector and watchdog
|
# Init metric collector and watchdog
|
||||||
self.init_metric_collector_watchdog()
|
self.init_metric_collector_watchdog()
|
||||||
@@ -522,13 +524,17 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
|||||||
for lora_ref in self.server_args.lora_paths:
|
for lora_ref in self.server_args.lora_paths:
|
||||||
self.lora_ref_cache[lora_ref.lora_name] = lora_ref
|
self.lora_ref_cache[lora_ref.lora_name] = lora_ref
|
||||||
|
|
||||||
def init_disaggregation(self):
|
def init_disaggregation(self, *, start_pd_bootstrap_service: bool = True):
|
||||||
# PD Disaggregation
|
# PD Disaggregation
|
||||||
self.disaggregation_mode = DisaggregationMode(
|
self.disaggregation_mode = DisaggregationMode(
|
||||||
self.server_args.disaggregation_mode
|
self.server_args.disaggregation_mode
|
||||||
)
|
)
|
||||||
# Keep a reference so the bootstrap server is not garbage-collected.
|
# Keep a reference so the bootstrap server is not garbage-collected.
|
||||||
self.bootstrap_server = start_disagg_service(self.server_args)
|
self.bootstrap_server = (
|
||||||
|
start_disagg_service(self.server_args)
|
||||||
|
if start_pd_bootstrap_service
|
||||||
|
else None
|
||||||
|
)
|
||||||
# Single-source counter for auto-assigning fake bootstrap_room.
|
# Single-source counter for auto-assigning fake bootstrap_room.
|
||||||
self.fake_bootstrap_room_counter = 0
|
self.fake_bootstrap_room_counter = 0
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user