diff --git a/tensorrt_llm/serve/cluster_storage.py b/tensorrt_llm/serve/cluster_storage.py index d2ba0d911717..bb4f1902b6f4 100644 --- a/tensorrt_llm/serve/cluster_storage.py +++ b/tensorrt_llm/serve/cluster_storage.py @@ -2,6 +2,7 @@ import asyncio import importlib.metadata as importlib_metadata import importlib.util +import ipaddress import sys import time from dataclasses import dataclass @@ -18,6 +19,30 @@ from tensorrt_llm.logger import logger +def is_loopback_host(host: Optional[str]) -> bool: + if not isinstance(host, str) or not host: + return False + if host.lower() == "localhost": + return True + try: + return ipaddress.ip_address(host).is_loopback + except ValueError: + return False + + +def validate_http_cluster_storage_scope(cluster_uri: str, + server_host: str) -> None: + parsed_uri = urlparse(cluster_uri) + if parsed_uri.scheme not in ("http", "https"): + return + if is_loopback_host(parsed_uri.hostname) and is_loopback_host(server_host): + return + raise ValueError( + "HTTP cluster storage is only supported for loopback-only " + "disaggregated serving. Use a loopback disagg_cluster.cluster_uri and " + "hostname, or use etcd for cluster storage.") + + def _find_module_file_in_distribution(dist, module_name: str): module_path = module_name.replace(".", "/") candidates = (f"{module_path}/__init__.py", f"{module_path}.py") diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index be5a65c0e45a..f57f251cab87 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -36,8 +36,9 @@ get_ctx_gen_server_addrs, get_global_disagg_request_id) from tensorrt_llm.logger import logger -from tensorrt_llm.serve.cluster_storage import (HttpClusterStorageServer, - create_cluster_storage) +from tensorrt_llm.serve.cluster_storage import ( + HttpClusterStorageServer, create_cluster_storage, + validate_http_cluster_storage_scope) from tensorrt_llm.serve.metadata_server import create_metadata_server from tensorrt_llm.serve.openai_client import OpenAIClient, OpenAIHttpClient from tensorrt_llm.serve.openai_disagg_service import ( @@ -105,7 +106,13 @@ def __init__(self, self._metadata_server = create_metadata_server(metadata_server_cfg) self._perf_metrics_collector = DisaggPerfMetricsCollector(config.perf_metrics_max_requests) - self._disagg_cluster_storage = create_cluster_storage(config.disagg_cluster_config.cluster_uri, config.disagg_cluster_config.cluster_name) if config.disagg_cluster_config else None + self._disagg_cluster_storage = None + if config.disagg_cluster_config: + validate_http_cluster_storage_scope( + config.disagg_cluster_config.cluster_uri, config.hostname) + self._disagg_cluster_storage = create_cluster_storage( + config.disagg_cluster_config.cluster_uri, + config.disagg_cluster_config.cluster_name) self._service = OpenAIDisaggregatedService( self._config, self._ctx_router, self._gen_router, self._create_client, diff --git a/tests/integration/defs/test_e2e.py b/tests/integration/defs/test_e2e.py index 9c1f5b681926..947c191394ba 100644 --- a/tests/integration/defs/test_e2e.py +++ b/tests/integration/defs/test_e2e.py @@ -1261,7 +1261,7 @@ def test_trtllm_multimodal_benchmark_serving(llm_root, llm_venv): @pytest.mark.skip_less_device(4) @pytest.mark.skip_less_device_memory(40000) -@pytest.mark.parametrize("service_discovery", ["etcd", "http"]) +@pytest.mark.parametrize("service_discovery", ["etcd"]) def test_openai_disagg_multi_nodes_completion_service_discovery( llm_root, llm_venv, service_discovery): test_root = unittest_path() / "llmapi" / "apps" diff --git a/tests/integration/test_lists/qa/llm_function_multinode.txt b/tests/integration/test_lists/qa/llm_function_multinode.txt index 898a65e59b89..a70db0ad96f3 100644 --- a/tests/integration/test_lists/qa/llm_function_multinode.txt +++ b/tests/integration/test_lists/qa/llm_function_multinode.txt @@ -5,4 +5,4 @@ test_e2e.py::test_multi_nodes_eval[DeepSeek-R1/DeepSeek-R1-0528-FP4-tp16-mmlu] test_e2e.py::test_multi_nodes_eval[Kimi-K2-Thinking-NVFP4-tp16-mmlu] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp2pp1-gen_tp2pp1] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp1pp2-gen_tp1pp2] -test_e2e.py::test_openai_disagg_multi_nodes_completion_service_discovery[http] +test_e2e.py::test_openai_disagg_multi_nodes_completion_service_discovery[etcd] diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 46312e23c7ca..58846ae9d161 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -327,7 +327,6 @@ test_e2e.py::test_multi_nodes_eval[Qwen3/saved_models_Qwen3-235B-A22B_nvfp4_hf-t test_e2e.py::test_openai_chat_example[trt] SKIP (https://nvbugs/5477444) test_e2e.py::test_openai_completions_example[trt] SKIP (https://nvbugs/5701450) test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp1pp2-gen_tp1pp2] SKIP (https://nvbugs/6190759) -test_e2e.py::test_openai_disagg_multi_nodes_completion_service_discovery[http] SKIP (https://nvbugs/6115562) test_e2e.py::test_openai_kv_cache_contamination SKIP (https://nvbugs/6227203) test_e2e.py::test_ptp_quickstart_advanced_deepseek_r1_w4afp8_8gpus[DeepSeek-R1-W4AFP8-DeepSeek-R1/DeepSeek-R1-W4AFP8] SKIP (https://nvbugs/5836830) test_e2e.py::test_trtllm_bench_iteration_log[TRT-streaming-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] SKIP (https://nvbugs/5448523) diff --git a/tests/unittest/disaggregated/test_cluster_storage.py b/tests/unittest/disaggregated/test_cluster_storage.py index 9385fdec71fe..d72c5a38d1ce 100644 --- a/tests/unittest/disaggregated/test_cluster_storage.py +++ b/tests/unittest/disaggregated/test_cluster_storage.py @@ -10,11 +10,10 @@ import uvicorn from fastapi import FastAPI -from tensorrt_llm.serve.cluster_storage import (HttpClusterStorageServer, - StorageItem, WatchEvent, - WatchEventType, - create_cluster_storage, - create_cluster_storage_client) +from tensorrt_llm.serve.cluster_storage import ( + HttpClusterStorageServer, StorageItem, WatchEvent, WatchEventType, + create_cluster_storage, create_cluster_storage_client, is_loopback_host, + validate_http_cluster_storage_scope) _counter = 0 @@ -44,6 +43,53 @@ def run_in_thread(self): timeout = pytest.mark.timeout +@pytest.mark.parametrize( + "host", ["localhost", "LOCALHOST", "LocalHost", "127.0.0.1", "::1"]) +def test_is_loopback_host_accepts_loopback_hosts(host): + assert is_loopback_host(host) + + +@pytest.mark.parametrize("host", + [None, "", "0.0.0.0", "10.0.0.1", "example.com"]) +def test_is_loopback_host_rejects_non_loopback_hosts(host): + assert not is_loopback_host(host) + + +@pytest.mark.parametrize("scheme", ["http", "https"]) +@pytest.mark.parametrize("uri_host", ["localhost", "127.0.0.1", "[::1]"]) +@pytest.mark.parametrize("server_host", ["localhost", "127.0.0.1", "::1"]) +def test_http_cluster_storage_scope_allows_loopback_only( + scheme, uri_host, server_host): + validate_http_cluster_storage_scope(f"{scheme}://{uri_host}:18000", + server_host) + + +@pytest.mark.parametrize( + "cluster_uri, server_host", + [ + ("http://10.0.0.1:18000", "localhost"), + ("http://localhost:18000", "0.0.0.0"), + ("https://example.com:18000", "127.0.0.1"), + ("https://127.0.0.1:18000", "10.0.0.1"), + ], +) +def test_http_cluster_storage_scope_rejects_non_loopback_scope( + cluster_uri, server_host): + with pytest.raises(ValueError, match="loopback-only"): + validate_http_cluster_storage_scope(cluster_uri, server_host) + + +@pytest.mark.parametrize( + "cluster_uri", + [ + "etcd://10.0.0.1:2379", + "etcd://example.com:2379", + ], +) +def test_etcd_cluster_storage_scope_is_unchanged(cluster_uri): + validate_http_cluster_storage_scope(cluster_uri, "0.0.0.0") + + @pytest_asyncio.fixture(scope="function") async def storage_client(storage_server): _, cluster_uri = storage_server diff --git a/tests/unittest/llmapi/apps/openai_server.py b/tests/unittest/llmapi/apps/openai_server.py index ebbe0d5627fe..c5326a3f291a 100644 --- a/tests/unittest/llmapi/apps/openai_server.py +++ b/tests/unittest/llmapi/apps/openai_server.py @@ -164,6 +164,7 @@ def __init__(self, self.disagg_config = self._get_extra_config() if disagg_config: self.disagg_config.update(disagg_config) + self.host = self.disagg_config.get("hostname", self.host) self.log_path = log_path self.log_file = None self.extra_config_file = os.path.join( diff --git a/tests/unittest/llmapi/apps/test_disagg_serving_perf_metrics.py b/tests/unittest/llmapi/apps/test_disagg_serving_perf_metrics.py index 5fc9bdc0f142..9b18e3025d73 100644 --- a/tests/unittest/llmapi/apps/test_disagg_serving_perf_metrics.py +++ b/tests/unittest/llmapi/apps/test_disagg_serving_perf_metrics.py @@ -94,6 +94,7 @@ def worker(server_role: str, port: int): @pytest.fixture def disagg_server(disagg_cluster_config: dict, workers, disagg_port: int): disagg_config = { + "hostname": "localhost", "port": disagg_port, "disagg_cluster": disagg_cluster_config, "perf_metrics_max_requests": 1000,