diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 026f838cbf99..7becb894ce86 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1,5 +1,6 @@ import asyncio import atexit +import contextlib import gc import importlib import inspect @@ -11,6 +12,7 @@ import socket import subprocess # nosec B404 import sys +import tempfile import time import uuid from importlib.util import find_spec @@ -347,6 +349,54 @@ def _build_llm_args_from_disagg_server_cfg(other_args: Dict) -> Dict: return update_llm_args_with_extra_dict(llm_args, llm_args_extra_dict) +def _publish_bound_address(report_addr: Optional[str], host: str, + port: int) -> None: + """Write the address this server actually bound to ``report_addr``. + + Lets a launcher pass ``--port 0`` and learn the kernel-assigned port + afterwards, instead of picking a port up front and racing whoever grabs it + before this process binds. The write is atomic (temp file in the same + directory, then rename) so a reader never observes a partial line, which + matters on the shared filesystems multi-node tests coordinate through. + + The caller is responsible for making the path unique per run: a stale file + from an earlier run points at a dead server, which fails far less obviously + than a port conflict. + + A wildcard bind host is replaced by this machine's hostname, since readers + use the published value as a URL authority and cannot dial 0.0.0.0 or ::. + """ + if not report_addr: + return + if host in ("0.0.0.0", "::", ""): # nosec B104 - reporting, not binding + resolved = socket.gethostname() + logger.info(f"Reporting hostname {resolved} instead of wildcard bind " + f"address {host!r}, which a reader cannot dial") + host = resolved + report_addr = os.path.abspath(report_addr) + parent = os.path.dirname(report_addr) + if parent: + os.makedirs(parent, exist_ok=True) + fd, tmp_path = tempfile.mkstemp(dir=parent or None, + prefix=os.path.basename(report_addr) + ".", + suffix=".tmp") + try: + with os.fdopen(fd, "w") as f: + # Bracket IPv6 literals so the value is a usable URL authority: + # readers build "http:///..." from it verbatim. + reported_host = f"[{host}]" if ":" in host else host + f.write(f"{reported_host}:{port}\n") + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, report_addr) + finally: + # A successful replace already consumed tmp_path; this only cleans up + # after a failed write so a partial file is not left behind. + with contextlib.suppress(OSError): + os.unlink(tmp_path) + logger.info(f"Reported bound address {host}:{port} to {report_addr}") + + def _diagnose_port_in_use(port: int) -> str: """Describe which process currently holds the given port, best effort.""" try: @@ -544,12 +594,23 @@ def launch_server( num_input_processor_workers: int = 8, num_media_load_workers: int = 8, multi_frontend_enabled: bool = True, - internal_disagg_auth_key: Optional[str] = None): + internal_disagg_auth_key: Optional[str] = None, + report_addr: Optional[str] = None): backend = llm_args["backend"] model = served_model_name or llm_args["model"] multi_frontend = _init_multi_frontend_mode(llm_args, multi_frontend_enabled) + # Same hazard the disaggregated fleet guard covers: _spawn_attached_frontends + # re-execs this command line verbatim, so with port 0 every frontend binds + # its own kernel-assigned port instead of sharing one, and every frontend + # also re-runs _publish_bound_address, leaving the reader with whichever + # child wrote last. + if (port == 0 or report_addr) and multi_frontend.num_frontends > 1: + raise click.BadParameter( + "port 0 and --report_addr are only supported with a single serving " + f"frontend, but num_serve_frontends={multi_frontend.num_frontends}." + ) if multi_frontend.is_launcher or multi_frontend.is_attached_frontend: # The Responses API store is per-process in-memory: with several # frontends behind one SO_REUSEPORT port, a follow-up request may @@ -567,8 +628,17 @@ def launch_server( 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: - # If disagg cluster config is provided and port is not specified, try to find a free port, otherwise try to bind to the specified port - assert port > 0 or disagg_cluster_config is not None, "Port must be specified if disagg cluster config is not provided" + # port == 0 lets the kernel pick the port; the caller then needs a way + # to learn it, either by service discovery or by report_addr. + assert port > 0 or disagg_cluster_config is not None or report_addr, ( + "Port must be specified unless disagg cluster config or " + "--report_addr is provided") + # Without SO_REUSEADDR a restart is refused for the whole TIME_WAIT + # window (~60s) by the tombstones of connections this server accepted. + # The flag has to be set on the socket that owns the port first, since + # the TIME_WAIT entry inherits it -- setting it only on the later bind + # is not enough. + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) if multi_frontend.is_launcher or multi_frontend.is_attached_frontend: # Every frontend process binds its own listening socket on the # same port; the kernel load-balances accepts across them. @@ -585,6 +655,10 @@ def launch_server( raise RuntimeError(f"Failed to bind socket to {host}:{port}: {e}. " f"Port holder(s): {holder}") + # Only now is the address final, and the socket stays bound from here + # until uvicorn takes it over, so no one can steal the port in between. + _publish_bound_address(report_addr, host, port) + if backend == 'pytorch': llm_args.pop("build_config", None) llm = PyTorchLLM(**llm_args) @@ -795,6 +869,9 @@ def launch_visual_gen_server( 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: + # See launch_server: without this, TIME_WAIT tombstones from the + # connections this server accepted refuse a restart for ~60s. + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) try: s.bind((host, port)) except OSError as e: @@ -1167,6 +1244,15 @@ def launch_visual_gen_server( help= "Types of agents to schedule. Now Only Support Open Deep Research agent.", status="prototype") +@stability_option( + "--report_addr", + type=str, + default=None, + help="Write the host:port this server actually bound to this file, " + "atomically, once the socket is bound. Lets --port 0 be used and have the " + "launcher read the kernel-assigned port back instead of reserving one up " + "front.", + status="prototype") def serve(model: str, tokenizer: Optional[str], custom_tokenizer: Optional[str], post_processor_hook: Optional[str], host: str, port: int, log_level: str, backend: str, generation_config: str, @@ -1191,7 +1277,8 @@ def serve(model: str, tokenizer: Optional[str], custom_tokenizer: Optional[str], telemetry: bool, custom_module_dirs: list[Path], chat_template: Optional[str], allow_request_chat_template: bool, middleware: tuple[str, ...], grpc: bool, enable_visual_gen: bool, - served_model_name: Optional[str], visual_gen_args: Optional[str]): + served_model_name: Optional[str], visual_gen_args: Optional[str], + report_addr: Optional[str]): """Running an OpenAI API compatible server MODEL: model name | HF checkpoint path | TensorRT engine path @@ -1408,7 +1495,8 @@ def _serve_llm(): allow_request_chat_template=allow_request_chat_template, num_input_processor_workers=num_input_processor_workers, num_media_load_workers=num_media_load_workers, - internal_disagg_auth_key=internal_disagg_auth_key) + internal_disagg_auth_key=internal_disagg_auth_key, + report_addr=report_addr) def _serve_visual_gen(): from tensorrt_llm.visual_gen.args import VisualGenArgs @@ -1424,6 +1512,12 @@ def _serve_visual_gen(): is_visual_gen = (enable_visual_gen or visual_gen_args is not None or get_is_diffusion_only_model(model)) + # Only the OpenAI HTTP path publishes the bound address. Fail loudly rather + # than leaving a launcher waiting forever on a file nobody writes. + if report_addr and (grpc or is_visual_gen): + raise click.BadParameter( + "--report_addr is only supported for the OpenAI HTTP server, not " + f"the {'gRPC' if grpc else 'VisualGen'} server.") if is_visual_gen: _serve_visual_gen() else: @@ -1750,6 +1844,15 @@ def serve_embedding( help="[Deprecated] The interval of logging metrics in seconds. " "This option is not connected to any functionality and will be removed in a future release.", status="deprecated") +@stability_option( + "--report_addr", + type=str, + default=None, + help="Write the host:port this server actually bound to this file, " + "atomically, once the socket is bound. Lets the config set port 0 and " + "have the launcher read the kernel-assigned port back instead of " + "reserving one up front.", + status="prototype") def disaggregated( config_file: Optional[str], metadata_server_config_file: Optional[str], @@ -1758,6 +1861,7 @@ def disaggregated( log_level: str, metrics_log_interval: int, schedule_style: str, + report_addr: Optional[str], ): """Running server in disaggregated mode""" @@ -1792,6 +1896,17 @@ def disaggregated( num_workers = disagg_cfg.num_workers coordinator_url = disagg_cfg.disagg_coordinator_url + # Only topology (c) below binds the public socket in this process. The fleet + # paths hand the port to N SO_REUSEPORT workers, which with port 0 would each + # get a *different* kernel-assigned port instead of sharing one, so reject + # the combination rather than publishing an address that serves 1/N requests. + if (disagg_cfg.port == 0 or report_addr) and (coordinator_url + or num_workers > 1): + raise click.BadParameter( + "port 0 and --report_addr are only supported for a single " + f"self-contained disaggregated server, but num_workers={num_workers} " + f"and disagg_coordinator_url={coordinator_url!r} select a fleet.") + if coordinator_url: # (a) External coordinator: fork a fleet of delegating servers (or a # single one) pointed at it; never start a coordinator in this process. @@ -1812,8 +1927,13 @@ def disaggregated( # (c) num_workers==1, no external coordinator: a single disagg server with an # in-process (local) coordinator. Pre-bind the socket (validates port), serve. with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + # See launch_server: without this, TIME_WAIT tombstones from the + # connections this server accepted refuse a restart for ~60s. + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) try: s.bind((disagg_cfg.hostname, disagg_cfg.port)) + if disagg_cfg.port == 0: + disagg_cfg.port = s.getsockname()[1] except OSError as e: holder = _diagnose_port_in_use(disagg_cfg.port) logger.error( @@ -1824,6 +1944,9 @@ def disaggregated( f"Failed to bind socket to {disagg_cfg.hostname}:{disagg_cfg.port}: {e}. " f"Port holder(s): {holder}") + _publish_bound_address(report_addr, disagg_cfg.hostname, + disagg_cfg.port) + server = OpenAIDisaggServer( config=disagg_cfg, req_timeout_secs=request_timeout, diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 0a970516d838..4287c43d7189 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -32,7 +32,7 @@ import pytest import requests import yaml -from defs.common import get_free_port_in_ci as get_free_port +from defs.common import wait_for_reported_addr from tensorrt_llm.executor.result import GenerationResultBase from tensorrt_llm.llmapi import CompletionOutput, RequestOutput, SamplingParams @@ -197,17 +197,18 @@ def _apply_perf_flags(cfg: Optional[Dict[str, Any]]): _apply_perf_flags(ctx_server_config) _apply_perf_flags(gen_server_config) - # Always assign free port dynamically for service discovery - serve_port = get_free_port() - disaggregated_server_config["port"] = serve_port + # Let the kernel assign the port inside trtllm-serve and report it back, + # rather than reserving one here and racing whoever takes it before the + # server binds. + disaggregated_server_config["port"] = 0 + disagg_addr_path = os.path.join(temp_dir.name, "disagg_server.addr") - # Use HTTP service discovery - cluster_uri = f"http://localhost:{serve_port}" - print(f"Using HTTP service discovery at {cluster_uri}") - - # Create service discovery config + # Create service discovery config. The server hosts the HTTP cluster + # storage on its own port, so the port in *its* copy of cluster_uri is + # never read (HttpClusterStorageServer.__init__ ignores the URI); only the + # workers dial it, and they get the resolved address below. disagg_cluster = { - "cluster_uri": cluster_uri, + "cluster_uri": "http://localhost:0", "cluster_name": "test_cluster", "heartbeat_interval_sec": 5, "inactive_timeout_sec": 10, @@ -230,15 +231,14 @@ def _apply_perf_flags(cfg: Optional[Dict[str, Any]]): disaggregated_server_config["internal_request_auth_key"] = ( internal_request_auth_key) - # Inject into worker configs + # Inject into worker configs. disagg_cluster is replaced in + # write_worker_configs below, once the server's real address is known. ctx_server_config = { **ctx_server_config, - "disagg_cluster": disagg_cluster, "internal_request_auth_key": internal_request_auth_key, } gen_server_config = { **gen_server_config, - "disagg_cluster": disagg_cluster, "internal_request_auth_key": internal_request_auth_key, } @@ -246,12 +246,20 @@ def _apply_perf_flags(cfg: Optional[Dict[str, Any]]): yaml.dump(disaggregated_server_config, f) ctx_server_config_path = os.path.join(temp_dir.name, "ctx_server_config.yaml") - with open(ctx_server_config_path, "w") as f: - yaml.dump(ctx_server_config, f) gen_server_config_path = os.path.join(temp_dir.name, "gen_server_config.yaml") - with open(gen_server_config_path, "w") as f: - yaml.dump(gen_server_config, f) + + def write_worker_configs(cluster_uri): + """Write the worker configs once the server's real address is known.""" + worker_cluster = {**disagg_cluster, "cluster_uri": cluster_uri} + with open(ctx_server_config_path, "w") as f: + yaml.dump({ + **ctx_server_config, "disagg_cluster": worker_cluster + }, f) + with open(gen_server_config_path, "w") as f: + yaml.dump({ + **gen_server_config, "disagg_cluster": worker_cluster + }, f) args = LlmArgs(model=model_name, tensor_parallel_size=tensor_parallel_size) @@ -402,15 +410,29 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): server_cmd = [ trtllm_serve_path, "disaggregated", "-c", disaggregated_serving_config_path, "--server_start_timeout", - str(server_waiting_timeout), "-r", "360000" + str(server_waiting_timeout), "-r", "360000", "--report_addr", + disagg_addr_path ] + # The disagg server must come up first: it owns the cluster storage the + # workers register with, and only it knows the port the kernel handed it. with ( MyThreadPoolExecutor(max_workers=max_workers) as thread_pool, temp_dir, - multi_popen(ctx_servers, "ctx") as ctx_processes, - multi_popen(gen_servers, "gen") as gen_processes, - multi_popen([(base_env, server_cmd)], "disagg") as server_processes, + contextlib.ExitStack() as server_stack, ): + server_processes = server_stack.enter_context( + multi_popen([(base_env, server_cmd)], "disagg")) + _, serve_port = wait_for_reported_addr(disagg_addr_path, + server_waiting_timeout, + server_processes[0]) + print(f"Using HTTP service discovery at http://localhost:{serve_port}") + write_worker_configs(f"http://localhost:{serve_port}") + + ctx_processes = server_stack.enter_context( + multi_popen(ctx_servers, "ctx")) + gen_processes = server_stack.enter_context( + multi_popen(gen_servers, "gen")) + start_time = time.time() server_is_ready = False while time.time() - start_time < server_waiting_timeout: diff --git a/tests/integration/defs/common.py b/tests/integration/defs/common.py index f6ddb0d5c7fc..31174b497797 100644 --- a/tests/integration/defs/common.py +++ b/tests/integration/defs/common.py @@ -21,7 +21,7 @@ import tempfile import time from difflib import SequenceMatcher -from typing import Any +from typing import Any, Optional import yaml from packaging import version @@ -671,13 +671,145 @@ def wait_for_server(host, port, timeout_seconds=180): return False +def wait_for_reported_addr(addr_path: str, + timeout: float, + process=None) -> tuple[str, int]: + """Read the address a server reported to its --report_addr file. + + The file only appears once trtllm-serve has bound its socket, and that + socket stays bound from then on, so the address cannot be stolen between + this read and its use -- unlike a port reserved before the server starts. + + Args: + addr_path: Path passed to the server's --report_addr. + timeout: Seconds to wait for the file to appear. + process: Optional Popen of the server, polled so that a crash fails + fast instead of burning the whole timeout. + + Returns: + tuple[str, int]: The host and port the server bound. + """ + deadline = time.time() + timeout + while time.time() < deadline: + if process is not None and process.poll() is not None: + raise RuntimeError( + f"server exited with code {process.returncode} before " + f"reporting its address to {addr_path}") + try: + with open(addr_path) as f: + reported = f.read().strip() + except FileNotFoundError: + reported = "" + if reported: + host, _, port = reported.rpartition(":") + return host, int(port) + time.sleep(0.5) + raise TimeoutError(f"server did not report its address to {addr_path} " + f"within {timeout}s") + + PORTS_IN_USE = set() +# Size of the window carved out just below the kernel's ephemeral range, used +# when CONTAINER_PORT_START is unset (e.g. the SLURM multi-node path). +STATIC_PORT_RANGE_SIZE = 4096 -def get_free_port_in_ci(max_attempts=100): + +def get_ephemeral_port_range() -> Optional[tuple[int, int]]: + """Return the kernel's ephemeral port range as (low, high), or None. + + These are the ports bind(('', 0)) hands out. None means the range could + not be read. """ - Get a free port in the range [CONTAINER_PORT_START, CONTAINER_PORT_START + CONTAINER_PORT_NUM - 1] - If CONTAINER_PORT_START and CONTAINER_PORT_NUM are not set or all ports are already in use, fallback to get_free_port + try: + with open("/proc/sys/net/ipv4/ip_local_port_range") as f: + low, high = (int(value) for value in f.read().split()) + except (OSError, ValueError) as e: + print_info(f"[get_free_port_in_ci] could not read the ephemeral port " + f"range ({e}); assuming none is reserved.") + return None + # Nonsense bounds would make get_static_port_range() hand out ports outside + # 1-65535, and bind() raises OverflowError (not OSError) for those, so + # reserve_port_from_range would propagate it instead of trying another port. + if not 1 <= low <= high <= 65535: + print_warning(f"[get_free_port_in_ci] ignoring implausible ephemeral " + f"port range ({low}, {high}).") + return None + return low, high + + +def get_static_port_range() -> Optional[tuple[int, int]]: + """Return a (low, high) window just below the kernel's ephemeral range. + + None is returned if the ephemeral range cannot be determined. + + Ports here are never handed out by bind(('', 0)), so a port reserved from + this window cannot be stolen by a sibling process launched with --port 0 -- + which is how a reserved disaggregated server port was lost to the test's + own worker. + """ + ephemeral_range = get_ephemeral_port_range() + if ephemeral_range is None: + return None + high = ephemeral_range[0] - 1 + low = max(1024, high - STATIC_PORT_RANGE_SIZE + 1) + if low > high: + return None + return low, high + + +def reserve_port_from_range(port_range: tuple[int, int], + source: str) -> Optional[int]: + """Probe-bind random ports from an inclusive (low, high) window. + + The first port found free is recorded in PORTS_IN_USE and returned; + None is returned once every candidate in the window is taken. + """ + global PORTS_IN_USE + + pid = os.getpid() + low, high = port_range + available_ports = [ + port for port in range(low, high + 1) if port not in PORTS_IN_USE + ] + num_candidates = len(available_ports) + + for attempt in range(1, num_candidates + 1): + # Get a random port from the available ports + port = random.choice(available_ports) + + # Check if the port is free + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + try: + s.bind(("localhost", port)) + PORTS_IN_USE.add(port) + print_info( + f"[get_free_port_in_ci] pid={pid} allocated port={port} " + f"from {source} range {port_range} after {attempt} " + f"attempt(s); {len(PORTS_IN_USE)} reserved in-process. The " + f"probe socket is now closed, so another process may take " + f"the port before the caller rebinds it (TOCTOU).") + return port + except OSError as e: + print_info( + f"[get_free_port_in_ci] pid={pid} candidate port={port} " + f"in {source} range {port_range} is busy ({e}); trying " + f"another.") + available_ports.remove(port) + continue + + print_warning( + f"[get_free_port_in_ci] pid={pid} exhausted all {num_candidates} " + f"candidate ports in {source} range {port_range}.") + return None + + +def get_free_port_in_ci(max_attempts: int = 100) -> int: + """Get a free port from the CI-assigned container port range. + + The range is [CONTAINER_PORT_START, CONTAINER_PORT_START + CONTAINER_PORT_NUM - 1]. + If those are unset, or every port in the range is already in use, fall back to + a port just below the kernel's ephemeral range, and only then to get_free_port. """ global PORTS_IN_USE @@ -685,45 +817,23 @@ def get_free_port_in_ci(max_attempts=100): container_port_start = int(os.environ.get("CONTAINER_PORT_START", -1)) container_port_num = int(os.environ.get("CONTAINER_PORT_NUM", -1)) if container_port_start != -1 and container_port_num != -1: - port_range = (container_port_start, - container_port_start + container_port_num - 1) - available_ports = [ - port for port in range(container_port_start, container_port_start + - container_port_num) - if port not in PORTS_IN_USE - ] - num_candidates = len(available_ports) - - for attempt in range(1, num_candidates + 1): - # Get a random port from the available ports - port = random.choice(available_ports) - - # Check if the port is free - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - try: - s.bind(("localhost", port)) - PORTS_IN_USE.add(port) - print_info( - f"[get_free_port_in_ci] pid={pid} allocated port={port} " - f"from CI range {port_range} after {attempt} attempt(s); " - f"{len(PORTS_IN_USE)} reserved in-process. The probe " - f"socket is now closed, so another process may take the " - f"port before the caller rebinds it (TOCTOU).") - return port - except OSError as e: - print_info( - f"[get_free_port_in_ci] pid={pid} candidate port={port} " - f"in CI range {port_range} is busy ({e}); trying another." - ) - available_ports.remove(port) - continue - - print_warning( - f"[get_free_port_in_ci] pid={pid} exhausted all {num_candidates} " - f"candidate ports in CI range {port_range}; falling back to a " - f"system-assigned ephemeral port.") - - # No port found in the range, try to get a random free port from the system + port = reserve_port_from_range( + (container_port_start, + container_port_start + container_port_num - 1), "CI") + if port is not None: + return port + + # No CI range configured, or every port in it is taken. Prefer a port below + # the ephemeral range over a system-assigned one: the latter is drawn from + # the same pool that trtllm-serve's own --port 0 workers bind from, so a + # sibling worker can take it before the caller rebinds it. + static_port_range = get_static_port_range() + if static_port_range is not None: + port = reserve_port_from_range(static_port_range, "static") + if port is not None: + return port + + # Last resort: a system-assigned ephemeral port. for _ in range(max_attempts): port = get_free_port() if port not in PORTS_IN_USE: diff --git a/tests/integration/defs/disaggregated/disagg_test_utils.py b/tests/integration/defs/disaggregated/disagg_test_utils.py index bf57abf0400f..8ea3dffe6fff 100644 --- a/tests/integration/defs/disaggregated/disagg_test_utils.py +++ b/tests/integration/defs/disaggregated/disagg_test_utils.py @@ -323,6 +323,32 @@ async def wait_for_port_released(port): return False +def get_registered_worker_urls(port: int) -> tuple[list[str], list[str]]: + """Return ``(ctx_urls, gen_urls)`` as registered with the disagg server. + + Workers launched with ``port=0`` pick their listening port inside the child + process, so the cluster registry is the only authoritative source for their + addresses. Call this once the server reports ready, i.e. once every worker + the test expects has registered. + + Args: + port: Disagg server port + + Returns: + tuple[list[str], list[str]]: Context and generation worker URLs, each + sorted so the ordering is stable across calls. + """ + assert port > 0, "port must be positive" + info_resp = requests.get(f"http://localhost:{port}/cluster_info", timeout=5) + assert info_resp.status_code == 200, f"cluster_info returned {info_resp.status_code}" + workers = info_resp.json().get("current_workers", {}) + + def _urls(role_key: str) -> list[str]: + return sorted(f"http://{w['host']}:{w['port']}" for w in workers.get(role_key, [])) + + return _urls("context_servers"), _urls("generation_servers") + + def verify_cluster_info(ready, ctx_workers=-1, gen_workers=-1, port=0, expected_code=200): """Verify cluster info from /cluster_info endpoint. diff --git a/tests/integration/defs/disaggregated/test_workers.py b/tests/integration/defs/disaggregated/test_workers.py index b6dfa488d5e8..9660178971a8 100644 --- a/tests/integration/defs/disaggregated/test_workers.py +++ b/tests/integration/defs/disaggregated/test_workers.py @@ -28,8 +28,8 @@ from defs.common import get_free_port_in_ci as get_free_port from defs.conftest import get_sm_version, skip_no_hopper from disagg_test_utils import (HEARTBEAT_INTERVAL, INACTIVE_TIMEOUT, - run_ctx_worker, run_disagg_server, - run_gen_worker, terminate, + get_registered_worker_urls, run_ctx_worker, + run_disagg_server, run_gen_worker, terminate, wait_for_disagg_server_ready) from transformers import AutoTokenizer @@ -614,35 +614,35 @@ def background_workers(llm_venv, config_file: str): ctx_workers = [] gen_workers = [] - ctx_urls = [] - gen_urls = [] next_device = 0 import torch num_gpus = torch.cuda.device_count() + # port=0 lets each worker bind an OS-assigned port in its own process and + # register it with the cluster, instead of pre-picking a port here and + # racing whoever takes it before the worker rebinds it. The real URLs are + # read back from the cluster registry once the server reports ready. for i in range(num_ctx): - port = get_free_port() - ctx_urls.append(f"http://localhost:{port}") ctx_workers.append( run_ctx_worker(model, ctx_worker_config, work_dir, - port=port, + port=0, device=next_device % num_gpus, - env=env)) + env=env, + worker_index=i)) next_device += gpus_per_ctx for i in range(num_gen): - port = get_free_port() - gen_urls.append(f"http://localhost:{port}") gen_workers.append( run_gen_worker(model, gen_worker_config, work_dir, - port=port, + port=0, device=next_device % num_gpus, - env=env)) + env=env, + worker_index=i)) next_device += gpus_per_gen server_config = { @@ -664,6 +664,10 @@ def background_workers(llm_venv, config_file: str): try: asyncio.run(wait_for_disagg_server_ready(disagg_port)) + ctx_urls, gen_urls = get_registered_worker_urls(disagg_port) + assert len(ctx_urls) == num_ctx and len(gen_urls) == num_gen, ( + f"Expected {num_ctx} ctx and {num_gen} gen workers registered, " + f"got {ctx_urls} and {gen_urls}") yield ctx_urls, gen_urls, disagg_port, internal_request_auth_key except Exception: logger.error("-------- Service discovery workers error --------") diff --git a/tests/integration/defs/perf/test_perf_sanity.py b/tests/integration/defs/perf/test_perf_sanity.py index a7aef4aee2fb..c5fb1594c2d0 100644 --- a/tests/integration/defs/perf/test_perf_sanity.py +++ b/tests/integration/defs/perf/test_perf_sanity.py @@ -32,8 +32,8 @@ from test_common.http_utils import fail_if_proc_died, wait_for_endpoint_ready from test_common.perf_sanity_matching import get_client_match_keys, get_server_match_keys +from defs.common import wait_for_reported_addr from defs.trt_test_alternative import print_info -from tensorrt_llm._utils import get_free_port from ..conftest import get_llm_root, llm_models_root from ._model_paths import MODEL_PATH_DICT as _MODEL_PATH_DICT_BASE @@ -1102,8 +1102,16 @@ def run_cmd(self, server_idx: int) -> List[str]: try: server_hostname = "localhost" - server_port = get_free_port() - server_cmd_with_port = add_host_port_to_cmd(server_cmd, server_hostname, server_port) + # port 0 + --report_addr: let the server bind a kernel-assigned + # port and tell us which one, instead of reserving one here and + # racing whoever takes it before the server binds. + server_addr_path = os.path.join(self.test_output_dir, f"trtllm-serve.{server_idx}.addr") + if os.path.exists(server_addr_path): + os.remove(server_addr_path) + server_cmd_with_port = add_host_port_to_cmd(server_cmd, server_hostname, 0) + [ + "--report_addr", + server_addr_path, + ] print_info(f"Starting server. cmd is {server_cmd_with_port}") server_file_path = os.path.join(self.test_output_dir, f"trtllm-serve.{server_idx}.log") @@ -1117,6 +1125,7 @@ def run_cmd(self, server_idx: int) -> List[str]: stdout=server_ctx, stderr=subprocess.STDOUT, ) + _, server_port = wait_for_reported_addr(server_addr_path, self.timeout, server_proc) wait_for_endpoint_ready( f"http://{server_hostname}:{server_port}/health", @@ -1207,19 +1216,28 @@ class DisaggTestCmds(NamedTuple): # receive env via SLURM env propagation set up by submit.py. server_configs: List[Tuple["ServerConfig", "ServerConfig", "DisaggConfig"]] = [] - def _generate_hostname_file(self, server_idx: int, port: int): - """Create hostname file for coordination.""" - hostnames_dir = os.path.join(self.test_output_dir, f"hostnames-{server_idx}") - if not os.path.exists(hostnames_dir): - os.makedirs(hostnames_dir, exist_ok=True) - hostname_file = os.path.join(hostnames_dir, f"{self.disagg_serving_type}.txt") - with open(hostname_file, "w") as f: - f.write(f"{self.hostname}:{port}") + def _hostnames_dir(self, server_idx: int) -> str: + """Directory the disagg tasks exchange bound addresses through. + + Scoped by SLURM job id so a rerun never reads the previous run's files: + test_output_dir is derived from the test case name alone and is created + with exist_ok=True, so it is reused across runs. The step id is + deliberately excluded -- each role runs as a separate srun step within + one job, and they must all agree on this path. + """ + run_id = os.environ.get("SLURM_JOB_ID", "local") + return os.path.join(self.test_output_dir, f"hostnames-{run_id}-{server_idx}") + + def _hostname_file(self, server_idx: int) -> str: + """Path this task's server reports its bound address to.""" + hostnames_dir = self._hostnames_dir(server_idx) + os.makedirs(hostnames_dir, exist_ok=True) + return os.path.join(hostnames_dir, f"{self.disagg_serving_type}.txt") def _generate_disagg_server_config(self, server_idx: int) -> str: """Generate disagg server config from hostname files.""" print_info(f"Generating disagg server config for server index {server_idx}") - hostnames_folder = os.path.join(self.test_output_dir, f"hostnames-{server_idx}") + hostnames_folder = self._hostnames_dir(server_idx) expected_count = self.num_ctx_servers + self.num_gen_servers start_time = time.time() hostnames = [] @@ -1236,7 +1254,11 @@ def _generate_disagg_server_config(self, server_idx: int) -> str: time.sleep(10) if not os.path.exists(hostnames_folder): continue - hostnames = os.listdir(hostnames_folder) + # Only completed files: trtllm-serve publishes its address by + # renaming a "..tmp" sibling into place, and counting + # those transient entries would both inflate the count and get + # parsed as a CTX/GEN url below. + hostnames = [f for f in os.listdir(hostnames_folder) if f.endswith(".txt")] if len(hostnames) >= expected_count: break @@ -1254,14 +1276,12 @@ def _generate_disagg_server_config(self, server_idx: int) -> str: elif hostname_file.startswith("GEN"): gen_hostnames.append(hostname_port) - # Allocate port here (after waiting) to minimize the window between - # port allocation and actual use, avoiding TOCTOU race conditions - # where another process on the same node grabs the port. - disagg_server_port = get_free_port() - + # port 0: the disagg server binds a kernel-assigned port and reports it + # back via --report_addr, so there is no window between choosing a port + # here and the server binding it. server_config = { "hostname": self.hostname, - "port": disagg_server_port, + "port": 0, "backend": "pytorch", "internal_request_auth_key": self.internal_request_auth_key, "context_servers": { @@ -1279,25 +1299,19 @@ def _generate_disagg_server_config(self, server_idx: int) -> str: print_info(f"Server config file {config_path} generated") return config_path + def _disagg_server_addr_file(self, server_idx: int) -> str: + """Path the disagg server reports its bound address to.""" + return os.path.join(self._hostnames_dir(server_idx), f"DISAGG_SERVER.{server_idx}.addr") + def _get_disagg_server_hostname_and_port(self, server_idx: int) -> Tuple[str, int]: - """Wait for and read disagg server config.""" - config_path = os.path.join(self.test_output_dir, f"server_config.{server_idx}.yaml") - start_time = time.time() - while True: - if os.path.exists(config_path): - print_info(f"Server config file found: {config_path}") - break - elapsed_time = time.time() - start_time - if elapsed_time > self.timeout: - raise RuntimeError( - f"Server config file {config_path} not found after {self.timeout}s" - ) - print_info(f"Waiting for server config file, elapsed time: {elapsed_time}s") - time.sleep(10) + """Wait for the disagg server to report the address it bound. - with open(config_path, "r") as f: - server_config = yaml.safe_load(f) - return server_config["hostname"], server_config["port"] + The config carries port 0, so the address is only known once the server + has bound; reading it from the config would yield 0. + """ + addr_path = self._disagg_server_addr_file(server_idx) + print_info(f"Waiting for disagg server address file {addr_path}") + return wait_for_reported_addr(addr_path, self.timeout) def wait_for_benchmark_ready( self, @@ -1451,8 +1465,11 @@ def run_cmd(self, server_idx: int) -> List[str]: self.server_configs[server_idx] if server_idx < len(self.server_configs) else None ) if "CTX" in self.disagg_serving_type or "GEN" in self.disagg_serving_type: - port = get_free_port() - self._generate_hostname_file(server_idx, port) + # port 0 + --report_addr: the worker binds a kernel-assigned port + # and publishes host:port itself, so no port is reserved here and + # left unbound while anything on the node could take it. The disagg + # server reads these files to build its config, exactly as before. + hostname_file = self._hostname_file(server_idx) is_ctx = "CTX" in self.disagg_serving_type server_cmd = ctx_cmd if is_ctx else gen_cmd @@ -1461,7 +1478,10 @@ def run_cmd(self, server_idx: int) -> List[str]: config_idx = server_cmd.index("--config") + 1 self._wait_for_config_file(server_cmd[config_idx]) - server_cmd = add_host_port_to_cmd(server_cmd, self.hostname, port) + server_cmd = add_host_port_to_cmd(server_cmd, self.hostname, 0) + [ + "--report_addr", + hostname_file, + ] try: print_info( f"Starting server. disagg_serving_type: {self.disagg_serving_type} cmd is {server_cmd}" @@ -1493,7 +1513,18 @@ def run_cmd(self, server_idx: int) -> List[str]: elif self.disagg_serving_type == "DISAGG_SERVER": try: + # _hostnames_dir is scoped by job, so a new job never sees an + # older one's files, but a retry within the same job and the + # same server_idx would. Drop the previous attempt's address + # first, or the BENCHMARK task connects to a dead port. This + # task owns the file exclusively, so removing it here is safe. + disagg_addr_path = self._disagg_server_addr_file(server_idx) + if os.path.exists(disagg_addr_path): + os.remove(disagg_addr_path) self._generate_disagg_server_config(server_idx) + # The config carries port 0; publish the resolved address so + # the BENCHMARK task can find the server. + disagg_cmd = disagg_cmd + ["--report_addr", disagg_addr_path] print_info(f"Starting disagg server. cmd is {disagg_cmd}") disagg_server_file_path = os.path.join( self.test_output_dir, diff --git a/tests/integration/defs/stress_test/disagg_cancel/harness.py b/tests/integration/defs/stress_test/disagg_cancel/harness.py index f2eb18cb5f7a..de7aed6c7fe1 100644 --- a/tests/integration/defs/stress_test/disagg_cancel/harness.py +++ b/tests/integration/defs/stress_test/disagg_cancel/harness.py @@ -422,6 +422,30 @@ def _parse_injection_schedule(raw_injections: list[Any]) -> list[_InjectionSpec] return specs +def _worker_id_pid(worker_id: str) -> Optional[int]: + """Extract the pid a cluster worker embedded in its worker id. + + DisaggClusterWorker builds the id as + ``{role}-{host}:{port}-{time_ms}-{pid}-{rand}``. Only the last two fields + are fixed-width in structure, so the pid is read from the tail: a substring + search for ``-{pid}-`` would also match inside a host name that contains a + dash-delimited digit run, such as ``node-1234-a``. + + Args: + worker_id: The ``worker_id`` reported by ``/cluster_info``. + + Returns: + The pid, or None if the id does not have the expected shape. + """ + fields = worker_id.rsplit("-", 2) + if len(fields) != 3: + return None + try: + return int(fields[1]) + except ValueError: + return None + + def _resolve_injection_target(target: str, tracked: list[_TrackedWorker]) -> _TrackedWorker: """Map a YAML target string to a tracked worker. @@ -1865,13 +1889,15 @@ def _respawn_tracked_worker(self, tracked: _TrackedWorker, *, timeout_s: float) Args: tracked: Worker to relaunch. - timeout_s: Maximum seconds to wait for HTTP 200 on ``/health``. + timeout_s: Total budget for the respawn, covering both the + cluster registration that reveals the worker's port and + the subsequent wait for HTTP 200 on ``/health``. Returns: True if the respawned worker reports healthy within the deadline; False otherwise. """ - from disagg_test_utils import _run_worker, get_free_port + from disagg_test_utils import _run_worker spec = tracked.spec old = tracked.wrapper @@ -1884,40 +1910,132 @@ def _respawn_tracked_worker(self, tracked: _TrackedWorker, *, timeout_s: float) role_key = "ctx" if spec.role == "ctx" else "gen" save_log = spec.log_path is not None - new_port = get_free_port() + deadline = time.monotonic() + timeout_s + # port=0 lets the respawned worker bind an OS-assigned port inside its + # own process and register it with the cluster. Pre-picking a port here + # would race whoever grabs it between the probe and the child's rebind. try: new_wrapper = _run_worker( spec.model_name, spec.worker_config, role_key, - port=new_port, + port=0, work_dir=spec.work_dir, device=spec.device, save_log=save_log, env=spec.env, + # Without this the respawn writes ctx_0/gen_0 filenames, so + # respawning worker N truncates worker 0's log out from under + # the log scanner. setup_disagg_cluster already passes it. + worker_index=spec.index, ) except Exception: - logger.exception( - "[injector] failed to respawn %s_%d on port %d", - spec.role, - spec.index, - new_port, - ) + logger.exception("[injector] failed to respawn %s_%d", spec.role, spec.index) return False tracked.wrapper = new_wrapper - spec.port = new_wrapper.port + # _teardown_cluster terminates the wrapper lists held in self._cluster, + # not self._tracked_workers, so the respawn has to replace its slot + # there or it survives teardown and keeps holding its GPUs. Done before + # the port wait below, so a respawn that never registers is cleaned up + # too. spec.index is per-role (see _make_worker_launch_spec callers). + if self._cluster is not None: + _, ctx_workers, gen_workers, *_ = self._cluster + workers = ctx_workers if spec.role == "ctx" else gen_workers + workers[spec.index] = new_wrapper if new_wrapper.log_path is not None: spec.log_path = new_wrapper.log_path + new_port = self._await_registered_worker_port( + role_key, + new_wrapper.process.pid, + deadline=deadline, + ) + if new_port is None: + logger.error( + "[injector] respawned %s_%d (pid %d) never registered a port with the cluster", + spec.role, + spec.index, + new_wrapper.process.pid, + ) + return False + spec.port = new_port + new_wrapper.port = new_port + logger.info( "[injector] respawned %s_%d on port %s; waiting up to %.0fs for /health", spec.role, spec.index, - new_wrapper.port, - timeout_s, + new_port, + max(0.0, deadline - time.monotonic()), + ) + + return self._wait_for_worker_health( + new_port, timeout_s=max(0.0, deadline - time.monotonic()) ) - return self._wait_for_worker_health(new_wrapper.port, timeout_s=timeout_s) + def _await_registered_worker_port( + self, role_key: str, pid: int, *, deadline: float + ) -> Optional[int]: + """Poll ``/cluster_info`` for the port a respawned worker registered. + + A worker launched with ``port=0`` only knows its port after it binds, + so the cluster registry is the sole authoritative source. Entries are + matched on the worker's pid, which ``WorkerInfo.worker_id`` embeds -- + matching on "a port we have not seen before" would be ambiguous while + the SIGKILLed worker's stale registration is still being reaped. See + ``_worker_id_pid`` for why the pid is parsed rather than substring-matched. + + Args: + role_key: ``"ctx"`` or ``"gen"``. + pid: Pid of the respawned ``trtllm-serve`` process. + deadline: ``time.monotonic()`` value to give up at. + + Returns: + The registered port, or None if it did not appear in time. + """ + if not self._server_url: + logger.error("[injector] no server URL bound; cannot resolve respawned worker port") + return None + + info_key = "context_servers" if role_key == "ctx" else "generation_servers" + seen_worker_ids: list[str] = [] + while time.monotonic() < deadline: + if self.stop_event.is_set() or self.failed_event.is_set(): + return None + try: + request_timeout = min(5.0, max(0.1, deadline - time.monotonic())) + with urllib.request.urlopen( + f"{self._server_url}/cluster_info", timeout=request_timeout + ) as response: + info = json.loads(response.read().decode("utf-8", errors="replace")) + seen_worker_ids = [] + for worker_info in (info.get("current_workers") or {}).get(info_key) or []: + if not isinstance(worker_info, dict): + continue + worker_id = str(worker_info.get("worker_id", "")) + seen_worker_ids.append(worker_id) + if _worker_id_pid(worker_id) != pid: + continue + try: + return int(worker_info["port"]) + except (KeyError, TypeError, ValueError): + logger.warning( + "[injector] cluster_info %s entry for pid %d has invalid port %r", + info_key, + pid, + worker_info.get("port"), + ) + except (json.JSONDecodeError, TimeoutError, OSError, urllib.error.URLError) as exc: + logger.debug("[injector] cluster_info poll failed: %s", exc) + self.stop_event.wait(timeout=min(1.0, max(0.0, deadline - time.monotonic()))) + + logger.warning( + "[injector] no %s worker matching pid %d in cluster_info; last saw %s", + info_key, + pid, + seen_worker_ids, + ) + return None def _wait_for_worker_health(self, port: int, *, timeout_s: float) -> bool: """Poll a worker's ``/health`` endpoint until healthy or timed out.""" diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 20d3543de415..9e51fde88bd2 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -89,6 +89,7 @@ l0_cpu: - unittest/llmapi/test_request_priority.py - unittest/llmapi/test_sampling_params.py - unittest/llmapi/test_serialization.py + - unittest/llmapi/test_serve_report_addr.py - unittest/llmapi/test_utils.py - unittest/metrics/test_collector.py - unittest/models/test_quant_config_utils.py diff --git a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml index ed37c2b310c4..4d4cabd2777f 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml @@ -385,6 +385,15 @@ commands: is_flag: false flags: - "--reasoning_parser" + report_addr: + type: str + default: null + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--report_addr" revision: type: str default: null @@ -520,6 +529,15 @@ commands: is_flag: false flags: - "--metrics-log-interval" + report_addr: + type: str + default: null + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--report_addr" request_timeout: type: int default: 180 diff --git a/tests/unittest/llmapi/test_serve_report_addr.py b/tests/unittest/llmapi/test_serve_report_addr.py new file mode 100644 index 000000000000..32462b08ee55 --- /dev/null +++ b/tests/unittest/llmapi/test_serve_report_addr.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit coverage for trtllm-serve's --report_addr publish/read contract. + +The integration migrations that consume this only run on GPU stages, so the +publisher, the reader and the guards that reject --report_addr where it cannot +work are covered here instead. +""" + +import os +import socket +import sys +import threading +import time +from pathlib import Path +from typing import Optional + +import click +import pytest +import yaml +from click.testing import CliRunner + +from tensorrt_llm.commands.serve import _publish_bound_address, disaggregated, launch_server + +# The reader half lives with the integration helpers, so both halves of the +# round trip are exercised together and a change to either side fails here. +_INTEGRATION_TESTS_DIR = Path(__file__).resolve().parents[2] / "integration" +if str(_INTEGRATION_TESTS_DIR) not in sys.path: + sys.path.insert(0, str(_INTEGRATION_TESTS_DIR)) +from defs.common import wait_for_reported_addr # noqa: E402 + +pytestmark = pytest.mark.cpu_only + + +class _FakeProcess: + """Stands in for a Popen handle; returncode None means still running.""" + + def __init__(self, returncode: Optional[int] = None) -> None: + self.returncode = returncode + + def poll(self) -> Optional[int]: + return self.returncode + + +@pytest.mark.parametrize( + "host,expected_host", + [ + ("localhost", "localhost"), + ("nvl72d066-T01", "nvl72d066-T01"), + ("10.67.24.211", "10.67.24.211"), + # Bracketed so the value is a usable URL authority. + ("::1", "[::1]"), + ("fe80::1", "[fe80::1]"), + ], +) +def test_publish_read_round_trip(tmp_path: Path, host: str, expected_host: str) -> None: + addr_path = str(tmp_path / "server.addr") + _publish_bound_address(addr_path, host, 22183) + + assert open(addr_path).read() == f"{expected_host}:22183\n" + assert wait_for_reported_addr(addr_path, timeout=5) == (expected_host, 22183) + + +@pytest.mark.parametrize("wildcard", ["0.0.0.0", "::", ""]) +def test_publish_replaces_wildcard_with_hostname(tmp_path: Path, wildcard: str) -> None: + """A reader cannot dial a wildcard bind address.""" + addr_path = str(tmp_path / "server.addr") + _publish_bound_address(addr_path, wildcard, 8000) + + reported_host = open(addr_path).read().strip().rsplit(":", 1)[0] + assert reported_host.strip("[]") == socket.gethostname() + + +def test_publish_creates_missing_parent_directories(tmp_path: Path) -> None: + addr_path = str(tmp_path / "deep" / "nested" / "server.addr") + _publish_bound_address(addr_path, "localhost", 8000) + assert open(addr_path).read() == "localhost:8000\n" + + +def test_publish_is_noop_without_a_path(tmp_path: Path) -> None: + """None/empty means the caller did not ask for the address.""" + before = set(os.listdir(tmp_path)) + _publish_bound_address(None, "localhost", 8000) + _publish_bound_address("", "localhost", 8000) + assert set(os.listdir(tmp_path)) == before + + +def test_publish_overwrites_without_leaking_temp_files(tmp_path: Path) -> None: + addr_path = str(tmp_path / "server.addr") + for port in (8000, 8001, 8002): + _publish_bound_address(addr_path, "localhost", port) + + assert open(addr_path).read() == "localhost:8002\n" + assert os.listdir(tmp_path) == ["server.addr"] + + +def test_concurrent_reader_never_sees_a_partial_line(tmp_path: Path) -> None: + """The rename must be atomic: readers poll this file while it is rewritten.""" + addr_path = str(tmp_path / "server.addr") + _publish_bound_address(addr_path, "localhost", 8000) + + # Every value the file is ever given. A reader must observe one of these + # exactly: a truncated write such as "localhost:\n" or "localhost:9\n" is + # what a non-atomic publisher would expose, and must fail this test. + published = {"localhost:8000\n"} | {f"localhost:{9000 + i}\n" for i in range(3)} + seen = set() + stop = threading.Event() + # Without this the publish loop can finish and set stop before the reader + # is ever scheduled, leaving seen empty and failing the test with no + # atomicity defect present. + reading = threading.Event() + + def read_loop() -> None: + while not stop.is_set(): + try: + seen.add(open(addr_path).read()) + except FileNotFoundError: + seen.add("") + reading.set() + + reader = threading.Thread(target=read_loop) + reader.start() + try: + assert reading.wait(timeout=30), "reader thread never ran" + for i in range(300): + _publish_bound_address(addr_path, "localhost", 9000 + (i % 3)) + finally: + stop.set() + reader.join() + + assert seen, "reader observed nothing" + assert seen <= published, f"observed partial or unexpected values: {seen - published}" + + +def test_reader_fails_fast_when_the_server_dies(tmp_path: Path) -> None: + """A crashed server must not burn the whole timeout.""" + addr_path = str(tmp_path / "server.addr") + started = time.monotonic() + with pytest.raises(RuntimeError, match="exited with code 1"): + wait_for_reported_addr(addr_path, timeout=30, process=_FakeProcess(1)) + assert time.monotonic() - started < 10 + + +def test_reader_times_out_when_nothing_is_published(tmp_path: Path) -> None: + addr_path = str(tmp_path / "server.addr") + with pytest.raises(TimeoutError, match="did not report its address"): + wait_for_reported_addr(addr_path, timeout=1) + + +def test_reader_waits_for_a_late_write(tmp_path: Path) -> None: + addr_path = str(tmp_path / "server.addr") + timer = threading.Timer(1.0, lambda: _publish_bound_address(addr_path, "localhost", 12345)) + timer.start() + try: + assert wait_for_reported_addr(addr_path, timeout=20, process=_FakeProcess()) == ( + "localhost", + 12345, + ) + finally: + timer.join() + + +@pytest.mark.parametrize("port,report_addr", [(0, "/tmp/a.addr"), (8000, "/tmp/a.addr"), (0, None)]) +def test_launch_server_rejects_multiple_frontends(port: int, report_addr: Optional[str]) -> None: + """Each frontend re-execs and would bind (and publish) its own port.""" + llm_args = {"backend": "pytorch", "model": "dummy", "num_serve_frontends": 2} + with pytest.raises(click.BadParameter, match="single serving frontend"): + launch_server("localhost", port, llm_args, report_addr=report_addr) + + +def test_launch_server_requires_a_way_to_learn_the_port() -> None: + """Port 0 with neither service discovery nor --report_addr is unusable.""" + llm_args = {"backend": "pytorch", "model": "dummy"} + with pytest.raises(AssertionError, match="Port must be specified"): + launch_server("localhost", 0, llm_args) + + +@pytest.mark.parametrize( + "port,extra_args", + [ + (0, []), + (0, ["--report_addr", "/tmp/a.addr"]), + (8000, ["--report_addr", "/tmp/a.addr"]), + ], +) +def test_disaggregated_rejects_a_worker_fleet( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, port: int, extra_args: list[str] +) -> None: + """A fleet spreads the port over N workers, so one published address is wrong. + + Only the rejected combinations are exercised: anything the guard lets + through goes on to bind a socket and serve. + """ + # set_prometheus_multiproc_dir, which disaggregated() calls before reaching + # the guard, deletes the directory its own environment variable points at on + # the second call in a process, so the third call raises FileNotFoundError. + # Pointing the variable at a directory that outlives every invocation keeps + # these parametrised cases independent of one another. + monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) + + config = tmp_path / "disagg.yaml" + config.write_text( + yaml.safe_dump( + { + "hostname": "localhost", + "port": port, + "num_workers": 2, + "context_servers": {}, + "generation_servers": {}, + } + ) + ) + + result = CliRunner().invoke(disaggregated, ["-c", str(config)] + extra_args) + + assert result.exit_code != 0 + assert "single self-contained disaggregated server" in result.output