Skip to content

[TRTLLM-17267][refactor] Move the completions handler into a serving class - #21640

Draft
JunyiXu-nv wants to merge 41 commits into
NVIDIA:mainfrom
JunyiXu-nv:dev-junyix-feat-serve-completions-serving
Draft

JunyiXu-nv wants to merge 41 commits into
NVIDIA:mainfrom
JunyiXu-nv:dev-junyix-feat-serve-completions-serving

Conversation

@JunyiXu-nv

Copy link
Copy Markdown
Collaborator

@coderabbitai summary

Stacked on #21575 ([TRTLLM-17267][refactor] Move the embeddings, mm_encoder and tokenize handlers into serving classes) and #21392 ([TRTLLM-17273][chore] Choose the trtllm-serve post-processing path once per server): the commits up to the integration merge 0d6a309791 (#21575's branch at 600bcd503b, merged with #21392's own commits e6df420796 and 5806927953) belong to those PRs and disappear from this diff when they merge; this PR's own change is the last 7 commits (d95c207e79 to e5bf1275f8). The reviewable range is git diff 0d6a309791...e5bf1275f8. Draft until both land; it is then regenerated on the latest main and marked ready.

Description

Part of the trtllm-serve refactor (epic TRTLLM-17056, task TRTLLM-17267: move the handlers off OpenAIServer into one serving class per endpoint). This PR moves openai_completion (POST /v1/completions) into CompletionServing (tensorrt_llm/serve/api/openai/completions/serving.py), served through completions/api_router.py and the routing table for the LLM role group (no role, CONTEXT, GENERATION). Its five nested functions become module-level functions in completions/_merge.py whose free variables are keyword-only parameters, so each one can be unit-tested on its own; that lift is the only non-verbatim step and is spelled out below. The six helpers the handler shares with the chat and Responses handlers (postproc_worker_enabled, _extract_metrics, _logit_bias_vocab_size, _validate_internal_disagg_request, _has_cache_transceiver_config, _apply_spec_decode_stats_opt_in) move into GenerateBaseServing; OpenAIServer keeps each as an alias of the moved object (never a copy) while chat and Responses still call it there.

It builds on #21575 (the serving-class bases, the first three serving classes and _generate_serving_kwargs()) and #21392 (one post-processing strategy object per server, which the serving object now receives as its postproc argument).

GenerateBaseServing gains ten keyword-only constructor arguments, one per attribute the moved code reads: collect_perf_metrics, metrics_collector, adjusted_steady_clock, iteration_stats_wakeup_event, iteration_stats_collector_task_getter, internal_disagg_auth_key, per_request_spec_decode_stats, spec_decode_num_spec_tokens, input_proc_executor, postproc. The composition root passes its own objects through _generate_serving_kwargs(), so the serving object shares the root's clock, wake-up event, thread pool, metrics collector and strategy. One argument needs a word: on main, _extract_metrics wakes the iteration-stats collector only while self._iteration_stats_collector_task is not None, and that task is created by the lifespan after the serving objects exist and is bound on the root. A value copied at construction would always be None and the collector would never be woken from completions, so the root passes a getter (iteration_stats_collector_task_getter=lambda: self._iteration_stats_collector_task) and GenerateBaseServing exposes it as the read-only property _iteration_stats_collector_task; the moved body stays verbatim and reads the root's current binding. The three getattr(self, ...) probes in the moved helpers (_internal_disagg_auth_key, _per_request_spec_decode_stats, _spec_decode_num_spec_tokens) are kept verbatim; the constructor makes the attributes exist, so the defaults are never taken, and a test pins that the auth key reaches the serving object.

Behavior

No behavior change. Status codes, error bodies, the SSE error frame of the streaming path, the msgpack route class, logging and the environment-variable special case (TRTLLM_KVCACHE_TIME_OUTPUT_PATH) are as before; the route set and the middleware order are unchanged (the route snapshot tests/unittest/api_stability/references/trtllm_serve_routes.yaml did not change).

Moves

Pure-move check
## Move spec

```yaml
moves:
- from: tensorrt_llm/serve/openai_server.py
  to: tensorrt_llm/serve/_base/serving.py
  group: B
  append: true
  methods:
    from_class: OpenAIServer
    to_class: GenerateBaseServing
    names:
    - postproc_worker_enabled
    - _extract_metrics
    - _logit_bias_vocab_size
    - _validate_internal_disagg_request
    - _has_cache_transceiver_config
    - _apply_spec_decode_stats_opt_in
- from: tensorrt_llm/serve/openai_server.py
  to: tensorrt_llm/serve/api/openai/completions/serving.py
  group: B
  append: true
  methods:
    from_class: OpenAIServer
    to_class: CompletionServing
    names:
    - openai_completion
lifted:
- from: tensorrt_llm/serve/openai_server.py::OpenAIServer.openai_completion.completion_response
  to: tensorrt_llm/serve/api/openai/completions/_merge.py::completion_response
  free_variables:
    self: serving
    disaggregated_params: disaggregated_params
    raw_request: raw_request
- from: tensorrt_llm/serve/openai_server.py::OpenAIServer.openai_completion.merge_completion_responses
  to: tensorrt_llm/serve/api/openai/completions/_merge.py::merge_completion_responses
  free_variables:
    self: serving
- from: tensorrt_llm/serve/openai_server.py::OpenAIServer.openai_completion.completion_generator
  to: tensorrt_llm/serve/api/openai/completions/_merge.py::completion_generator
  free_variables:
    self: serving
    raw_request: raw_request
- from: tensorrt_llm/serve/openai_server.py::OpenAIServer.openai_completion.merge_generators
  to: tensorrt_llm/serve/api/openai/completions/_merge.py::merge_generators
  free_variables: {}
- from: tensorrt_llm/serve/openai_server.py::OpenAIServer.openai_completion.generator_wrapper
  to: tensorrt_llm/serve/api/openai/completions/_merge.py::generator_wrapper
  free_variables:
    self: serving
    raw_request: raw_request
lifted_call_sites:
- symbol: tensorrt_llm/serve/openai_server.py::OpenAIServer.openai_completion
borrowed:
- name: postproc_worker_enabled
  alias: GenerateBaseServing.postproc_worker_enabled
- name: _extract_metrics
  alias: GenerateBaseServing._extract_metrics
- name: _logit_bias_vocab_size
  alias: GenerateBaseServing._logit_bias_vocab_size
- name: _validate_internal_disagg_request
  alias: GenerateBaseServing._validate_internal_disagg_request
- name: _has_cache_transceiver_config
  alias: GenerateBaseServing._has_cache_transceiver_config
- name: _apply_spec_decode_stats_opt_in
  alias: GenerateBaseServing._apply_spec_decode_stats_opt_in
modified:
- symbol: tensorrt_llm/serve/_base/serving.py::GenerateBaseServing.__init__
  reason: new keyword-only arguments for the moved helpers
- symbol: tensorrt_llm/serve/openai_server.py::OpenAIServer._generate_serving_kwargs
  reason: the new GenerateBaseServing arguments
- symbol: tensorrt_llm/serve/openai_server.py::OpenAIServer.__init__
  reason: builds CompletionServing in the LLM branch
- symbol: tensorrt_llm/serve/openai_server.py::OpenAIServer.register_routes
  reason: /v1/completions comes from the completions router
- symbol: tensorrt_llm/serve/_base/routing.py::ROUTERS
  reason: the completions router joins the routing table for the LLM role group
added:
- symbol: tensorrt_llm/serve/api/openai/completions/serving.py::CompletionServing
  reason: serving class of /v1/completions (no constructor of its own)
- symbol: tensorrt_llm/serve/api/openai/completions/api_router.py::*
  reason: new router module
- symbol: tensorrt_llm/serve/_base/serving.py::GenerateBaseServing._iteration_stats_collector_task
  reason: 'property behind the moved _extract_metrics: the composition root''s lifespan
    creates the iteration-stats collector task after the serving objects exist, so
    they read the root''s current binding through an injected getter'
rewrite_internal_imports: true

Pure-move check: base 0d6a309, head e5bf127, tools c1d68659fb46

Symbol From To Result
ROUTERS tensorrt_llm/serve/_base/routing.py tensorrt_llm/serve/_base/routing.py modified (declared)
GenerateBaseServing.__init__ tensorrt_llm/serve/_base/serving.py tensorrt_llm/serve/_base/serving.py modified (declared)
OpenAIServer._apply_spec_decode_stats_opt_in tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/_base/serving.py verbatim
OpenAIServer._extract_metrics tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/_base/serving.py verbatim
OpenAIServer._has_cache_transceiver_config tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/_base/serving.py verbatim
OpenAIServer._logit_bias_vocab_size tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/_base/serving.py verbatim
OpenAIServer._validate_internal_disagg_request tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/_base/serving.py verbatim
OpenAIServer.postproc_worker_enabled tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/_base/serving.py verbatim
OpenAIServer.openai_completion.completion_generator tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/api/openai/completions/_merge.py lifted
OpenAIServer.openai_completion.completion_response tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/api/openai/completions/_merge.py lifted
OpenAIServer.openai_completion.generator_wrapper tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/api/openai/completions/_merge.py lifted
OpenAIServer.openai_completion.merge_completion_responses tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/api/openai/completions/_merge.py lifted
OpenAIServer.openai_completion.merge_generators tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/api/openai/completions/_merge.py lifted
OpenAIServer.openai_completion tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/api/openai/completions/serving.py verbatim
OpenAIServer.__init__ tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/openai_server.py modified (declared)
OpenAIServer._apply_spec_decode_stats_opt_in tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/openai_server.py borrowed
OpenAIServer._extract_metrics tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/openai_server.py borrowed
OpenAIServer._generate_serving_kwargs tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/openai_server.py modified (declared)
OpenAIServer._has_cache_transceiver_config tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/openai_server.py borrowed
OpenAIServer._logit_bias_vocab_size tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/openai_server.py borrowed
OpenAIServer._validate_internal_disagg_request tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/openai_server.py borrowed
OpenAIServer.postproc_worker_enabled tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/openai_server.py borrowed
OpenAIServer.register_routes tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/openai_server.py modified (declared)
GenerateBaseServing._iteration_stats_collector_task — tensorrt_llm/serve/_base/serving.py added
tensorrt_llm/serve/api/openai/completions/api_router.py::* — tensorrt_llm/serve/api/openai/completions/api_router.py added
CompletionServing — tensorrt_llm/serve/api/openai/completions/serving.py added

Added

  • CompletionServing: serving class of /v1/completions (no constructor of its own)
  • tensorrt_llm/serve/api/openai/completions/api_router.py::*: new router module
  • GenerateBaseServing._iteration_stats_collector_task: property behind the moved _extract_metrics: the composition root's lifespan creates the iteration-stats collector task after the serving objects exist, so they read the root's current binding through an injected getter

Modified (declared)

  • OpenAIServer.__init__: builds CompletionServing in the LLM branch
    --- base
    +++ head
    @@ -188,6 +188,9 @@
    from tensorrt_llm.serve.api.openai.tokenize import api_router as tokenize_router
    from tensorrt_llm.serve.api.openai.tokenize.serving import TokenizeServing
    setattr(self.app.state, tokenize_router.APP_STATE_KEY, TokenizeServing(model=self.model, tokenizer=self.tokenizer))
    •    from tensorrt_llm.serve.api.openai.completions import api_router as completions_router
      
    •    from tensorrt_llm.serve.api.openai.completions.serving import CompletionServing
      
    •    setattr(self.app.state, completions_router.APP_STATE_KEY, CompletionServing(**self._generate_serving_kwargs()))
         self._register_launcher_forwarding_routes()
         self.register_routes()
      
      attach_routers(self.app, ServerCapabilities.for_role(self.server_role))
  • OpenAIServer._generate_serving_kwargs: the new GenerateBaseServing arguments
    --- base
    +++ head
    @@ -1,3 +1,3 @@
    def _generate_serving_kwargs(self) -> Dict[str, Any]:
    """Constructor arguments of every GenerateBaseServing, read once."""
    • return dict(model=self.model, generator=self.generator, tokenizer=self.tokenizer, model_config=self.model_config)
    • return dict(model=self.model, generator=self.generator, tokenizer=self.tokenizer, model_config=self.model_config, collect_perf_metrics=self._collect_perf_metrics, metrics_collector=self.metrics_collector, adjusted_steady_clock=self._adjusted_steady_clock, iteration_stats_wakeup_event=self._iteration_stats_wakeup_event, iteration_stats_collector_task_getter=lambda: self._iteration_stats_collector_task, internal_disagg_auth_key=self._internal_disagg_auth_key, per_request_spec_decode_stats=self._per_request_spec_decode_stats, spec_decode_num_spec_tokens=self._spec_decode_num_spec_tokens, input_proc_executor=self._input_proc_executor, postproc=self._postproc)
  • OpenAIServer.register_routes: /v1/completions comes from the completions router
    --- base
    +++ head
    @@ -14,7 +14,6 @@
    else:
    logger.warning('Resource governor is disabled: the executor backend does not provide a resource_governor_queue (e.g. RPC orchestrator mode or explicit opt-out). Endpoints under /_resource_governor/ will not be available.')
    self.resource_governor = None
    • self.app.add_api_route('/v1/completions', self.openai_completion, methods=['POST'])
      self.app.add_api_route('/v1/chat/completions', self.openai_chat if not self.use_harmony else self.chat_harmony, methods=['POST'])
      if not self.use_harmony:
      self.app.add_api_route('/v1/messages', self.anthropic_messages, methods=['POST'])
  • GenerateBaseServing.__init__: new keyword-only arguments for the moved helpers
    --- base
    +++ head
    @@ -1,5 +1,16 @@
    -def init(self, *, model: str, generator: Any, tokenizer: Any, model_config: Any) -> None:
    +def init(self, *, model: str, generator: Any, tokenizer: Any, model_config: Any, collect_perf_metrics: bool, metrics_collector: Optional[MetricsCollector], adjusted_steady_clock: AdjustedSteadyClock, iteration_stats_wakeup_event: asyncio.Event, iteration_stats_collector_task_getter: Callable[[], Optional[asyncio.Task]], internal_disagg_auth_key: Optional[str], per_request_spec_decode_stats: bool, spec_decode_num_spec_tokens: Optional[int], input_proc_executor: ThreadPoolExecutor, postproc: PostprocStrategy) -> None:
    self.model = model
    self.generator = generator
    self.tokenizer = tokenizer
    self.model_config = model_config
    • self._collect_perf_metrics = collect_perf_metrics
    • self.metrics_collector = metrics_collector
    • self._adjusted_steady_clock = adjusted_steady_clock
    • self._iteration_stats_wakeup_event = iteration_stats_wakeup_event
    • self._iteration_stats_collector_task_getter = iteration_stats_collector_task_getter
    • self._internal_disagg_auth_key = internal_disagg_auth_key
    • self._per_request_spec_decode_stats = per_request_spec_decode_stats
    • self._spec_decode_num_spec_tokens = spec_decode_num_spec_tokens
    • self._input_proc_executor = input_proc_executor
    • self._postproc = postproc
    • self._is_visual_gen = False
  • ROUTERS: the completions router joins the routing table for the LLM role group
    --- base
    +++ head
    @@ -1 +1 @@
    -ROUTERS: tuple[RouterSpec, ...] = (RouterSpec('tensorrt_llm.serve.api.ops.system.api_router', roles=ALL_ROLE_GROUPS), RouterSpec('tensorrt_llm.serve.api.openai.embeddings.api_router', roles=frozenset({RoleGroup.EMBEDDING})), RouterSpec('tensorrt_llm.serve.api.openai.mm_encoder.api_router', roles=frozenset({RoleGroup.MM_ENCODER})), RouterSpec('tensorrt_llm.serve.api.openai.tokenize.api_router', roles=frozenset({RoleGroup.LLM})))
    +ROUTERS: tuple[RouterSpec, ...] = (RouterSpec('tensorrt_llm.serve.api.ops.system.api_router', roles=ALL_ROLE_GROUPS), RouterSpec('tensorrt_llm.serve.api.openai.embeddings.api_router', roles=frozenset({RoleGroup.EMBEDDING})), RouterSpec('tensorrt_llm.serve.api.openai.mm_encoder.api_router', roles=frozenset({RoleGroup.MM_ENCODER})), RouterSpec('tensorrt_llm.serve.api.openai.tokenize.api_router', roles=frozenset({RoleGroup.LLM})), RouterSpec('tensorrt_llm.serve.api.openai.completions.api_router', roles=frozenset({RoleGroup.LLM})))

7 definitions moved verbatim; 0 differences; 5 lifted; 6 borrowed; 3 added; 5 modified (declared); shim exports 0 names


</details>

The `verbatim` row of `openai_completion` comes from the comparison that removes the five nested definitions from the original and the keyword arguments appended at the five call sites from the moved method; the independent comparison listed under Test Coverage agrees with it.

The lifted functions are allowed to differ from the original nested functions in exactly these ways, and in nothing else:

1. Each nested function becomes a module-level function in `completions/_merge.py`, in the original order, with its `async` keyword, parameter list, annotations, defaults and return annotation unchanged (none had decorators or docstrings).
2. A function that had free variables gets a bare `*` appended to its parameter list, then one keyword-only parameter per free variable, in a fixed order, with no annotation and no default: `serving` for `self`, and `raw_request` and `disaggregated_params` under their own names.
3. In the body, every `self` becomes `serving`; `raw_request` and `disaggregated_params` keep their names. Nothing else in the body changes.
4. Nested functions move inside their parent unchanged: `producer` stays inside `merge_generators`; its free variables `result_queue` and `finished` are locals of `merge_generators`.
5. In `CompletionServing.openai_completion` the five definitions are removed, and every call of a lifted function gets the matching keyword arguments appended: `serving=self`, `raw_request=raw_request`, `disaggregated_params=disaggregated_params`. No other change to the method (yapf re-wrapped lines; the comparison is on syntax trees).
6. `completions/serving.py` imports the five functions from `completions/_merge.py`; `_merge.py` imports what the functions' signatures and bodies name (`RequestOutput`, `PostprocParams`, the completion response models, `logger`, `asyncio`, `json`, `traceback`, `typing`). `import tensorrt_llm.serve.api.openai.completions._merge` succeeds on its own.

Free variables of the nested functions at the integration base (`tensorrt_llm/serve/openai_server.py`):

| Nested function | Lines | Free variables | Called at |
|---|---|---|---|
| `completion_response(promise, postproc_params)` | 2612-2624 | `self`, `raw_request`, `disaggregated_params` (assigned once, at 2742, before the prompt loop) | 2811 |
| `merge_completion_responses(responses)` | 2626-2654 | `self` (for `self.model`) | 2814 |
| `completion_generator(promise, params)` | 2656-2679 | `self`, `raw_request` | 2801 |
| `merge_generators(generators)` | 2681-2698 | none | 2804 |
| `merge_generators.producer(generator, idx)` | 2685-2688 | `result_queue`, `finished` (locals of `merge_generators`) | inside `merge_generators` |
| `generator_wrapper(generator)` | 2700-2707 | `self`, `raw_request` | 2807 |

Resulting signatures:

```python
async def completion_response(promise: RequestOutput, postproc_params: Optional[PostprocParams], *,
                              serving, disaggregated_params, raw_request) -> CompletionResponse
