Set up the idle sleeper outside of the IPC channel initialization (#25713)
This commit is contained in:
@@ -384,6 +384,7 @@ class Scheduler(
|
|||||||
|
|
||||||
# Init inter-process communication
|
# Init inter-process communication
|
||||||
self.init_ipc_channels(port_args)
|
self.init_ipc_channels(port_args)
|
||||||
|
self.init_idle_sleeper()
|
||||||
|
|
||||||
self.mm_receiver = None
|
self.mm_receiver = None
|
||||||
self.disagg_prefill_bootstrap_queue = None
|
self.disagg_prefill_bootstrap_queue = None
|
||||||
@@ -726,7 +727,6 @@ class Scheduler(
|
|||||||
|
|
||||||
def init_ipc_channels(self, port_args: PortArgs):
|
def init_ipc_channels(self, port_args: PortArgs):
|
||||||
context = zmq.Context(2)
|
context = zmq.Context(2)
|
||||||
self.idle_sleeper = None
|
|
||||||
self.send_metrics_from_scheduler = None
|
self.send_metrics_from_scheduler = None
|
||||||
|
|
||||||
if (
|
if (
|
||||||
@@ -757,14 +757,6 @@ class Scheduler(
|
|||||||
|
|
||||||
self.send_to_tokenizer = SenderWrapper(send_to_tokenizer)
|
self.send_to_tokenizer = SenderWrapper(send_to_tokenizer)
|
||||||
self.send_to_detokenizer = SenderWrapper(send_to_detokenizer)
|
self.send_to_detokenizer = SenderWrapper(send_to_detokenizer)
|
||||||
|
|
||||||
if self.server_args.sleep_on_idle:
|
|
||||||
self.idle_sleeper = IdleSleeper(
|
|
||||||
[
|
|
||||||
self.recv_from_tokenizer,
|
|
||||||
self.recv_from_rpc,
|
|
||||||
]
|
|
||||||
)
|
|
||||||
else:
|
else:
|
||||||
self.recv_from_tokenizer = None
|
self.recv_from_tokenizer = None
|
||||||
self.recv_from_rpc = None
|
self.recv_from_rpc = None
|
||||||
@@ -779,6 +771,22 @@ class Scheduler(
|
|||||||
context, zmq.PUSH, port_args.metrics_ipc_name, False
|
context, zmq.PUSH, port_args.metrics_ipc_name, False
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def init_idle_sleeper(self) -> None:
|
||||||
|
if (
|
||||||
|
self.ps.pp_rank == 0
|
||||||
|
and self.ps.attn_tp_rank == 0
|
||||||
|
and self.ps.attn_cp_rank == 0
|
||||||
|
and self.server_args.sleep_on_idle
|
||||||
|
):
|
||||||
|
self.idle_sleeper = IdleSleeper(
|
||||||
|
sockets=[
|
||||||
|
self.recv_from_tokenizer,
|
||||||
|
self.recv_from_rpc,
|
||||||
|
],
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.idle_sleeper = None
|
||||||
|
|
||||||
def init_tokenizer(self):
|
def init_tokenizer(self):
|
||||||
server_args = self.server_args
|
server_args = self.server_args
|
||||||
self.is_generation = self.model_config.is_generation
|
self.is_generation = self.model_config.is_generation
|
||||||
|
|||||||
Reference in New Issue
Block a user