[TRTLLM-17267][refactor] Move the operational endpoints into serving classes and routers

.claude/skills/perf-regression-cookbook/references/kv-cache-transceiver/disagg-gen-retokenize-and-memcpy.md

@@ -73,7 +73,7 @@ failed_prs: []
 - **Detection signal:** gen-server host time before the first forward step
   grows with prompt length (tokenizer + block-key construction spans in a
   py-spy/nsys host profile); verify the guard exists with
-  `grep -n "prompt_token_ids is not None" tensorrt_llm/serve/openai_server.py`
+  `grep -n "prompt_token_ids is not None" tensorrt_llm/serve/api/openai/chat/serving.py`
   and compare ctx-vs-gen TTFT via the server's `/perf_metrics` endpoint
   (requires `TRTLLM_KVCACHE_TIME_OUTPUT_PATH` to be set on this path).
 - **Prevention/guard:** the PR adds

.github/CODEOWNERS

@@ -371,6 +371,7 @@
 /cpp/tensorrt_llm/thop/*lysses* @NVIDIA/trt-llm-torch-visual-gen-devs
 /cpp/tensorrt_llm/kernels/ulysses* @NVIDIA/trt-llm-torch-visual-gen-devs
 /tensorrt_llm/media @NVIDIA/trt-llm-torch-visual-gen-devs
+/tensorrt_llm/serve/api/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs
 /tensorrt_llm/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs @NVIDIA/trt-llm-runtime-devs
 /tests/integration/defs/examples/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs
 /tests/integration/defs/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs

.pre-commit-config.yaml

@@ -417,13 +417,46 @@ common-files: &common_files |
         tensorrt_llm/scaffolding/worker.py |
         tensorrt_llm/scheduling_params.py |
         tensorrt_llm/serve/__init__.py |
+        tensorrt_llm/serve/_base/msgspec_route.py |
+        tensorrt_llm/serve/_base/serving.py |
+        tensorrt_llm/serve/_base/sse.py |
+        tensorrt_llm/serve/api/anthropic/serving.py |
+        tensorrt_llm/serve/api/openai/chat/serving.py |
+        tensorrt_llm/serve/api/openai/completions/_merge.py |
+        tensorrt_llm/serve/api/openai/completions/serving.py |
+        tensorrt_llm/serve/api/openai/embeddings/serving.py |
+        tensorrt_llm/serve/api/openai/mm_encoder/serving.py |
+        tensorrt_llm/serve/api/openai/tokenize/serving.py |
+        tensorrt_llm/serve/api/ops/metrics/prometheus.py |
+        tensorrt_llm/serve/api/ops/metrics/serving.py |
+        tensorrt_llm/serve/api/ops/profile/serving.py |
+        tensorrt_llm/serve/api/ops/rl/serving.py |
+        tensorrt_llm/serve/api/ops/runtime_control/serving.py |
+        tensorrt_llm/serve/api/ops/worker/serving.py |
+        tensorrt_llm/serve/api/visual_gen/media.py |
+        tensorrt_llm/serve/api/visual_gen/protocol.py |
+        tensorrt_llm/serve/api/visual_gen/serving_base.py |
+        tensorrt_llm/serve/api/visual_gen/serving_images.py |
         tensorrt_llm/serve/chat_utils.py |
         tensorrt_llm/serve/cluster_storage.py |
         tensorrt_llm/serve/disagg_auto_scaling.py |
         tensorrt_llm/serve/harmony_adapter.py |
         tensorrt_llm/serve/metadata_server.py |
+        tensorrt_llm/serve/middleware/arrival_time.py |
+        tensorrt_llm/serve/middleware/user_middleware.py |
         tensorrt_llm/serve/openai_disagg_server.py |
-        tensorrt_llm/serve/openai_protocol.py |
+        tensorrt_llm/serve/openai_protocol/chat.py |
+        tensorrt_llm/serve/openai_protocol/common.py |
+        tensorrt_llm/serve/openai_protocol/completions.py |
+        tensorrt_llm/serve/openai_protocol/control.py |
+        tensorrt_llm/serve/openai_protocol/disagg.py |
+        tensorrt_llm/serve/openai_protocol/embeddings.py |
+        tensorrt_llm/serve/openai_protocol/kv_cache.py |
+        tensorrt_llm/serve/openai_protocol/models.py |
+        tensorrt_llm/serve/openai_protocol/responses.py |
+        tensorrt_llm/serve/openai_protocol/sampling.py |
+        tensorrt_llm/serve/openai_protocol/tokenize.py |
+        tensorrt_llm/serve/openai_protocol/tools.py |
         tensorrt_llm/serve/openai_server.py |
         tensorrt_llm/serve/postprocess_handlers.py |
         tensorrt_llm/serve/responses_utils.py |
@@ -1144,13 +1177,46 @@ legacy-files: &legacy_files |
         tensorrt_llm/scaffolding/worker.py |
         tensorrt_llm/scheduling_params.py |
         tensorrt_llm/serve/__init__.py |
+        tensorrt_llm/serve/_base/msgspec_route.py |
+        tensorrt_llm/serve/_base/serving.py |
+        tensorrt_llm/serve/_base/sse.py |
+        tensorrt_llm/serve/api/anthropic/serving.py |
+        tensorrt_llm/serve/api/openai/chat/serving.py |
+        tensorrt_llm/serve/api/openai/completions/_merge.py |
+        tensorrt_llm/serve/api/openai/completions/serving.py |
+        tensorrt_llm/serve/api/openai/embeddings/serving.py |
+        tensorrt_llm/serve/api/openai/mm_encoder/serving.py |
+        tensorrt_llm/serve/api/openai/tokenize/serving.py |
+        tensorrt_llm/serve/api/ops/metrics/prometheus.py |
+        tensorrt_llm/serve/api/ops/metrics/serving.py |
+        tensorrt_llm/serve/api/ops/profile/serving.py |
+        tensorrt_llm/serve/api/ops/rl/serving.py |
+        tensorrt_llm/serve/api/ops/runtime_control/serving.py |
+        tensorrt_llm/serve/api/ops/worker/serving.py |
+        tensorrt_llm/serve/api/visual_gen/media.py |
+        tensorrt_llm/serve/api/visual_gen/protocol.py |
+        tensorrt_llm/serve/api/visual_gen/serving_base.py |
+        tensorrt_llm/serve/api/visual_gen/serving_images.py |
         tensorrt_llm/serve/chat_utils.py |
         tensorrt_llm/serve/cluster_storage.py |
         tensorrt_llm/serve/disagg_auto_scaling.py |
         tensorrt_llm/serve/harmony_adapter.py |
         tensorrt_llm/serve/metadata_server.py |
+        tensorrt_llm/serve/middleware/arrival_time.py |
+        tensorrt_llm/serve/middleware/user_middleware.py |
         tensorrt_llm/serve/openai_disagg_server.py |
-        tensorrt_llm/serve/openai_protocol.py |
+        tensorrt_llm/serve/openai_protocol/chat.py |
+        tensorrt_llm/serve/openai_protocol/common.py |
+        tensorrt_llm/serve/openai_protocol/completions.py |
+        tensorrt_llm/serve/openai_protocol/control.py |
+        tensorrt_llm/serve/openai_protocol/disagg.py |
+        tensorrt_llm/serve/openai_protocol/embeddings.py |
+        tensorrt_llm/serve/openai_protocol/kv_cache.py |
+        tensorrt_llm/serve/openai_protocol/models.py |
+        tensorrt_llm/serve/openai_protocol/responses.py |
+        tensorrt_llm/serve/openai_protocol/sampling.py |
+        tensorrt_llm/serve/openai_protocol/tokenize.py |
+        tensorrt_llm/serve/openai_protocol/tools.py |
         tensorrt_llm/serve/openai_server.py |
         tensorrt_llm/serve/postprocess_handlers.py |
         tensorrt_llm/serve/responses_utils.py |
@@ -1620,6 +1686,7 @@ repos:
         # Run mypy with no explicit paths so it type-checks the set declared in
         # pyproject.toml [tool.mypy] `files` (single source of truth). This
         # `files:` pattern only decides when the hook triggers — any change under
-        # the sampler source/test dirs.
+        # the sampler source/test dirs or the serving frontends. A PR that adds
+        # a typed module elsewhere extends this pattern in the same PR.
         pass_filenames: false
-        files: ^(tensorrt_llm/_torch/pyexecutor/sampler/|tests/unittest/_torch/sampler/)
+        files: ^(tensorrt_llm/_torch/pyexecutor/sampler/|tests/unittest/_torch/sampler/|tensorrt_llm/serve/|tensorrt_llm/commands/serve_cli/)

docs/source/features/sleep-wakeup.md

@@ -350,7 +350,7 @@ Principal implementation locations are:
 
 - Synchronous API: `tensorrt_llm/llmapi/llm.py`
 - Asynchronous API: `tensorrt_llm/_torch/async_llm.py`
-- Serving and authentication: `tensorrt_llm/serve/openai_server.py` and
+- Serving and authentication: `tensorrt_llm/serve/api/ops/runtime_control/` and
   `tensorrt_llm/serve/runtime_control_auth.py`
 - Workers: `tensorrt_llm/executor/base_worker.py` and
   `tensorrt_llm/executor/ray/gpu_worker.py`

jenkins/L0_MergeRequest.groovy

@@ -1948,10 +1948,24 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars)
         "tensorrt_llm/models/llama/",
         "tensorrt_llm/parameter.py",
         "tensorrt_llm/serve/__init__.py",
+        // _MsgspecRoute decodes the disaggregated frontend's msgpack bodies.
+        "tensorrt_llm/serve/_base/",
+        // The disaggregated /v1/messages runs through the Anthropic serving class.
+        "tensorrt_llm/serve/api/anthropic/api_router.py",
+        "tensorrt_llm/serve/api/anthropic/serving.py",
+        // Code moved out of tensorrt_llm/serve/openai_server.py keeps its trigger.
+        "tensorrt_llm/serve/api/openai/chat/",
+        "tensorrt_llm/serve/api/openai/completions/",
+        "tensorrt_llm/serve/api/openai/embeddings/",
+        "tensorrt_llm/serve/api/openai/mm_encoder/",
+        "tensorrt_llm/serve/api/openai/tokenize/",
+        // /health is the disaggregated readiness probe.
+        "tensorrt_llm/serve/api/ops/",
         "tensorrt_llm/serve/cluster_storage.py",
         "tensorrt_llm/serve/disagg_auto_scaling.py",
         "tensorrt_llm/serve/extensions/",
         "tensorrt_llm/serve/metadata_server.py",
+        "tensorrt_llm/serve/middleware/",
         "tensorrt_llm/serve/openai_client.py",
         "tensorrt_llm/serve/openai_disagg_server.py",
         "tensorrt_llm/serve/openai_disagg_service.py",

legacy-files.txt

@@ -409,13 +409,46 @@ tensorrt_llm/scaffolding/task_collection.py
 tensorrt_llm/scaffolding/worker.py
 tensorrt_llm/scheduling_params.py
 tensorrt_llm/serve/__init__.py
+tensorrt_llm/serve/_base/msgspec_route.py
+tensorrt_llm/serve/_base/serving.py
+tensorrt_llm/serve/_base/sse.py
+tensorrt_llm/serve/api/anthropic/serving.py
+tensorrt_llm/serve/api/openai/chat/serving.py
+tensorrt_llm/serve/api/openai/completions/_merge.py
+tensorrt_llm/serve/api/openai/completions/serving.py
+tensorrt_llm/serve/api/openai/embeddings/serving.py
+tensorrt_llm/serve/api/openai/mm_encoder/serving.py
+tensorrt_llm/serve/api/openai/tokenize/serving.py
+tensorrt_llm/serve/api/ops/metrics/prometheus.py
+tensorrt_llm/serve/api/ops/metrics/serving.py
+tensorrt_llm/serve/api/ops/profile/serving.py
+tensorrt_llm/serve/api/ops/rl/serving.py
+tensorrt_llm/serve/api/ops/runtime_control/serving.py
+tensorrt_llm/serve/api/ops/worker/serving.py
+tensorrt_llm/serve/api/visual_gen/media.py
+tensorrt_llm/serve/api/visual_gen/protocol.py
+tensorrt_llm/serve/api/visual_gen/serving_base.py
+tensorrt_llm/serve/api/visual_gen/serving_images.py
 tensorrt_llm/serve/chat_utils.py
 tensorrt_llm/serve/cluster_storage.py
 tensorrt_llm/serve/disagg_auto_scaling.py
 tensorrt_llm/serve/harmony_adapter.py
 tensorrt_llm/serve/metadata_server.py
+tensorrt_llm/serve/middleware/arrival_time.py
+tensorrt_llm/serve/middleware/user_middleware.py
 tensorrt_llm/serve/openai_disagg_server.py
-tensorrt_llm/serve/openai_protocol.py
+tensorrt_llm/serve/openai_protocol/chat.py
+tensorrt_llm/serve/openai_protocol/common.py
+tensorrt_llm/serve/openai_protocol/completions.py
+tensorrt_llm/serve/openai_protocol/control.py
+tensorrt_llm/serve/openai_protocol/disagg.py
+tensorrt_llm/serve/openai_protocol/embeddings.py
+tensorrt_llm/serve/openai_protocol/kv_cache.py
+tensorrt_llm/serve/openai_protocol/models.py
+tensorrt_llm/serve/openai_protocol/responses.py
+tensorrt_llm/serve/openai_protocol/sampling.py
+tensorrt_llm/serve/openai_protocol/tokenize.py
+tensorrt_llm/serve/openai_protocol/tools.py
 tensorrt_llm/serve/openai_server.py
 tensorrt_llm/serve/postprocess_handlers.py
 tensorrt_llm/serve/responses_utils.py

pyproject.toml

@@ -474,13 +474,46 @@ exclude = [
     "tensorrt_llm/scaffolding/worker.py",
     "tensorrt_llm/scheduling_params.py",
     "tensorrt_llm/serve/__init__.py",
+    "tensorrt_llm/serve/_base/msgspec_route.py",
+    "tensorrt_llm/serve/_base/serving.py",
+    "tensorrt_llm/serve/_base/sse.py",
+    "tensorrt_llm/serve/api/anthropic/serving.py",
+    "tensorrt_llm/serve/api/openai/chat/serving.py",
+    "tensorrt_llm/serve/api/openai/completions/_merge.py",
+    "tensorrt_llm/serve/api/openai/completions/serving.py",
+    "tensorrt_llm/serve/api/openai/embeddings/serving.py",
+    "tensorrt_llm/serve/api/openai/mm_encoder/serving.py",
+    "tensorrt_llm/serve/api/openai/tokenize/serving.py",
+    "tensorrt_llm/serve/api/ops/metrics/prometheus.py",
+    "tensorrt_llm/serve/api/ops/metrics/serving.py",
+    "tensorrt_llm/serve/api/ops/profile/serving.py",
+    "tensorrt_llm/serve/api/ops/rl/serving.py",
+    "tensorrt_llm/serve/api/ops/runtime_control/serving.py",
+    "tensorrt_llm/serve/api/ops/worker/serving.py",
+    "tensorrt_llm/serve/api/visual_gen/media.py",
+    "tensorrt_llm/serve/api/visual_gen/protocol.py",
+    "tensorrt_llm/serve/api/visual_gen/serving_base.py",
+    "tensorrt_llm/serve/api/visual_gen/serving_images.py",
     "tensorrt_llm/serve/chat_utils.py",
     "tensorrt_llm/serve/cluster_storage.py",
     "tensorrt_llm/serve/disagg_auto_scaling.py",
     "tensorrt_llm/serve/harmony_adapter.py",
     "tensorrt_llm/serve/metadata_server.py",
+    "tensorrt_llm/serve/middleware/arrival_time.py",
+    "tensorrt_llm/serve/middleware/user_middleware.py",
     "tensorrt_llm/serve/openai_disagg_server.py",
-    "tensorrt_llm/serve/openai_protocol.py",
+    "tensorrt_llm/serve/openai_protocol/chat.py",
+    "tensorrt_llm/serve/openai_protocol/common.py",
+    "tensorrt_llm/serve/openai_protocol/completions.py",
+    "tensorrt_llm/serve/openai_protocol/control.py",
+    "tensorrt_llm/serve/openai_protocol/disagg.py",
+    "tensorrt_llm/serve/openai_protocol/embeddings.py",
+    "tensorrt_llm/serve/openai_protocol/kv_cache.py",
+    "tensorrt_llm/serve/openai_protocol/models.py",
+    "tensorrt_llm/serve/openai_protocol/responses.py",
+    "tensorrt_llm/serve/openai_protocol/sampling.py",
+    "tensorrt_llm/serve/openai_protocol/tokenize.py",
+    "tensorrt_llm/serve/openai_protocol/tools.py",
     "tensorrt_llm/serve/openai_server.py",
     "tensorrt_llm/serve/postprocess_handlers.py",
     "tensorrt_llm/serve/responses_utils.py",
@@ -862,6 +895,29 @@ files = [
     "tensorrt_llm/_torch/pyexecutor/sampler/token_ban.py",
     "tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py",
     "tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py",
+    "tensorrt_llm/llmapi/engine_client.py",
+    "tensorrt_llm/serve/_base/app_builder.py",
+    "tensorrt_llm/serve/_base/chat_backend.py",
+    "tensorrt_llm/serve/_base/errors.py",
+    "tensorrt_llm/serve/_base/request_context.py",
+    "tensorrt_llm/serve/_base/routing.py",
+    "tensorrt_llm/serve/api/anthropic/api_router.py",
+    "tensorrt_llm/serve/api/openai/chat/api_router.py",
+    "tensorrt_llm/serve/api/openai/completions/api_router.py",
+    "tensorrt_llm/serve/api/openai/embeddings/api_router.py",
+    "tensorrt_llm/serve/api/openai/mm_encoder/api_router.py",
+    "tensorrt_llm/serve/api/openai/tokenize/api_router.py",
+    "tensorrt_llm/serve/api/ops/control/api_router.py",
+    "tensorrt_llm/serve/api/ops/metrics/api_router.py",
+    "tensorrt_llm/serve/api/ops/profile/api_router.py",
+    "tensorrt_llm/serve/api/ops/rl/api_router.py",
+    "tensorrt_llm/serve/api/ops/runtime_control/api_router.py",
+    "tensorrt_llm/serve/api/ops/system/api_router.py",
+    "tensorrt_llm/serve/api/ops/system/serving.py",
+    "tensorrt_llm/serve/api/ops/worker/api_router.py",
+    "tensorrt_llm/serve/api/visual_gen/api_router.py",
+    "tensorrt_llm/serve/api/visual_gen/errors.py",
+    "tensorrt_llm/serve/middleware/register.py",
     "tests/unittest/_torch/sampler/test_torch_sampler.py",
     "tests/unittest/_torch/sampler/test_token_ban.py",
     "tests/unittest/_torch/sampler/test_beam_search.py",

ruff-legacy-baseline.json

@@ -528,6 +528,15 @@
   "tensorrt_llm/scaffolding/task.py": {
     "F821": 1
   },
+  "tensorrt_llm/serve/_base/sse.py": {
+    "F821": 1
+  },
+  "tensorrt_llm/serve/api/openai/chat/serving.py": {
+    "F821": 1
+  },
+  "tensorrt_llm/serve/api/openai/completions/_merge.py": {
+    "F821": 1
+  },
   "tensorrt_llm/serve/disagg_auto_scaling.py": {
     "D205": 2,
     "D212": 2
@@ -539,6 +548,9 @@
     "D212": 19,
     "D415": 1
   },
+  "tensorrt_llm/serve/middleware/arrival_time.py": {
+    "D212": 1
+  },
   "tensorrt_llm/serve/openai_disagg_server.py": {
     "D210": 1,
     "F811": 1

ruff-legacy.toml

@@ -426,13 +426,46 @@ include = [
     "tensorrt_llm/scaffolding/worker.py",
     "tensorrt_llm/scheduling_params.py",
     "tensorrt_llm/serve/__init__.py",
+    "tensorrt_llm/serve/_base/msgspec_route.py",
+    "tensorrt_llm/serve/_base/serving.py",
+    "tensorrt_llm/serve/_base/sse.py",
+    "tensorrt_llm/serve/api/anthropic/serving.py",
+    "tensorrt_llm/serve/api/openai/chat/serving.py",
+    "tensorrt_llm/serve/api/openai/completions/_merge.py",
+    "tensorrt_llm/serve/api/openai/completions/serving.py",
+    "tensorrt_llm/serve/api/openai/embeddings/serving.py",
+    "tensorrt_llm/serve/api/openai/mm_encoder/serving.py",
+    "tensorrt_llm/serve/api/openai/tokenize/serving.py",
+    "tensorrt_llm/serve/api/ops/metrics/prometheus.py",
+    "tensorrt_llm/serve/api/ops/metrics/serving.py",
+    "tensorrt_llm/serve/api/ops/profile/serving.py",
+    "tensorrt_llm/serve/api/ops/rl/serving.py",
+    "tensorrt_llm/serve/api/ops/runtime_control/serving.py",
+    "tensorrt_llm/serve/api/ops/worker/serving.py",
+    "tensorrt_llm/serve/api/visual_gen/media.py",
+    "tensorrt_llm/serve/api/visual_gen/protocol.py",
+    "tensorrt_llm/serve/api/visual_gen/serving_base.py",
+    "tensorrt_llm/serve/api/visual_gen/serving_images.py",
     "tensorrt_llm/serve/chat_utils.py",
     "tensorrt_llm/serve/cluster_storage.py",
     "tensorrt_llm/serve/disagg_auto_scaling.py",
     "tensorrt_llm/serve/harmony_adapter.py",
     "tensorrt_llm/serve/metadata_server.py",
+    "tensorrt_llm/serve/middleware/arrival_time.py",
+    "tensorrt_llm/serve/middleware/user_middleware.py",
     "tensorrt_llm/serve/openai_disagg_server.py",
-    "tensorrt_llm/serve/openai_protocol.py",
+    "tensorrt_llm/serve/openai_protocol/chat.py",
+    "tensorrt_llm/serve/openai_protocol/common.py",
+    "tensorrt_llm/serve/openai_protocol/completions.py",
+    "tensorrt_llm/serve/openai_protocol/control.py",
+    "tensorrt_llm/serve/openai_protocol/disagg.py",
+    "tensorrt_llm/serve/openai_protocol/embeddings.py",
+    "tensorrt_llm/serve/openai_protocol/kv_cache.py",
+    "tensorrt_llm/serve/openai_protocol/models.py",
+    "tensorrt_llm/serve/openai_protocol/responses.py",
+    "tensorrt_llm/serve/openai_protocol/sampling.py",
+    "tensorrt_llm/serve/openai_protocol/tokenize.py",
+    "tensorrt_llm/serve/openai_protocol/tools.py",
     "tensorrt_llm/serve/openai_server.py",
     "tensorrt_llm/serve/postprocess_handlers.py",
     "tensorrt_llm/serve/responses_utils.py",

tensorrt_llm/_torch/pyexecutor/profiling.py

@@ -390,7 +390,8 @@ def stop_profile(self) -> None:
           HTTP handler the server layer is responsible for tickling
           the executor loop so the stop iteration actually runs even
           when the engine would otherwise be idle (see
-          ``OpenAIServer.stop_profile``).
+          ``ProfileServing.stop_profile`` in
+          ``tensorrt_llm/serve/api/ops/profile/serving.py``).
         """
         executor = self._executor
         # Apply locally so diagnostic state (profile_stop_iters,

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -14,8 +14,8 @@
 from contextlib import ExitStack, contextmanager, nullcontext
 from enum import IntEnum
 from queue import Queue
-from typing import (TYPE_CHECKING, Callable, Dict, Iterable, Iterator, List,
-                    Literal, NoReturn, Optional, Tuple, Union)
+from typing import (TYPE_CHECKING, Any, Callable, Dict, Iterable, Iterator,
+                    List, Literal, NoReturn, Optional, Tuple, Union)
 
 import torch
 from strenum import StrEnum
@@ -2749,6 +2749,19 @@ def get_openengine_discovery(self) -> dict:
         """Return immutable-at-startup OpenEngine node and publisher discovery."""
         return self._openengine_discovery
 
+    def set_openengine_discovery(self, discovery: dict) -> None:
+        """Record OpenEngine node and KV-event publisher discovery metadata.
+
+        Called once at startup by the OpenEngine node-discovery hook
+        (tensorrt_llm/grpc/openengine/node.py); get_openengine_discovery()
+        returns it.
+        """
+        self._openengine_discovery = discovery
+
+    def set_openengine_node_server(self, server: Any) -> None:
+        """Adopt the OpenEngine node metadata server; shutdown closes it."""
+        self._openengine_node_server = server
+
     def get_kv_cache_load(self) -> dict:
         """Return the scheduler-owned KV load snapshot.