def merge_completion_responses(responses: List[CompletionResponse], *, serving) -> CompletionResponse
async def completion_generator(promise: RequestOutput, params: Optional[PostprocParams], *, serving, raw_request)
async def merge_generators(generators: List[AsyncIterator[Any]])
async def generator_wrapper(generator: AsyncIterator[Any], *, serving, raw_request)

Why this preserves behavior: a closure reads its free variables when its body runs; a keyword argument is bound when the call creates the coroutine or async generator. All five call sites run after the last assignment of each free variable (disaggregated_params is assigned once, before the prompt loop, and self and raw_request are never reassigned in the method), so each function sees the same objects the closure would have seen.

Kept shims: none. Release notes: this PR moves methods and nested functions only, no module-level name, so no import path changes; the methods that left OpenAIServer are listed below. History: git blame -C -C tensorrt_llm/serve/api/openai/completions/serving.py and ... completions/_merge.py attribute the moved lines to their original commits.

Removed from OpenAIServer

  • openai_completion → CompletionServing.openai_completion (tensorrt_llm/serve/api/openai/completions/serving.py), reached through app.state.openai_serving_completions on an LLM-role server; its nested functions → tensorrt_llm/serve/api/openai/completions/_merge.py.
  • postproc_worker_enabled (property), _extract_metrics, _logit_bias_vocab_size, _validate_internal_disagg_request, _has_cache_transceiver_config, _apply_spec_decode_stats_opt_in → GenerateBaseServing (tensorrt_llm/serve/_base/serving.py); still reachable on OpenAIServer as aliases of the moved objects (name = GenerateBaseServing.name) while the chat and Responses handlers call them there; tests/unittest/serve/test_openai_server_borrowed_helpers.py pins the identity of all six and that postproc_worker_enabled is still a property (the composition root evaluates it while building the server).

