From d451807f02148e8cead61095b4ebae8b5faf678f Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 11 Jun 2026 14:07:48 -0500 Subject: [PATCH 1/2] Allow for pre-bound socket to be passed in server Signed-off-by: Teodor-Dumitru Ene --- .../text_generation_server.py | 35 ++++++++++++------- 1 file changed, 23 insertions(+), 12 deletions(-) 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 8f91644afdd..25619532284 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 @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. import asyncio import logging @@ -150,6 +150,7 @@ def start_text_gen_server( verbose: bool = False, num_replicas: int = 4, hostname: Optional[str] = None, + sock: Optional[socket.socket] = None, ): """Start the text generation server.""" global _SERVER_PROCESSES @@ -159,18 +160,28 @@ def start_text_gen_server( logger.warning("Text gen server processes are already running.") return - _SHARED_SOCKET = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - - if hasattr(socket, 'SO_REUSEPORT'): - try: - _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) - except OSError: - pass + if sock is not None: + bound_port = sock.getsockname()[1] + if bound_port == 0: + raise ValueError( + "socket must be bound to a real port before being passed to start_text_gen_server" + ) + _SHARED_SOCKET = sock + server_port = bound_port + _SHARED_SOCKET.setblocking(False) + else: + _SHARED_SOCKET = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + + if hasattr(socket, 'SO_REUSEPORT'): + try: + _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) + except OSError: + pass - bind_address = hostname if hostname is not None else "0.0.0.0" - _SHARED_SOCKET.bind((bind_address, server_port)) - _SHARED_SOCKET.setblocking(False) + 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) fd = _SHARED_SOCKET.fileno() From 9c69e5186258a5ea58cfcfd59b4068a74b69dd30 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 11 Jun 2026 17:35:00 -0500 Subject: [PATCH 2/2] Address reviewer comments Signed-off-by: Teodor-Dumitru Ene --- .../dynamic_text_gen_server/text_generation_server.py | 1 + 1 file changed, 1 insertion(+) 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 25619532284..af12bc5bbc1 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 @@ -160,6 +160,7 @@ def start_text_gen_server( logger.warning("Text gen server processes are already running.") return + # The caller may pass in a socket it has already bound ahead of time. if sock is not None: bound_port = sock.getsockname()[1] if bound_port == 0: