diff --git a/megatron/core/inference/data_parallel_inference_coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator.py index 146ecf1f1dc..fab9b4acfe6 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator.py +++ b/megatron/core/inference/data_parallel_inference_coordinator.py @@ -95,6 +95,7 @@ def __init__( PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK ), schedule_output_path: str | None = None, + hostname: str | None = None, ): """ Initializes the inference coordinator. @@ -130,7 +131,7 @@ def __init__( # the user that had submitted the request originally. # Get local IP. - local_ip = socket.gethostname() + local_ip = hostname or socket.gethostname() self.router_socket = self.context.socket(zmq.ROUTER) # Raise error if the other side of the connection has dropped. @@ -533,6 +534,7 @@ def entrypoint( PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK ), schedule_output_path: str | None = None, + hostname: str | None = None, ): """ Class method to instantiate and run the coordinator, for use in a separate process. @@ -562,6 +564,7 @@ def entrypoint( enable_prefix_caching=enable_prefix_caching, prefix_caching_coordinator_policy=prefix_caching_coordinator_policy, schedule_output_path=schedule_output_path, + hostname=hostname, ) ready_event.set() try: diff --git a/megatron/core/inference/engines/async_zmq_communicator.py b/megatron/core/inference/engines/async_zmq_communicator.py index 6af35be7fb4..52570845d61 100644 --- a/megatron/core/inference/engines/async_zmq_communicator.py +++ b/megatron/core/inference/engines/async_zmq_communicator.py @@ -26,13 +26,20 @@ class AsyncZMQCommunicator: on the CPU. """ - def __init__(self, zmq_context: zmq.Context, process_group: dist.ProcessGroup): + def __init__( + self, + zmq_context: zmq.Context, + process_group: dist.ProcessGroup, + hostname: str | None = None, + ): """ Constructor for AsyncZMQCommunicator. Sets up ZMQ sockets for communication among ranks in the given process group. Args: zmq_context (zmq.Context): ZMQ context to create sockets. process_group (dist.ProcessGroup): Process group for communication. + hostname (str | None): Hostname or IP address to use for ZMQ socket binding. + If None, defaults to socket.gethostname(). """ self.rank = dist.get_rank(process_group) self.world_size = dist.get_world_size(process_group) @@ -41,7 +48,7 @@ def __init__(self, zmq_context: zmq.Context, process_group: dist.ProcessGroup): src_rank = dist.get_process_group_ranks(process_group)[0] if self.is_leader: - local_ip = socket.gethostname() + local_ip = hostname or socket.gethostname() self.gather_sock = zmq_context.socket(zmq.PULL) self.gather_sock.bind_to_random_port(f"tcp://{local_ip}") gather_socket_addr = self.gather_sock.getsockopt_string(zmq.LAST_ENDPOINT) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 534a342174e..4d5562bf39d 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -426,6 +426,7 @@ async def start_listening_to_data_parallel_coordinator( inference_coordinator_port: int | None = None, launch_inference_coordinator: bool = True, *, + hostname: str | None = None, coordinator_schedule_output_path: str | None = None, loop: Optional[asyncio.AbstractEventLoop] = None, ): @@ -463,6 +464,9 @@ async def start_listening_to_data_parallel_coordinator( launch_inference_coordinator (bool, optional): If True, the global rank 0 process will spawn and manage the `InferenceCoordinator` process. Defaults to True. + hostname (str | None): Hostname or IP address to use for ZMQ socket binding. + If None, defaults to `socket.gethostname()`. Should be set to a routable + address in multi-node settings where gethostname() may return 127.0.0.1. Returns: inference_coordinator_addresss (str): The network address of the central @@ -495,7 +499,7 @@ async def start_listening_to_data_parallel_coordinator( self.is_mp_coordinator = tp_rank == 0 and pp_rank == 0 self.is_dp_coordinator = (dp_rank == 0) and self.is_mp_coordinator - local_ip = socket.gethostname() + local_ip = hostname or socket.gethostname() # Spawn a DP coordinator process and get the connection info. if launch_inference_coordinator and self.is_dp_coordinator: @@ -516,6 +520,7 @@ async def start_listening_to_data_parallel_coordinator( self.context.enable_prefix_caching, self.context.prefix_caching_coordinator_policy, coordinator_schedule_output_path, + hostname, ), ) self.inference_coordinator_process.start() @@ -601,13 +606,15 @@ async def start_listening_to_data_parallel_coordinator( self.ep_world_size = get_pg_size(self.pg_collection.ep) if self.ep_world_size > 1: self.expert_parallel_zmq_communicator = AsyncZMQCommunicator( - self.zmq_context, process_group=self.pg_collection.ep + self.zmq_context, process_group=self.pg_collection.ep, hostname=hostname ) # initialize zmq-based world communicator for consensus barriers total_world_size = torch.distributed.get_world_size() if total_world_size > 1: - self.world_zmq_communicator = AsyncZMQCommunicator(self.zmq_context, process_group=None) + self.world_zmq_communicator = AsyncZMQCommunicator( + self.zmq_context, process_group=None, hostname=hostname + ) if launch_inference_coordinator and self.is_dp_coordinator: await await_process_call( diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py index 77ef410f538..8f91644afdd 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py @@ -48,6 +48,7 @@ async def _run_text_gen_server( parsers: Optional[List[str]] = None, verbose: bool = False, fd: Optional[int] = None, + hostname: Optional[str] = None, ): """ Initializes and runs the async web server. Automatically starts and @@ -62,11 +63,12 @@ async def _run_text_gen_server( logger.info(f"Rank {rank}: InferenceClient connected.") try: - try: - hostname = socket.gethostname() - except Exception as e: - logger.warning(f"Could not get hostname: {e}") - hostname = "0.0.0.0" + if hostname is None: + try: + hostname = socket.gethostname() + except Exception as e: + logger.warning(f"Could not get hostname: {e}") + hostname = "0.0.0.0" app = Quart(__name__) @@ -93,7 +95,7 @@ async def _run_text_gen_server( if fd is not None: config.bind = [f"fd://{fd}"] else: - config.bind = [f"0.0.0.0:{server_port}"] + config.bind = [f"{hostname}:{server_port}"] with temp_log_level(logging.INFO, logger): logger.info(f"Starting text generation server on http://{hostname}:{server_port}") @@ -117,6 +119,7 @@ def _server_process_worker( parsers: Optional[List[str]] = None, verbose: bool = False, fd: Optional[int] = None, + hostname: Optional[str] = None, ): """Synchronous worker function that sets up a new event loop for the separate process.""" loop = asyncio.new_event_loop() @@ -124,7 +127,7 @@ def _server_process_worker( try: loop.run_until_complete( _run_text_gen_server( - coordinator_addr, tokenizer, rank, server_port, parsers, verbose, fd + coordinator_addr, tokenizer, rank, server_port, parsers, verbose, fd, hostname ) ) except KeyboardInterrupt: @@ -146,6 +149,7 @@ def start_text_gen_server( parsers: Optional[List[str]] = None, verbose: bool = False, num_replicas: int = 4, + hostname: Optional[str] = None, ): """Start the text generation server.""" global _SERVER_PROCESSES @@ -164,7 +168,8 @@ def start_text_gen_server( except OSError: pass - _SHARED_SOCKET.bind(("0.0.0.0", server_port)) + bind_address = hostname if hostname is not None else "0.0.0.0" + _SHARED_SOCKET.bind((bind_address, server_port)) _SHARED_SOCKET.setblocking(False) _SHARED_SOCKET.set_inheritable(True) @@ -173,7 +178,7 @@ def start_text_gen_server( for i in range(num_replicas): p = mp.Process( target=_server_process_worker, - args=(coordinator_addr, tokenizer, rank, server_port, parsers, verbose, fd), + args=(coordinator_addr, tokenizer, rank, server_port, parsers, verbose, fd, hostname), daemon=True, ) p.start() diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py index 1f99939a184..74fbfd8e3d3 100644 --- a/tools/run_dynamic_text_generation_server.py +++ b/tools/run_dynamic_text_generation_server.py @@ -22,6 +22,10 @@ def add_text_generation_server_args(parser: argparse.ArgumentParser): parser = add_modelopt_args(parser) parser = add_inference_args(parser) parser.add_argument("--port", type=int, default=5000, help="Port for Flask server to run on") + parser.add_argument( + "--host", type=str, default=None, + help="Hostname or IP address to bind the server to. Defaults to 0.0.0.0 (all interfaces)." + ) parser.add_argument( "--parsers", type=str, nargs="+", default=[], help="Parsers to use for parsing the response" ) @@ -30,7 +34,7 @@ def add_text_generation_server_args(parser: argparse.ArgumentParser): @trace_async_exceptions async def run_text_generation_server( - engine: DynamicInferenceEngine, coordinator_port: int, server_port: int + engine: DynamicInferenceEngine, coordinator_port: int, server_port: int, hostname: str | None = None, ): """ Runs the text generation server from rank 0 and initializes the @@ -45,7 +49,8 @@ async def run_text_generation_server( rank = torch.distributed.get_rank() coordinator_addr = await engine.start_listening_to_data_parallel_coordinator( - inference_coordinator_port=coordinator_port, launch_inference_coordinator=True + inference_coordinator_port=coordinator_port, launch_inference_coordinator=True, + hostname=hostname, ) try: @@ -57,6 +62,7 @@ async def run_text_generation_server( rank=rank, server_port=server_port, verbose=args.inference_text_gen_server_logging, + hostname=hostname, ) # Await the engine loop directly since the server is running in a separate process @@ -85,7 +91,7 @@ async def run_text_generation_server( try: asyncio.run( - run_text_generation_server(engine, args.inference_coordinator_port, args.port) + run_text_generation_server(engine, args.inference_coordinator_port, args.port, args.host) ) except KeyboardInterrupt: # Catching at the top level ensures clean stdout without spamming the traceback