CI bookkeeping

  • Multi-GPU table (getMultiGpuFileChanged in jenkins/L0_MergeRequest.groovy): added tensorrt_llm/serve/api/openai/completions/, so the completions code that left tensorrt_llm/serve/openai_server.py keeps triggering the disaggregated stages.
  • CODEOWNERS: none (/tensorrt_llm/serve covers the new directory).
  • legacy-files.txt: tensorrt_llm/serve/api/openai/completions/serving.py and tensorrt_llm/serve/api/openai/completions/_merge.py (Group B, like openai_server.py their code comes from); api_router.py and __init__.py are new code (Group A, ruff); derived configs regenerated with scripts/legacy_utils.py gen-configs (ruff-legacy.toml, the managed blocks of pyproject.toml and .pre-commit-config.yaml); ruff-legacy-baseline.json: one new entry, tensorrt_llm/serve/api/openai/completions/_merge.py: F821 1 (the anext of generator_wrapper, one of the three F821 recorded under openai_server.py); the openai_server.py entry is left as it is (the lint prints a tip to tighten it) because lint-update-violations rewrites about 50 unrelated entries and is not run in a move. Needs infra-devs.
  • mypy: tensorrt_llm/serve/api/openai/completions/api_router.py joins [tool.mypy] files; serving.py and _merge.py stay out like the other modules that hold moved code without annotations.
  • Test lists: no change. tests/unittest/serve/api/openai/completions/test_merge.py and test_serving.py are collected by the unittest/serve entry of tests/integration/test_lists/test-db/l0_cpu.yml (both carry pytest.mark.cpu_only; python scripts/test_to_stage_mapping.py --tests "unittest/serve" names the CPU-Generic stages). The rewritten tests keep their registrations: tests/unittest/disaggregated/test_openai_server_stream_errors.py and test_disagg_internal_auth.py (unittest/disaggregated, CPU), tests/unittest/llmapi/test_kv_hints.py (CPU), tests/unittest/llmapi/test_llm.py::test_openai_completion_list_prompt_stream_reuses_stream_metadata (A100 stage) and tests/unittest/executor/test_spec_decode_stats_payload.py (A10 stage); the two GPU-stage files were run locally, they need no GPU for these tests.
  • tests/unittest/serve/test_no_object_new.py (the ratchet against servers built with object.__new__): the rows of the four rewritten sites are deleted (test_openai_server_stream_errors.py and test_spec_decode_stats_payload.py leave the baseline entirely); the fifth site this PR touched, test_kv_hints.py::test_http_handlers_forward_hints, keeps its row because its chat, harmony and Responses cases still build an OpenAIServer (only the completions case now builds a CompletionServing).
  • tests/test_common/serve/app_factory.py: AggConfig gains internal_disagg_auth_key and cache_transceiver_backend, so a test can build a worker whose auth key must reach the serving object (test-only).
  • Tests that build GenerateBaseServing directly or re-enact the composition root by hand follow the constructor: tests/unittest/serve/_base/test_serving.py requires every argument, pins the attribute each one is stored as and checks the collector-task property; tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py and test_serve_multi_frontend_helpers.py put a CompletionServing with placeholder dependencies on app.state before attaching the role's routers (the completions router refuses to attach without it); the OpenEngine Prometheus test's constructor stub (tests/unittest/grpc/openengine/test_metrics.py) sets model_config too, since the LLM role now reads it when it builds the completions serving object.

What changes in /openapi.json

