Make text generation server hostname configurable

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:

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)

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(

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,14 +119,15 @@ 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()
     asyncio.set_event_loop(loop)
     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()

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