Use ipc instead of tcp in zmq (#1566)

This commit is contained in:
Lianmin Zheng
2024-10-04 00:45:52 -07:00
committed by GitHub
parent 32eb6e96f2
commit 114bbc8651
9 changed files with 48 additions and 96 deletions
@@ -59,10 +59,10 @@ class DetokenizerManager:
# Init inter-process communication
context = zmq.Context(2)
self.recv_from_scheduler = context.socket(zmq.PULL)
self.recv_from_scheduler.bind(f"tcp://127.0.0.1:{port_args.detokenizer_port}")
self.recv_from_scheduler.bind(f"ipc://{port_args.detokenizer_ipc_name}")
self.send_to_tokenizer = context.socket(zmq.PUSH)
self.send_to_tokenizer.connect(f"tcp://127.0.0.1:{port_args.tokenizer_port}")
self.send_to_tokenizer.connect(f"ipc://{port_args.tokenizer_ipc_name}")
if server_args.skip_tokenizer_init:
self.tokenizer = None
+2 -6
View File
@@ -96,14 +96,10 @@ class Scheduler:
if self.tp_rank == 0:
self.recv_from_tokenizer = context.socket(zmq.PULL)
self.recv_from_tokenizer.bind(
f"tcp://127.0.0.1:{port_args.scheduler_input_port}"
)
self.recv_from_tokenizer.bind(f"ipc://{port_args.scheduler_input_ipc_name}")
self.send_to_detokenizer = context.socket(zmq.PUSH)
self.send_to_detokenizer.connect(
f"tcp://127.0.0.1:{port_args.detokenizer_port}"
)
self.send_to_detokenizer.connect(f"ipc://{port_args.detokenizer_ipc_name}")
else:
self.recv_from_tokenizer = self.send_to_detokenizer = None
@@ -84,12 +84,10 @@ class TokenizerManager:
# Init inter-process communication
context = zmq.asyncio.Context(2)
self.recv_from_detokenizer = context.socket(zmq.PULL)
self.recv_from_detokenizer.bind(f"tcp://127.0.0.1:{port_args.tokenizer_port}")
self.recv_from_detokenizer.bind(f"ipc://{port_args.tokenizer_ipc_name}")
self.send_to_scheduler = context.socket(zmq.PUSH)
self.send_to_scheduler.connect(
f"tcp://127.0.0.1:{port_args.scheduler_input_port}"
)
self.send_to_scheduler.connect(f"ipc://{port_args.scheduler_input_ipc_name}")
# Read model args
self.model_path = server_args.model_path