Only the order of paths: /v1/completions is now registered by its router after the routes register_routes adds and after /_internal/tokenize, so it moves from before /v1/chat/completions to after them. Operation id, parameters, body schema and response schema are unchanged (the route function keeps the handler's name and signature).

Test Coverage

All commands ran from the worktree against this branch's code, on CPU (no GPU on the host), with base = the integration merge 0d6a309791.

  • import tensorrt_llm; tensorrt_llm.__file__ printed the worktree path; python -c "import tensorrt_llm.serve.api.openai.completions._merge" succeeds.
  • Pure-move check against 0d6a309791 (output above): 7 definitions moved verbatim, 0 differences, 5 lifted, 6 borrowed, 3 added, 5 modified (declared).
  • Independent comparison of every moved definition, outside the check above: before formatting, the six helper methods are byte-identical after dedent, openai_completion differs only by the five removed definitions and the keyword arguments appended at its five call sites, and each lifted function differs from its nested original only by the appended keyword-only parameters and self → serving; after formatting, syntax-tree comparison with imports dropped: 12/12 identical (the six helpers; openai_completion after removing the nested definitions on one side and the appended keywords on the other; the five lifted functions after reversing the lift).
  • New tests: pytest -m cpu_only tests/unittest/serve/api/openai/completions -q → 24 passed (test_merge.py: usage and prompt-id merging, generator merging with producers finished, the SSE error frame and [DONE] with nothing raised, streaming chunks plus one metrics record, the first-token stamp before the second item, context-only prompt ids for three request types, the error path; test_serving.py: all fourteen constructor arguments required and keyword-only, the shared root objects (_adjusted_steady_clock, _iteration_stats_wakeup_event, _input_proc_executor, metrics_collector, _postproc) plus the live collector-task binding, the route on the three LLM roles, several prompts in one non-streaming request merged into one response (one choice per prompt, summed usage; this test fails when the merge call loses its serving argument or the merge drops a prompt), several streamed prompts with one stream id and a single [DONE], msgpack bodies, the wired internal disaggregated auth key (400 "Invalid internal") with its control (no key: accepted), a finished request waking the stats collector with its control, 404 on the multimodal-encoder and embedding roles, attach_router without the serving object).
  • Rewritten tests and guards: pytest tests/unittest/disaggregated/test_openai_server_stream_errors.py tests/unittest/disaggregated/test_disagg_internal_auth.py tests/unittest/executor/test_spec_decode_stats_payload.py tests/unittest/llmapi/test_kv_hints.py "tests/unittest/llmapi/test_llm.py::test_openai_completion_list_prompt_stream_reuses_stream_metadata" tests/unittest/serve/test_openai_server_borrowed_helpers.py tests/unittest/serve/test_no_object_new.py tests/unittest/serve/api/openai/mm_encoder/test_serving.py tests/unittest/serve/test_serve_test_helpers.py tests/unittest/serve/test_layering.py tests/unittest/api_stability/test_serve_routes.py -q → 393 passed (together with the new tests).
  • Adapted tests: pytest -m cpu_only tests/unittest/serve/_base/test_serving.py tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py tests/unittest/llmapi/test_serve_multi_frontend_helpers.py tests/unittest/llmapi/test_serve_frontend_affinity.py -q → 26 passed; pytest -m cpu_only "tests/unittest/grpc/openengine/test_metrics.py::test_prometheus_contract[http_enablement]" → 1 passed.
  • Two FastAPI versions (fastapi==0.136.3 and 0.139.2): pytest tests/unittest/api_stability/test_serve_routes.py tests/unittest/serve/api/openai/completions/test_serving.py tests/unittest/serve/api/openai/tokenize/test_serving.py tests/unittest/serve/test_middleware_order.py -q → 80 passed on each.
  • Patch sites: git grep over tests/ for openai_server.<name>, the quoted names, setattr(<module>, "<name>" and patch( followed by the module for the seven moved and five lifted names → none (the only patch( targets on the module are _is_visual_gen_instance, untouched).
  • python scripts/legacy_utils.py check-configs → up to date; python scripts/legacy_utils.py lint-precommit tensorrt_llm/serve/api/openai/completions/_merge.py tensorrt_llm/serve/api/openai/completions/serving.py tensorrt_llm/serve/openai_server.py → passes (tip: openai_server.py F821 -1, as expected); git diff 0d6a309791 -- ruff-legacy-baseline.json → only the _merge.py entry; bash scripts/run_mypy.sh → no issues in 24 source files.
  • pre-commit run --files <all 25 changed files> twice → clean on both runs.
  • Left to CI (GPU): tests/integration/defs/test_e2e.py::test_openai_completions_example[pytorch] and the single-GPU disaggregated block of l0_a10.yml (completions through the disaggregated frontend, msgpack to the workers); the multi-GPU stages run once the approval label is applied, since the trigger table changed.

Full suites, A/B against the integration base (same commands on the base merge 0d6a309791 first, then on this head; a failure counts only if its test id is new):

  • pytest -m cpu_only tests/unittest/serve tests/unittest/llmapi/apps tests/unittest/disaggregated -q: base 52 failed / 4164 passed; head 52 failed / 4196 passed, identical failing ids. pytest -m cpu_only tests/unittest/api_stability tests/unittest/serve -q on the head: 8 failed / 312 passed, the 8 being the base's tests/unittest/serve failures below. The 52 are not this PR's: 44 environment failures in tests/unittest/disaggregated (test_context_first_profile.py, test_lifecycle_activation.py, test_peer.py), 6 in tests/unittest/serve/test_serve_health.py (OpenAIServer has no attribute health) and 2 private-access lint tests (stale baseline rows for the _postprocess_result accesses that [TRTLLM-17273][chore] Choose the trtllm-serve post-processing path once per server #21392 removed, plus _PreprocessingWork._release in tensorrt_llm/grpc/openengine/servicer.py), all present on the base.
  • pytest tests/unittest/api_stability -q: base 108 passed; head 108 passed (route snapshot, protocol and CLI references unchanged).
  • pytest "tests/unittest/llmapi/test_llm.py::test_openai_completion_list_prompt_stream_reuses_stream_metadata" tests/unittest/executor/test_spec_decode_stats_payload.py tests/unittest/llmapi/test_kv_hints.py -q (the GPU-stage files, run locally): base 210 passed; head 210 passed.
  • LLM_MODELS_ROOT=... pytest -rs tests/unittest/llmapi/apps/test_harmony_parsing.py tests/unittest/llmapi/apps/test_harmony_channel_validation.py: base 79 passed; head 79 passed, no "not importable" skip.
  • LLM_MODELS_ROOT=... PYTHONPATH=tests python3 -m test_common.serve.cpu_tests (every serve CPU test, one pytest process per test-db entry, 76 entries): the failing ids were the 8 base failures above plus 9 tests in tests/unittest/serve/_base/test_serving.py, test_serve_multi_frontend_route_gate.py, test_serve_multi_frontend_helpers.py and the OpenEngine test_prometheus_contract[http_enablement], which the last commit adapts (re-run: 26 passed and 1 passed, listed above); no other failure.
  • PYTHONPATH=tests python3 -m test_common.serve.routes → exit 0.

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

🤖 Generated with Claude Code

…ce per server

OpenAIServer decided nine times, in four handlers, whether a response's
post-processed value comes from a postproc worker (read the value the
worker stored on outputs[0]) or must be computed in the server process
(call the post-processor inline). The flag never changes after
construction, so the server now chooses a strategy object once:
select_postproc_strategy picks WorkerPostprocStrategy or
InlinePostprocStrategy from num_postprocess_workers; __init__ pre-caches
it and a cached property keeps server instances that unit tests build
without __init__ working. The nine decision points become one strategy
call each, and the conditional post-processor unpacks they fed are gone.

CompletionOutput gains the read-only public accessor postprocess_result
so the server no longer reads the executor's private field; the
api-stability reference records the new property. The submission-side
uses of postproc_worker_enabled (what generate_async receives, filling
postproc_args for in-server post-processing, the Responses store
decision) are deliberately unchanged.

No behavior change. With more than one candidate the two modes still
differ in what the frontend records (worker mode: only outputs[0]
carries the result and metrics arrive merged; inline mode records
metrics per candidate); that known difference is documented on the
strategy class and left for a separate change.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…ss lint for the serving frontends

The HTTP server, the SMG gRPC frontend and the OpenEngine gRPC frontend
learned engine health by reading private members (_check_health, _executor,
_fatal_error, _engine_dead). Add the public health API on the classes that
own those members and switch all three frontends to it, with no behavior
change:

- BaseLLM.check_health() (public rename of _check_health; the private name
  stays for one release as an alias for out-of-tree callers that look it
  up), BaseLLM.fatal_error, BaseLLM.shutdown_started, BaseLLM.engine_dead.
- GenerationExecutor.fatal_error and engine_dead;
  GenerationExecutorProxy.engine_dead also reads the sticky worker-death
  flag set before the fatal error is recorded.

Add a private-access lint (tests/unittest/serve/test_private_access_lint.py)
that counts cross-package private accesses in tensorrt_llm/serve,
tensorrt_llm/commands and tensorrt_llm/grpc against a checked-in baseline
(tests/unittest/serve/private_access_baseline.yaml) that only shrinks, with
companion test files under tests/unittest/serve/api/visual_gen/ and
tests/unittest/grpc/openengine/ so VisualGen-only and OpenEngine-only pull
requests keep running the check.

Replace the three frontend health mirror tests with tests of the production
code, rename the _check_health test doubles to the public name, and register
the new CPU test files in l0_cpu.yml.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…erving-frontend private accesses

The serving frontends (tensorrt_llm/serve, tensorrt_llm/commands, the
OpenEngine gRPC frontend) read private members of LLM, the executors,
CompletionOutput, SamplingParams, MultimodalDataTracker,
OpenAIDisaggServer and PyExecutor. Give each remaining read a thin
public entry point on the class that owns the member, switch the
frontends to it with no behavior change, and shrink the private-access
baseline from 50 recorded accesses to 19. The reads of the two private
LlmArgs attributes stay (exposing them is an LLM-args change for its
own PR); the remaining baseline entries belong to other in-flight work.

New public members (status prototype where api-stability tracks them):
LLM.record_fatal_error / multi_frontend_attach_info /
multi_frontend_ipc_dir / engine_in_worker_process / hf_model_dir /
hf_model_config / resource_governor_queue / encode_limits /
get_kv_cache_load / get_openengine_discovery;
GenerationExecutor.record_fatal_error and multi_frontend_ipc_dir;
CompletionOutput.spec_dec_counters and text_diff_offset (and the
SpecDecCounters rename); SamplingParams.set_request_provided_fields and
stop_word_ids; MultimodalDataTracker.multimodal_server_config;
OpenAIDisaggServer.address; PyExecutor.set_openengine_discovery and
set_openengine_node_server; registered_model_type() in inputs/registry.

Also adds tensorrt_llm/llmapi/engine_client.py: the EngineClient
protocol naming the engine operations every frontend shares (generate,
health, load), checked by tests rather than at runtime.

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

The shut-down-LLM case of the launcher-lost watchdog test now calls the
production method unbound on a stand-in object whose generator double
defines the public members the method reads, the form the file's other
stand-in tests use. New tests must not build servers with object.__new__.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…apshots and a CPU test seam for trtllm-serve

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
… baselines at the merge base

The route reference is generated from the live apps; the object.__new__
baseline and the gRPC import allowance record what exists on main today
and only shrink.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…n the snapshot matrix

An attached frontend of a multi-frontend server registers ten launcher
forwarding routes before the local table; seven share a method and path
with local handlers and three use a path parameter that matches every
batch id, so registration order decides which handler serves them. The
new configuration records all of it, a live invariant pins the
forwarding-before-local order independently of the reference, and the
route-class failure message now says why the msgpack route class matters.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
- AGENTS.md: the CPU test command sets LLM_MODELS_ROOT, which the runner
  requires, and the stage-mapping note says the mapper matches list-entry
  text, so a file listed through its directory maps by the directory.
- check_contracts matches each from-import as <module>.<name>, so
  'from pkg import sub' and 'from .. import sub' count as imports of the
  submodule; new test for both forms and for the allowance entry.
- The VisualGen /v1/models test asserts the 200 status.
- The AGENTS.md command check also recognizes commands that start with
  LLM_MODELS_ROOT=.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…on the current main

The disaggregated frontend now registers its routes with KvHintsRoute
(NVIDIA#20269), which strips kv_hints from request bodies unless the server
trusts its ingress. Regenerate the route reference: the six
disaggregated configurations record KvHintsRoute where they recorded
APIRoute; their route sets and order are unchanged.

The route-class invariant, which a regenerated reference cannot hide,
now requires KvHintsRoute on every API route of the disaggregated
frontend and APIRoute on the coordinator app. Dropping the KV-hint
policy from the frontend fails it even after a regeneration, and both
apps still reject the aggregated server's msgpack route class.

The object.__new__ baseline records the sites on main at the merge base:
test_kv_hints.py builds one server that way (NVIDIA#20269), and
test_harmony_parsing.py no longer builds any (NVIDIA#21176).

The harmony parsing tests are now registered as a whole file (NVIDIA#21411),
so the node-id entry that the CPU-selection test used as its example is
gone. The node-id check now runs select_entries on a one-entry list of
its own, so it no longer depends on which node-id entries l0_cpu.yml
holds.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…ail on silent regressions

The route snapshot now requires every configuration other than the three
invalid ones to build, so a configuration that starts failing at construction
cannot pass through a regenerated reference. The object.__new__ ratchet
resolves 'from ... import <server class> as <name>' aliases and records each
baseline site as its enclosing function or class plus the server class, so
swapping one site for another in the same file no longer passes. The baseline
keeps the same 52 sites in the same 26 files.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…lm/serve/_base/msgspec_route.py

Generated move: _msgpack_decoder, _MsgspecRequest and _MsgspecRoute leave
tensorrt_llm/serve/openai_server.py byte for byte; the module imports
_MsgspecRoute back for its app router and keeps no re-export (the three
names are private). The two disaggregated test importers now import the
new module. OpenAIServer.health, .version and .get_model are removed
here and replaced by the system router added in the following commits.

tensorrt_llm/serve/_base/msgspec_route.py joins legacy-files.txt (its
code comes from a legacy-formatted file), and the generated lint configs
are regenerated with scripts/legacy_utils.py gen-configs.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…ve the system endpoints through it

New structure: tensorrt_llm/serve/_base/ (shared serving infrastructure)
and tensorrt_llm/serve/api/ (endpoint packages, one router plus one
serving class each; the package __init__ imports nothing so routers load
lazily, by role). serve/_base/routing.py provides make_router(), which
always uses the msgpack-decoding route class, and attach_routers(),
which imports and attaches only the routers whose RouterSpec matches the
server role.

The first endpoint package, serve/api/ops/system/, serves /health,
/version and /v1/models for every role. OpenAIServer builds a
SystemServing with the injected _health_probe (today's probe body,
SIGINT side effect included) and attaches the routers before the role
dispatch; the 12 add_api_route lines for the three paths leave the four
role tables. Same routes, methods, route classes, bodies and
conflict-relevant order: the recorded route reference is unchanged.

The layering tests gain the new packages' contracts, the route-snapshot
failure message points new endpoints at the router pattern, and the
serve agent docs describe where a new endpoint goes.

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

tensorrt_llm/serve/_base/ and tensorrt_llm/serve/api/ops/ join the
multi-GPU trigger list: the msgpack route class decodes the
disaggregated frontend's worker requests, and /health is the
disaggregated readiness probe, so changes under these directories keep
running the multi-GPU stages that exercise that path.

The three new modules join the strict mypy file set, and the type-check
pre-commit hook now also triggers on changes under tensorrt_llm/serve/
and tensorrt_llm/commands/serve_cli/, so edits to typed serving modules
get local type checking; the hook still checks only the declared file
set.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
… tests that build routes by hand

Two multi-frontend tests build an OpenAIServer without its constructor
and call register_routes() themselves. Since /health, /version and
/v1/models are attached through the system router by the constructor,
those tests no longer saw the three routes: one asserted that /version
stays local on an attached frontend (404), the other reported the three
per-frontend-safe entries as removed. Both helpers now set the serving
object and attach the routers before the hand registration, in the
constructor's order, and the route enumeration walks included routers,
which FastAPI 0.139 keeps as objects without a path in app.routes.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…n the OpenEngine metrics test

The HTTP-enablement case builds a real OpenAIServer through its
constructor and looked for /prometheus/metrics by reading .path off
every entry of app.routes. Since the constructor attaches the system
router, FastAPI 0.139 keeps an included-router object there that has no
.path, and the child process failed with AttributeError. The test now
matches a GET /prometheus/metrics request against the routes with the
public route-matching API, which works on both ends of the pinned
FastAPI range.

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

The check that endpoint packages under tensorrt_llm/serve/api do not
import each other read the package off the from-module alone, so
`from tensorrt_llm.serve import api` inside an endpoint package passed
and `from tensorrt_llm.serve.api.ops import system` inside ops/system
failed. It now reads the package off <module>.<name> per imported name,
as the layering contracts do, and a synthetic-edge test pins both forms.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…e/api/visual_gen

Move the eight VisualGen request and response models (ImageGenerationRequest,
ImageEditRequest, ImageObject, ImageGenerationResponse, MediaReferenceItem,
VideoGenerationRequest, VideoJob, VideoJobList) out of
tensorrt_llm/serve/openai_protocol.py into the new package
tensorrt_llm/serve/api/visual_gen/ (protocol.py), unchanged. openai_protocol
keeps the eight names importable through a lazy module __getattr__: importing
it no longer loads the VisualGen models or tensorrt_llm.visual_gen, and the
re-export is part of its own interface, so it never warns and has no removal
release. The imports only the moved block used (UploadFile, AliasChoices,
MediaRole) leave openai_protocol with it.

Every importer in the repository now imports the new module; the benchmarking
guide links the new file. protocol.py joins legacy-files.txt (its code comes
from a legacy-formatted file) and the derived lint configs are regenerated.

No behavior change. OpenAPI component names may change with the move; they are
not a contract.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…isualGen team in CODEOWNERS

VisualGen serving code now lives under tensorrt_llm/serve/api/visual_gen/. The
new rule, placed after the /tensorrt_llm/serve rule so that it wins, makes
@NVIDIA/trt-llm-torch-visual-gen-devs the owner of that directory: changes there
need the VisualGen team's approval and no longer need runtime-devs.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…y re-export under a visual_gen path

Add tests/unittest/api_stability/visual_gen/test_visual_gen_serve_api.py: it
runs the checks of test_serve_api.py on the models of
references/trtllm_serve_api.yaml that are defined in
tensorrt_llm.serve.api.visual_gen.protocol, and checks that
tensorrt_llm.serve.openai_protocol re-exports the eight VisualGen names lazily
(importing it does not load the VisualGen module; unknown names still raise
AttributeError). The file lives under a visual_gen/ path and has its own
l0_cpu.yml entry so that CI selection for VisualGen-only pull requests keeps
running it; the reference YAML is unchanged.

Add the eight identity rows of the VisualGen names to the VisualGen shim
identity table: identity only, no warning and no ban on the old path, because
the re-export is openai_protocol's permanent interface.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…its no-warning contract

Resolve every pinned model with the full gate's own loader, so a VisualGen
model pinned in the reference YAML that is later renamed or removed fails this
gate too instead of being dropped silently; the four known models stay a
secondary check. Make the lazy-load probe's diagnostic best-effort and
identity-based, so a failure prints which modules hold the VisualGen classes
instead of crashing on module-level __getattr__ hooks while walking
sys.modules. In the same fresh-interpreter probe, assert that reading each of
the eight names from openai_protocol emits no FutureWarning or
DeprecationWarning.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
… one submodule per API

tensorrt_llm/serve/openai_protocol.py held every OpenAI-dialect request and
response model plus the helpers that turn them into engine parameters. It
becomes the package tensorrt_llm/serve/openai_protocol/ with one submodule per
API: common, sampling, tools, chat, completions, embeddings, responses,
tokenize, models, kv_cache, control and disagg. Every definition moves
unchanged (93 definitions; generated by an AST tool and checked definition by
definition against the former module).

The package __init__.py is a lazy facade: a name-to-module table and a module
__getattr__ import a submodule only when one of its names is first read, so
every import from tensorrt_llm.serve.openai_protocol keeps working for every
name the former module had, the underscore helpers that tests import included,
and the VisualGen models re-exported from
tensorrt_llm.serve.api.visual_gen.protocol. The facade is the package's own
interface, not a compatibility shim: no warning, no removal release, and no
importer in the repository changes. Its TYPE_CHECKING block imports the public
names from their submodules so type checkers resolve them; the VisualGen names
are left out of it on purpose and appear only as strings in the table, so that
importing the facade never loads the VisualGen package. DisaggScheduleStyle,
which the former module imported and two importers took from it, stays
reachable through the facade.

Submodules import each other directly, never through the package, along an
acyclic graph: common imports no sibling; sampling, tools, models, tokenize,
embeddings and control import common; completions and responses import common
and sampling; chat imports common, sampling and tools; kv_cache imports
common, tools and chat; disagg imports common, completions and chat. torch is
imported only by sampling and disagg, openai_harmony only by chat, and the
serving-extension registry only by sampling, so code that needs only the tool
types no longer loads the chat or Responses models.

The 12 submodules are registered in legacy-files.txt (their code comes from a
legacy-formatted file) and the deleted module's line is removed; the derived
lint configs are regenerated with scripts/legacy_utils.py gen-configs. The
facade is new code and stays under the default formatter.

No behavior change. OpenAPI component names may change for classes that share
a short name with an openai type, because pydantic qualifies them with the
module path; they are not a contract.

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

tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py keeps the
package tensorrt_llm/serve/openai_protocol/ honest: every name the former
single-file module defined (plus DisaggScheduleStyle) is reachable through the
facade and is the object its submodule defines; the facade table is frozen at
those names; __all__, dir() and the TYPE_CHECKING block agree with the table;
unknown names raise AttributeError; the submodule import graph stays inside
the allowed acyclic graph and no submodule imports the package itself;
importing the package loads no submodule, importing tools loads neither chat
nor responses nor completions, and reading ChatCompletionRequest does not load
the serving extensions; the private helpers that other tests import through
the package stay importable.

tests/unittest/serve/api/visual_gen/test_openai_protocol_facade_visual_gen.py
holds the VisualGen assertions under a visual_gen/ path so that CI selection
runs them for VisualGen-only changes: the eight VisualGen models map to
tensorrt_llm.serve.api.visual_gen.protocol, nothing but a lookup of one of
them loads that module, and the facade names it only as strings in its table,
never in an import statement.

tests/unittest/serve/test_layering.py gains the contract that the protocol
package imports no handler module (servers, handlers and routers import the
models, never the reverse), and tensorrt_llm/serve/AGENTS.md tells new code to
import the submodule that defines a model rather than the package.

Both test files are collected by the existing unittest/serve and
unittest/serve/api/visual_gen/ entries of l0_cpu.yml and carry the cpu_only
marker.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…s as strings and bind the test's culprit list unconditionally

The table comment in tensorrt_llm/serve/openai_protocol/__init__.py said that a
TYPE_CHECKING import of the VisualGen models would load their package at run
time, which a TYPE_CHECKING block never does. The real reason is the static
import-isolation check that lets CI run only the VisualGen tests for
VisualGen-only changes: it forbids an import edge from this package to the
VisualGen package even under TYPE_CHECKING, so the facade names that module
only as a string.

In tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py the
list of submodules that import the serving extensions was bound only on the
failing path and read in the assertion message; it is now bound before the
branch, so an unconditional message can never raise UnboundLocalError.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…junyix-feat-serve-embeddings-mm-encoder

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

# Conflicts:
#	pyproject.toml
#	tensorrt_llm/grpc/openengine/streaming.py
#	tensorrt_llm/serve/openai_protocol.py
#	tensorrt_llm/serve/openai_server.py
#	tests/integration/test_lists/test-db/l0_cpu.yml
#	tests/unittest/serve/__init__.py
#	tests/unittest/serve/api/__init__.py
#	tests/unittest/serve/api/visual_gen/__init__.py
…gs, multimodal encoder and tokenize

Create tensorrt_llm/serve/_base/serving.py with BaseServing and
GenerateBaseServing, and the EmbeddingServing, MMEncoderServing and
TokenizeServing skeletons under tensorrt_llm/serve/api/openai/. Every
dependency is a keyword-only constructor argument. The new Group B files
are registered in legacy-files.txt and the generated lint configs are
regenerated.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…tokenize handlers into their serving classes

Generated move. OpenAIServer loses openai_embedding,
_normalize_embedding_input, _init_embedding_batcher (to EmbeddingServing),
openai_mm_encoder (to MMEncoderServing) and tokenize (to TokenizeServing);
register_embedding_routes is deleted, its route is served by the
embeddings router. create_error_response, _attach_multimodal_data and
await_disconnected move to the serving bases in
tensorrt_llm/serve/_base/serving.py and stay reachable on OpenAIServer
as aliases of the moved objects while chat and completions still call
them. _record_generator_termination moves to the same module.
tensorrt_llm/serve/encode_batcher.py becomes
tensorrt_llm/serve/api/openai/embeddings/batcher.py; the old module is
deleted and every import of it is rewritten.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…tokenize routers and serving objects

Each new package gets an api_router.py built with make_router() and a
RouterSpec in ROUTERS for exactly the roles the route had. The
composition root builds RenderResources once after _init_llm and
injects it, constructs the serving objects inside the role branches,
puts them on app.state, and attaches the routers after the role switch,
where every attach_router can check its app.state key. The two route
registrations the routers now serve leave register_routes and
register_mm_encoder_routes.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…enize serving classes on CPU

Construction tests (every dependency is a required keyword-only argument),
route tests through the real composition root and routers (msgpack bodies
included), the batcher reference the health probe reads, the injected
RenderResources carrying the server chat template, the serving bases and
the borrowed-helper identity on OpenAIServer. The encode-only fake now
encodes fixed embeddings and scripted outputs can carry disaggregated
params; the app factory passes the chat-template knobs. The runtime-API
tests build EmbeddingServing and patch the termination telemetry where
the moved helper looks it up; the lazy-import guard also checks that
importing tensorrt_llm.serve loads no endpoint package.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
… and the serve guide

The multi-GPU trigger table gains the three package prefixes, so code
moved out of tensorrt_llm/serve/openai_server.py keeps its trigger; the
three router modules join the mypy file list; the serve guide states the
constructor rule of serving classes.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
… stub provide the tokenizer

The test runs the real OpenAIServer constructor with _init_llm replaced
by a stub. The constructor now builds the tokenize serving object for
the LLM role from self.tokenizer, so the stub sets it too; a missing
dependency fails at construction on purpose instead of at request time.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…ition root with the tokenize serving object

Both tests build a partial OpenAIServer and register its routes the way
the constructor does. The constructor now puts the tokenize serving object
on app.state and attaches the role's routers after the local table, so the
re-enactments do the same; the tokenize router refuses to attach without
its serving object.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…unyix-feat-serve-completions-serving (integration base, not part of the reviewable change)

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…nd the generate-pipeline constructor arguments

Create tensorrt_llm/serve/api/openai/completions/ with the CompletionServing
skeleton, and give GenerateBaseServing the keyword-only constructor arguments
that the completions handler and the helpers it shares with chat and Responses
read: the perf-metrics switch and collector, the adjusted steady clock, the
iteration-stats wake-up event and a getter for the collector task that the
lifespan creates later, the internal disaggregated auth key, the per-request
speculative-decoding opt-in and its draft-token bound, the input-processor
thread pool and the post-processing strategy. The new Group B files are
registered in legacy-files.txt and the generated lint configs are regenerated.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…Serving and lift its closures

Generated move: openai_completion leaves OpenAIServer for
CompletionServing in tensorrt_llm/serve/api/openai/completions/serving.py;
its five nested functions become module-level functions of
tensorrt_llm/serve/api/openai/completions/_merge.py whose free variables
(the serving object, the raw request, the disaggregated parameters) are
keyword-only parameters, so they can be unit-tested; the six helpers it
shares with the chat and Responses handlers (postproc_worker_enabled,
_extract_metrics, _logit_bias_vocab_size, _validate_internal_disagg_request,
_has_cache_transceiver_config, _apply_spec_decode_stats_opt_in) move into
GenerateBaseServing and stay reachable on OpenAIServer as aliases of the
moved objects while chat and Responses call them there. The one recorded
F821 (anext) of the lifted generator_wrapper follows it into the lint
baseline under its new path.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…s router and serving object

The LLM role group builds CompletionServing in the composition root, next
to TokenizeServing, with the generate-pipeline arguments read once by
_generate_serving_kwargs(), and puts it on app.state; the new
tensorrt_llm/serve/api/openai/completions/api_router.py maps POST
/v1/completions to it through make_router(), so the route keeps the
msgpack route class the disaggregated frontend relies on, and joins the
routing table for the LLM role group. register_routes no longer registers
the route itself.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…, and build the completions tests on serving objects

New CPU tests: tests/unittest/serve/api/openai/completions/test_merge.py
exercises each lifted function alone (usage merging, generator merging,
the SSE error frame, the first-token stamp, context-only prompt ids, the
error path); test_serving.py covers the constructor contract, the shared
composition-root objects, the route for one and several prompts, msgpack
bodies, the wired internal disaggregated auth key, the stats-collector
wake-up and the role gate. The tests that built an OpenAIServer with
object.__new__ to reach the completions code (the streaming error test,
the list-prompt streaming test, the cache-transceiver auth test, the
spec-decode opt-in test and the KV-hints handler test) now build a
CompletionServing with explicit arguments, and the object.__new__ ratchet
drops their rows. The direct MMEncoderServing constructions pass the new
base arguments; the borrowed-helper table gains the six helpers; the
fixture factory can set the internal disaggregated auth key and a cache
transceiver backend.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…-GPU trigger table and type checking

The multi-GPU trigger table gains the tensorrt_llm/serve/api/openai/completions/
prefix, so the completions code that left tensorrt_llm/serve/openai_server.py
keeps triggering the disaggregated stages; the new router module joins the
mypy file list. The moved serving and lifted-function modules stay out of it,
like the other modules that hold moved legacy code.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…se-class test the new generate-pipeline arguments

The multi-frontend tests re-enact the composition root by hand: they now
put a CompletionServing with placeholder dependencies on app.state before
attaching the role's routers, as the constructor does, because the
completions router refuses to attach without it. The OpenEngine Prometheus
test's constructor stub sets model_config too, since the LLM role now reads
it when it builds the completions serving object. The base-class test
requires every GenerateBaseServing argument, pins the attribute each one is
stored as, and checks that the collector-task property follows the root's
current binding.

Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com>
…r what it is and cover the merged several-prompt response

GenerateBaseServing's constructor argument that carries the getter of the
composition root's iteration-stats collector task is now called
iteration_stats_collector_task_getter at every site: the base class, the
root's _generate_serving_kwargs(), the direct constructions in tests and the
hand-built route tables. The stored attribute and the read-only property
keep their names. The completions route test suite gains the non-streaming
several-prompt case: one merged response with one choice per prompt and the
summed usage, which fails when the merge call loses its serving argument or
the merge drops a prompt.

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

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant