Make text generation server hostname configurable

megatron/core/inference/engines/dynamic_engine.py

@@ -453,6 +453,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,
     ):
@@ -490,6 +491,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
@@ -522,7 +526,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:

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