[https://nvbugs/6813945][fix] Use stored metadata server URLs as is and support etcd credentials and signed server records

examples/disaggregated/README.md

@@ -583,10 +583,29 @@ trtllm-serve disaggregated -c disagg_config.yaml -m ./metadata_config.yaml
 
 The metadata_config file looks like
 ```yaml
+server_type: etcd
 hostname: "localhost"
 port: 2379
 health_check_timeout: 5.0
-refersh_interval: 10.0
+refresh_interval: 10.0
+```
+
+To use an ETCD server with TLS, add ```ca_cert``` (with optional ```cert_key``` and ```cert_cert``` for client certificates) to the metadata_config file. To use ETCD authentication, add ```user``` and ```password``` together with ```ca_cert```: the etcd client sends its authentication token only over a TLS connection, and the config is rejected if ```user``` and ```password``` are set without ```ca_cert```. The ETCD user needs read and write permission on the ```trtllm/``` key prefix only. The etcd client logs in once, when the server starts, and does not renew its token: after ETCD stops accepting the token, for example following an ETCD restart, restart the context, generation and disaggregated servers so that they log in again. Set ```record_auth_key``` to the same non-empty secret in the metadata_config file of every context, generation and disaggregated server: servers then sign the records they register, and the disaggregated server ignores records without a valid signature.
+
+Without these fields the connection to the ETCD server is anonymous and unencrypted, and the disaggregated server uses every ```trtllm/*``` record found in the store. When the ETCD server is reachable from hosts other than the context, generation and disaggregated servers, enable ETCD authentication with TLS together with ```record_auth_key```. The signature is checked only by readers, so access control of the store itself still comes from the ETCD server. A metadata_config file with all optional fields looks like
+
+```yaml
+server_type: etcd
+hostname: "etcd.internal"
+port: 2379
+health_check_timeout: 5.0
+refresh_interval: 10.0
+user: "trtllm"
+password: "<etcd user credential>"
+ca_cert: "/etc/trtllm/etcd-ca.pem"
+cert_key: "/etc/trtllm/etcd-client-key.pem"
+cert_cert: "/etc/trtllm/etcd-client.pem"
+record_auth_key: "<shared secret, identical on every server>"
 ```
 
 The ```hostname``` and ```port``` must match those used when starting the ETCD server. The ```health_check_timeout``` parameter specifies how long a server will be considered dead if no healthy response is received. By default, trtllm will perform two checks before marking a server as dead. The ```refresh_interval``` parameter determines how often the latest server list is fetched from the ETCD server.

tensorrt_llm/llmapi/disagg_utils.py

@@ -13,9 +13,11 @@
 import yaml
 from mpi4py.MPI import COMM_WORLD, Comm
 from mpi4py.util import pkl5
-from pydantic import Field, TypeAdapter, ValidationError
+from pydantic import (ConfigDict, Field, TypeAdapter, ValidationError,
+                      model_validator)
 
 from .._utils import global_mpi_rank, global_mpi_size
+from .utils import StrictBaseModel
 
 __all__ = [
     'InternalRequestAuthKey',
@@ -29,7 +31,8 @@
     'validate_internal_request_auth_key',
 ]
 
-InternalRequestAuthKey = Annotated[str, Field(strict=True, min_length=1)] | None
+_AuthKey = Annotated[str, Field(strict=True, min_length=1)]
+InternalRequestAuthKey = _AuthKey | None
 _AUTH_KEY_ADAPTER = TypeAdapter(InternalRequestAuthKey)
 
 
@@ -186,13 +189,68 @@ class DisaggServerConfig():
     subagent_affinity_scope: Literal['context', 'both'] = 'context'
 
 
-@dataclass
-class MetadataServerConfig():
-    server_type: Literal['etcd']
-    hostname: str = "localhost"
-    port: int = 2379
-    health_check_timeout: float = 5.0
-    refresh_interval: float = 10.0
+class MetadataServerConfig(StrictBaseModel):
+    """Settings of the metadata server that context and generation servers register with."""
+
+    # Keeps the values of rejected fields, such as credentials, out of errors.
+    model_config = ConfigDict(hide_input_in_errors=True)
+
+    server_type: Literal['etcd'] = Field(
+        description="Type of the metadata server.")
+    hostname: str = Field(default="localhost",
+                          description="Host of the metadata server.")
+    port: int = Field(default=2379, description="Port of the metadata server.")
+    health_check_timeout: float = Field(
+        default=5.0,
+        description="Seconds without a healthy response after which a server "
+        "is considered dead.")
+    refresh_interval: float = Field(
+        default=10.0,
+        description="Seconds between fetches of the latest server list.")
+    user: Optional[str] = Field(
+        default=None,
+        description="User name for the metadata server. Requires `password` "
+        "and `ca_cert`.")
+    password: Optional[str] = Field(
+        default=None,
+        repr=False,
+        description="Password for the metadata server. Requires `user` and "
+        "`ca_cert`.")
+    ca_cert: Optional[str] = Field(
+        default=None,
+        description="Path to the CA certificate used to verify the metadata "
+        "server. Enables TLS.")
+    cert_key: Optional[str] = Field(
+        default=None,
+        description="Path to the client key. Requires `cert_cert` and "
+        "`ca_cert`.")
+    cert_cert: Optional[str] = Field(
+        default=None,
+        description="Path to the client certificate. Requires `cert_key` and "
+        "`ca_cert`.")
+    record_auth_key: Optional[_AuthKey] = Field(
+        default=None,
+        repr=False,
+        description="Shared secret used to sign and verify server records. "
+        "When set, records without a valid signature are ignored by readers.")
+
+    @model_validator(mode="after")
+    def _validate_credentials(self) -> "MetadataServerConfig":
+        if (self.user is None) != (self.password is None):
+            raise ValueError(
+                "metadata server user and password must be set together")
+        if self.user is not None and self.ca_cert is None:
+            raise ValueError(
+                "metadata server user and password require ca_cert: the etcd "
+                "client sends its authentication token only over a TLS "
+                "connection")
+        if (self.cert_key is None) != (self.cert_cert is None):
+            raise ValueError(
+                "metadata server cert_key and cert_cert must be set together")
+        if self.cert_key is not None and self.ca_cert is None:
+            raise ValueError(
+                "metadata server cert_key and cert_cert require ca_cert")
+        return self
 
 
 def get_usage_tokens_from_ctx(

tensorrt_llm/serve/metadata_server.py

@@ -1,10 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+import hashlib
+import hmac
 import json
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Any, Optional
 
 from tensorrt_llm.llmapi.disagg_utils import MetadataServerConfig
+from tensorrt_llm.logger import logger
 from tensorrt_llm.serve.cluster_storage import etcd3
 
+METADATA_KEY_PREFIX = "trtllm/"
+RECORD_SIGNATURE_FIELD = "signature"
+_SIGNATURE_PREFIX = "sha256="
+
 
 class RemoteDictionary(ABC):
 
@@ -27,8 +36,21 @@ def keys(self) -> list[str]:
 
 class EtcdDictionary(RemoteDictionary):
 
-    def __init__(self, host: str, port: int):
-        self._client = etcd3.client(host, port)
+    def __init__(self,
+                 host: str,
+                 port: int,
+                 user: Optional[str] = None,
+                 password: Optional[str] = None,
+                 ca_cert: Optional[str] = None,
+                 cert_key: Optional[str] = None,
+                 cert_cert: Optional[str] = None):
+        self._client = etcd3.client(host,
+                                    port,
+                                    ca_cert=ca_cert,
+                                    cert_key=cert_key,
+                                    cert_cert=cert_cert,
+                                    user=user,
+                                    password=password)
 
     def get(self, key: str) -> str:
         return self._client.get(key)
@@ -44,7 +66,8 @@ def keys(self) -> list[str]:
         # This implementation assumes that key is in the
         # format of "trtllm/executor_name/key"
         unique_keys = set()
-        for _, metadata in self._client.get_all():
+        for _, metadata in self._client.get_prefix(METADATA_KEY_PREFIX,
+                                                   keys_only=True):
             key = metadata.key.decode('utf-8')
             sub_keys = key.split('/')
             if len(sub_keys) >= 2:
@@ -53,16 +76,76 @@ def keys(self) -> list[str]:
         return list(unique_keys)
 
 
+_SIGNATURE_HEX_LENGTH = 64  # SHA-256 hex digest
+
+
+def _is_well_formed_signature(value: Any) -> bool:
+    """Whether ``value`` has the form ``_sign_record`` produces."""
+    if not isinstance(value, str) or not value.startswith(_SIGNATURE_PREFIX):
+        return False
+    digest = value[len(_SIGNATURE_PREFIX):]
+    return len(digest) == _SIGNATURE_HEX_LENGTH and all(c in "0123456789abcdef"
+                                                        for c in digest)
+
+
+def _sign_record(record_auth_key: str, key: str, value: dict) -> str:
+    unsigned = {k: v for k, v in value.items() if k != RECORD_SIGNATURE_FIELD}
+    payload = json.dumps({
+        "key": key,
+        "value": unsigned
+    },
+                         sort_keys=True,
+                         separators=(",", ":")).encode("utf-8")
+    digest = hmac.new(record_auth_key.encode("utf-8"), payload,
+                      hashlib.sha256).hexdigest()
+    return f"{_SIGNATURE_PREFIX}{digest}"
+
+
 class JsonDictionary:
+    """JSON view of a RemoteDictionary.
+
+    With ``record_auth_key`` set, every dict value written is signed with an
+    HMAC bound to its key, and ``get`` returns None for a value that cannot be
+    parsed or whose signature is missing or does not match.
+    """
 
-    def __init__(self, dict: RemoteDictionary):
+    def __init__(self,
+                 dict: RemoteDictionary,
+                 record_auth_key: Optional[str] = None):
         self._dict = dict
+        self._record_auth_key = record_auth_key
 
-    def get(self, key: str) -> str:
+    def get(self, key: str) -> Any:
         bytes_data, _ = self._dict.get(key)
-        return json.loads(bytes_data.decode('utf-8'))
-
-    def put(self, key: str, value: str):
+        if bytes_data is None:
+            return None
+        try:
+            value = json.loads(bytes_data.decode('utf-8'))
+        except (ValueError, RecursionError):
+            if self._record_auth_key is None:
+                raise
+            logger.warning(f"Ignoring unreadable metadata record {key}")
+            return None
+        if self._record_auth_key is None:
+            return value
+        if not isinstance(value, dict):
+            logger.warning(f"Ignoring unsigned metadata record {key}")
+            return None
+        provided = value.get(RECORD_SIGNATURE_FIELD)
+        if not _is_well_formed_signature(provided) or not hmac.compare_digest(
+                provided, _sign_record(self._record_auth_key, key, value)):
+            logger.warning(
+                f"Ignoring metadata record {key} with an invalid signature")
+            return None
+        return {k: v for k, v in value.items() if k != RECORD_SIGNATURE_FIELD}
+
+    def put(self, key: str, value: Any):
+        if self._record_auth_key is not None:
+            if not isinstance(value, dict):
+                raise ValueError("Signed metadata records must be dicts")
+            value = dict(value)
+            value[RECORD_SIGNATURE_FIELD] = _sign_record(
+                self._record_auth_key, key, value)
         self._dict.put(key, json.dumps(value))
 
     def remove(self, key: str):
@@ -81,10 +164,19 @@ def create_metadata_server(
 
     if metadata_server_cfg.server_type == 'etcd':
         dict = EtcdDictionary(host=metadata_server_cfg.hostname,
-                              port=metadata_server_cfg.port)
+                              port=metadata_server_cfg.port,
+                              user=metadata_server_cfg.user,
+                              password=metadata_server_cfg.password,
+                              ca_cert=metadata_server_cfg.ca_cert,
+                              cert_key=metadata_server_cfg.cert_key,
+                              cert_cert=metadata_server_cfg.cert_cert)
     else:
         raise ValueError(
             f"Unsupported metadata server type: {metadata_server_cfg.server_type}"
         )
 
-    return JsonDictionary(dict)
+    if metadata_server_cfg.record_auth_key is None:
+        logger.warning(
+            "Metadata server record_auth_key is not set; server records are "
+            "not signed or verified.")
+    return JsonDictionary(dict, metadata_server_cfg.record_auth_key)

tensorrt_llm/serve/openai_client.py

@@ -212,6 +212,13 @@ async def _finish_request(
         ...
 
 
+def _server_base_url(server: str) -> str:
+    """Return the base URL of a router server entry, with or without a scheme."""
+    if server.startswith(("http://", "https://")):
+        return server.rstrip("/")
+    return f"http://{server}".rstrip("/")
+
+
 class OpenAIHttpClient(OpenAIClient):
     def __init__(
         self,
@@ -270,8 +277,7 @@ async def post_json(
         response_type: Type[_ResponseT],
         server: str,
     ) -> _ResponseT:
-        server_url = server if server.startswith("http") else f"http://{server}"
-        url = f"{server_url.rstrip('/')}/{endpoint}"
+        url = f"{_server_base_url(server)}/{endpoint}"
         # A count-tokens body carries the whole conversation, so it is as large
         # as a chat body and encoding it costs the orchestrator's event loop the
         # same way -- hence the same msgpack transport as _send_request. The
@@ -297,8 +303,7 @@ async def get_json(
         response_type: Type[_ResponseT],
         server: str,
     ) -> _ResponseT:
-        server_url = server if server.startswith("http") else f"http://{server}"
-        url = f"{server_url.rstrip('/')}/{endpoint}"
+        url = f"{_server_base_url(server)}/{endpoint}"
         async with self._session.get(url) as response:
             if response.status >= 400:
                 error_body = await response.text()
@@ -340,7 +345,7 @@ async def _send_request(
             pinned_dp = request.disaggregated_params
             if pinned_dp is not None:
                 req_id = pinned_dp.disagg_request_id
-        url = f"http://{server}/{endpoint}"
+        url = f"{_server_base_url(server)}/{endpoint}"
         # disaggregated_params is None when conditional_disagg bypasses ctx.
         _dp = request.disaggregated_params
         _ctx_rid = _dp.ctx_request_id if _dp is not None else None
@@ -697,11 +702,7 @@ async def check_ready_for_servers(
     ) -> Tuple[List[str], List[str]]:
         async def check_server_ready(server: str) -> bool:
             try:
-                url = (
-                    f"{server}/health"
-                    if server.startswith("http://")
-                    else f"http://{server}/health"
-                )
+                url = f"{_server_base_url(server)}/health"
                 async with session.get(url) as response:
                     return response.status == 200
             except Exception:

tensorrt_llm/serve/openai_disagg_server.py

@@ -56,7 +56,8 @@
                                                 resolve_request_conversation_id)
 from tensorrt_llm.serve.disagg_coordinator import (CoordinatorClient,
                                                    DisaggCoordinatorService)
-from tensorrt_llm.serve.openai_client import OpenAIClient, OpenAIHttpClient
+from tensorrt_llm.serve.openai_client import (OpenAIClient, OpenAIHttpClient,
+                                              _server_base_url)
 from tensorrt_llm.serve.openai_disagg_service import (
     OpenAIDisaggregatedService, ResponseHooks)
 from tensorrt_llm.serve.openai_protocol import (
@@ -713,8 +714,7 @@ async def align_steady_clock_offset(session: aiohttp.ClientSession, server_url:
             # Negate the offset so that worker servers can adjust their steady clock by adding the new offset
             await set_steady_clock_offset(session, server_url, -offset)
 
-        server_scheme = "http://" if not server.startswith("http://") else ""
-        server_url = f"{server_scheme}{server}/steady_clock_offset"
+        server_url = f"{_server_base_url(server)}/steady_clock_offset"
 
         try:
             async with aiohttp.ClientSession(