[https://nvbugs/6143883][fix] Preserve ip:port for trtllm-serve visual-gen (#14355)

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>

Author: JunyiXu-nv

tensorrt_llm/commands/serve.py

@@ -560,23 +560,40 @@ def launch_visual_gen_server(
         visual_gen_args: Optional validated VisualGenArgs for model configuration.
         metadata_server_cfg: Optional metadata server configuration.
     """
-    logger.info(f"Initializing VisualGen ({model})")
+    # Reserve the listening (host, port) by binding the socket *before*
+    # constructing the VisualGen pipeline, then hand the bound socket to
+    # uvicorn. VisualGen initialization can take many minutes; if we deferred
+    # the bind until uvicorn started, anything else on the host could grab the
+    # port in that window and trtllm-serve would die at bind() time.
+    addr_info = socket.getaddrinfo(host, port, socket.AF_UNSPEC,
+                                   socket.SOCK_STREAM)
+    address_family = socket.AF_INET6 if all(
+        [info[0] == socket.AF_INET6 for info in addr_info]) else socket.AF_INET
+    with socket.socket(address_family, socket.SOCK_STREAM) as s:
+        try:
+            s.bind((host, port))
+        except OSError as e:
+            raise RuntimeError(f"Failed to bind socket to {host}:{port}: {e}")
 
-    visual_gen_model = VisualGen(model=model, args=visual_gen_args)
+        logger.info(f"Initializing VisualGen ({model})")
 
-    n_workers = visual_gen_model.args.parallel_config.n_workers
-    logger.info(f"World size: {n_workers}")
-    logger.info(f"CFG size: {visual_gen_model.args.parallel_config.cfg_size}")
-    logger.info(
-        f"Ulysses size: {visual_gen_model.args.parallel_config.ulysses_size}")
+        visual_gen_model = VisualGen(model=model, args=visual_gen_args)
 
-    server = OpenAIServer(generator=visual_gen_model,
-                          model=model,
-                          server_role=ServerRole.VISUAL_GEN,
-                          metadata_server_cfg=metadata_server_cfg,
-                          tool_parser=None)
-    _apply_fastapi_middlewares(server.app, middleware)
-    asyncio.run(server(host, port))
+        n_workers = visual_gen_model.args.parallel_config.n_workers
+        logger.info(f"World size: {n_workers}")
+        logger.info(
+            f"CFG size: {visual_gen_model.args.parallel_config.cfg_size}")
+        logger.info(
+            f"Ulysses size: {visual_gen_model.args.parallel_config.ulysses_size}"
+        )
+
+        server = OpenAIServer(generator=visual_gen_model,
+                              model=model,
+                              server_role=ServerRole.VISUAL_GEN,
+                              metadata_server_cfg=metadata_server_cfg,
+                              tool_parser=None)
+        _apply_fastapi_middlewares(server.app, middleware)
+        asyncio.run(server(host, port, sockets=[s]))
 
 
 class ChoiceWithAlias(click.Choice):

tests/unittest/_torch/visual_gen/test_trtllm_serve_e2e.py

@@ -40,8 +40,6 @@
 import requests
 import yaml
 
-from tensorrt_llm._utils import get_free_port
-
 # ---------------------------------------------------------------------------
 # Model paths
 # ---------------------------------------------------------------------------
@@ -69,6 +67,14 @@ def _llm_models_root() -> str:
 _PROJECT_ROOT = Path(__file__).resolve().parents[4]  # repo root
 _REF_IMAGE_PATH = _PROJECT_ROOT / "examples" / "visual_gen" / "cat_piano.png"
 
+# Use the CI-aware port allocator from tests/integration/defs/common.py so
+# parallel pytest sessions on the same OCI node fall into disjoint port
+# sections (CONTAINER_PORT_START / CONTAINER_PORT_NUM). It transparently falls
+# back to the plain free-port scan when those env vars are not set.
+_INTEGRATION_TESTS_DIR = _PROJECT_ROOT / "tests" / "integration"
+if str(_INTEGRATION_TESTS_DIR) not in sys.path:
+    sys.path.insert(0, str(_INTEGRATION_TESTS_DIR))
+from defs.common import get_free_port_in_ci  # noqa: E402
 
 # ---------------------------------------------------------------------------
 # Remote server helper (follows RemoteOpenAIServer pattern)
@@ -94,7 +100,7 @@ def __init__(
         env: Optional[dict] = None,
     ) -> None:
         self.host = host
-        self.port = port if port is not None else get_free_port()
+        self.port = port if port is not None else get_free_port_in_ci()
         self._config_file: Optional[str] = None
         self.proc: Optional[subprocess.Popen] = None