[TRTLLM-8921][feat] implement gen-first disagg_service

[reasonsolo]

Summary by CodeRabbit

Release Notes

• New Features

  • Added /server_info endpoint providing server metadata and context state information
  • Introduced generation-first scheduling style for disaggregated execution alongside context-first via environment configuration
  • New llm_info property exposing model information and encoded context state
  • Router enhancements now collect and coordinate per-server metadata for improved scheduling

• Tests

  • Added comprehensive unit tests for disaggregated service request scheduling across both modes

_{✏️ Tip: You can customize this high-level summary in your review settings.}

Description

Test Coverage

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)

• 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

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces disaggregated request scheduling style support (context-first vs. generation-first) across the system. Key changes include: updating the KV cache transceiver API from batch to single-request handling, adding scheduling style selection in disaggregated executor paths, exposing server context state via new LLM and OpenAI APIs, and refactoring router metadata collection to support per-server state tracking.

Changes

Cohort / File(s)	Summary
KV Cache & Scheduling Infrastructure tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py, tensorrt_llm/llmapi/disagg_utils.py	Updated KV cache transceiver interface from batch-oriented (List[LlmRequest]) to single-request handling (LlmRequest). Introduced new DisaggScheduleStyle enum (CONTEXT_FIRST, GENERATION_FIRST) and get_disagg_schedule_style() function that reads TRTLLM_DISAGG_SCHEDULE_STYLE environment variable.
Executor & Request Preparation tensorrt_llm/_torch/pyexecutor/py_executor.py	Added scheduling style-aware context request preparation via new _check_disagg_ctx_schedulable_status() method. Integrated environment-based scheduling style selection to determine when context-only requests should be prepared ahead of disaggregated scheduling.
Worker & Executor APIs tensorrt_llm/executor/base_worker.py, tensorrt_llm/executor/executor.py	Added get_disagg_context_state() methods to both BaseWorker (delegates to engine.kv_cache_transceiver) and GenerationExecutor (returns empty dict stub).
LLM API tensorrt_llm/llmapi/llm.py	Added llm_info read-only property to BaseLLM that exposes model name and encoded_opaque_state from executor context state.
OpenAI Disaggregated Service tensorrt_llm/serve/openai_disagg_service.py	Significant refactoring to support dual scheduling paths: renamed original _send_disagg_request to _send_disagg_request_ctx_first, added new _send_disagg_request_gen_first for generation-first strategy with concurrent request dispatch. Updated _get_gen_request to accept optional context response and server info, enabling both context-first and generation-only request flows.
OpenAI Server tensorrt_llm/serve/openai_server.py	Added GET /server_info endpoint and get_server_info() method that returns llm.llm_info for retrieving server metadata.
Routing & Metadata Collection tensorrt_llm/serve/router.py	Introduced per-server metadata caching via _server_info dict and _fetch_server_info() coroutine. Refactored abstract method structure: replaced public get_next_server with internal _get_next_server (implemented by router subclasses) and new public wrapper that merges cached server_info into results. All router variants (RoundRobinRouter, LoadBalancingRouter, KvCacheAwareRouter) updated to implement _get_next_server.
Tests tests/integration/test_lists/test-db/l0_a10.yml, tests/unittest/disaggregated/test_openai_disagg_service.py	Added comprehensive unit tests for disaggregated service covering both scheduling styles, streaming/non-streaming modes, context state propagation, and request construction validation.

Sequence Diagram(s)

sequenceDiagram
    participant Client
    participant OpenAIDisaggService as OpenAI Disagg Service
    participant ContextRouter as Context Router
    participant GenerationRouter as Generation Router
    participant Executor as Executor

    rect rgba(100, 200, 100, 0.5)
    Note over Client,Executor: Context-First Scheduling Style
    Client->>OpenAIDisaggService: UCompletionRequest
    OpenAIDisaggService->>ContextRouter: context_only=true
    ContextRouter->>Executor: prepare_context_request(request)
    Executor-->>ContextRouter: UCompletionResponse (context)
    OpenAIDisaggService->>GenerationRouter: build generation request + encoded_opaque_state
    GenerationRouter->>Executor: generate(request)
    Executor-->>GenerationRouter: UCompletionResponse (generation)
    OpenAIDisaggService-->>Client: Final UCompletionResponse
    end

    rect rgba(100, 100, 200, 0.5)
    Note over Client,Executor: Generation-First Scheduling Style
    Client->>OpenAIDisaggService: UCompletionRequest
    par Concurrent
        OpenAIDisaggService->>ContextRouter: prepare_context_request
        ContextRouter->>Executor: context_only=true
        Executor-->>ContextRouter: UCompletionResponse (context)
    and
        OpenAIDisaggService->>GenerationRouter: generation_only=true
        GenerationRouter->>Executor: generate(request)
        Executor-->>GenerationRouter: UCompletionResponse (generation)
    end
    OpenAIDisaggService->>OpenAIDisaggService: merge context state + generation result
    OpenAIDisaggService-->>Client: Final UCompletionResponse
    end

Loading

sequenceDiagram
    participant Server1 as Server 1
    participant Server2 as Server 2
    participant Router
    participant Client

    rect rgba(180, 120, 200, 0.5)
    Note over Server1,Client: Router Initialization & Metadata Collection
    Router->>Server1: GET /server_info
    Server1-->>Router: {model, encoded_opaque_state}
    Router->>Server2: GET /server_info
    Server2-->>Router: {model, encoded_opaque_state}
    Router->>Router: cache _server_info[Server1], _server_info[Server2]
    end

    rect rgba(200, 180, 120, 0.5)
    Note over Server1,Client: Request Routing with Metadata
    Client->>Router: OpenAI request
    Router->>Router: _get_next_server(request)
    Router->>Router: merge server_info from cache
    Router-->>Client: selected server + server_info
    end

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is empty—it only contains the repository template with no actual implementation details, test information, or explanation of the changes made.	Fill in the Description, Test Coverage, and other checklist sections with details about the gen-first disagg service implementation, affected components, and relevant test files (e.g., test_openai_disagg_service.py).
Docstring Coverage	⚠️ Warning	Docstring coverage is 7.69% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title clearly summarizes the main change: implementing a gen-first disaggregation service feature for TRTLLM, matching the substantial changes across multiple files.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 5

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py (1)

99-107: Abstract method signature inconsistent with concrete implementation.

The abstract method prepare_context_request at line 100 declares requests: List[LlmRequest] (plural, list), but the concrete implementation at line 163 accepts request: LlmRequest (singular). This violates the interface contract.

🐛 Proposed fix - update abstract method signature to match implementation

     `@abstractmethod`
-    def prepare_context_request(self, requests: List[LlmRequest]):
+    def prepare_context_request(self, request: LlmRequest):
         """
         Prepare the context request for the cache transceiver in generation-first mode.
-        This method should set the context request state to DISAGG_CONTEXT_WAIT_SCHEDULER
+        This method should set the context request's state to DISAGG_CONTEXT_WAIT_SCHEDULER
         so that it won't be scheduled if the responding generation kvcache request is not
         yet received otherwise set it to CONTEXT_INIT.
         """
         ...

tensorrt_llm/serve/router.py (1)

159-195: Fix _fetch_server_info always returning None and swallowing errors

A return inside finally overrides the JSON response and suppresses exceptions, so _server_info is always None (and downstream routing info breaks).

🛠 Proposed fix

 async def _fetch_server_info(self, server: str, timeout: float) -> dict:
     session = aiohttp.ClientSession()
     try:
         async with session.get(f"{server}/server_info",
                                timeout=timeout) as response:
             return await response.json()
-    finally:
-        await session.close()
-        return None
+    except Exception:
+        return {}
+    finally:
+        await session.close()

🤖 Fix all issues with AI agents

In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 2206-2218: The method _check_disagg_ctx_schedulable_status wrongly
calls kv_cache_transceiver.prepare_context_requests with a list
(ctx_only_requests) while the KvCacheTransceiver interface and
BindKvCacheTransceiver implement prepare_context_request (singular); change the
call to iterate over ctx_only_requests and call
self.kv_cache_transceiver.prepare_context_request(req) for each LlmRequest (or,
if intended, update the interface/implementations to accept a list), and also
either use the new_requests parameter in the filtering logic or remove it to
avoid an unused parameter; update references to BindKvCacheTransceiver and
prepare_context_request accordingly.

In `@tensorrt_llm/executor/base_worker.py`:
- Around line 654-655: The get_disagg_context_state method should guard against
missing engine or kv_cache_transceiver to avoid raising when disagg is disabled
or the worker is shutting down; update get_disagg_context_state to check
self.engine is not None and self.engine.kv_cache_transceiver is not None before
calling get_context_state(), and return a safe fallback (e.g., an empty dict)
when either is missing, preserving the original return type and avoiding side
effects.

In `@tensorrt_llm/executor/executor.py`:
- Around line 360-362: Add overrides of get_disagg_context_state in
GenerationExecutorProxy and RayExecutor so they forward the worker's actual
context state instead of returning {}. Locate the GenerationExecutorProxy and
RayExecutor classes and implement get_disagg_context_state to call through to
the remote worker (the same RPC/path used for other delegated methods) and
return the worker's response; ensure the returned value matches what BaseWorker
provides (i.e., the result of engine.kv_cache_transceiver.get_context_state())
and propagate any errors consistently with existing RPC handling.

In `@tensorrt_llm/llmapi/llm.py`:
- Around line 266-275: The property llm_info is constructing
encoded_opaque_state as a set because of the curly braces; change
encoded_opaque_state to be the direct value returned by
self._executor.get_disagg_context_state() (or None when self._executor is falsy)
instead of wrapping it in braces so the result is JSON-serializable—update the
llm_info property to assign "encoded_opaque_state":
self._executor.get_disagg_context_state() if self._executor else None,
referencing the llm_info property and _executor.get_disagg_context_state() to
locate the code.

In `@tensorrt_llm/serve/router.py`:
- Around line 229-236: The get_next_server wrapper currently mutates and
flattens the cached server info into the returned info dict
(info.update(self._server_info.get(server, {}))) which can cause inconsistent
shapes and a TypeError if the cache holds None; change the behavior to always
return a consistent structure by nesting the cached data under a dedicated key
(e.g., "server_info") instead of updating top-level keys: in get_next_server
(and the analogous places you saw) fetch cached = self._server_info.get(server)
and set info["server_info"] = cached or {} (or cached if truthy) so you never
call update on None and consumers always see a single "server_info" object
rather than flattened fields.

🧹 Nitpick comments (2)

tensorrt_llm/llmapi/disagg_utils.py (1)

15-20: Add new public entities to __all__.

The new DisaggScheduleStyle enum and get_disagg_schedule_style function are used by other modules (e.g., py_executor.py, openai_disagg_service.py) but are not exported in __all__. This can affect documentation generation and explicit import handling.

♻️ Proposed fix

 __all__ = [
     'ServerConfig',
     'parse_disagg_config_file',
     'extract_server_configs',
     'split_world_comm',
+    'DisaggScheduleStyle',
+    'get_disagg_schedule_style',
 ]

tests/unittest/disaggregated/test_openai_disagg_service.py (1)

94-116: Loop variable binding issue in async closures.

The async functions _delayed_ctx_response and _delayed_gen_response capture loop variables (ctx_delay, gen_delay, resp_text, stream_chunks) by reference, not value. While this test likely works because the functions are awaited within the same loop iteration, this pattern is fragile and could cause subtle bugs if the test is refactored.

♻️ Proposed fix - bind loop variables as default arguments

-        async def _delayed_ctx_response(*_args, **_kwargs):
+        async def _delayed_ctx_response(*_args, _ctx_delay=ctx_delay, _resp_text=resp_text, **_kwargs):
             request = _args[0]
-            server, info = await service._ctx_router.get_next_server(request)
-            await asyncio.sleep(ctx_delay)
+            _, _ = await service._ctx_router.get_next_server(request)
+            await asyncio.sleep(_ctx_delay)
             return _make_completion_response(
-                resp_text,
+                _resp_text,
                 finish_reason="length",
                 disagg_request_id=request.disaggregated_params.disagg_request_id,
                 context_only=True,
             )

-        async def _delayed_gen_response(*_args, **_kwargs):
+        async def _delayed_gen_response(*_args, _gen_delay=gen_delay, _resp_text=resp_text, _stream=stream, _stream_chunks=stream_chunks, **_kwargs):
             request = _args[0]
-            server, info = await service._gen_router.get_next_server(request)
-            await asyncio.sleep(gen_delay)
-            if stream:
-                return _mock_streaming_response(stream_chunks)
+            _, _ = await service._gen_router.get_next_server(request)
+            await asyncio.sleep(_gen_delay)
+            if _stream:
+                return _mock_streaming_response(_stream_chunks)
             return _make_completion_response(
-                resp_text,
+                _resp_text,
                 finish_reason="stop",
                 disagg_request_id=request.disaggregated_params.disagg_request_id,
                 context_only=False,
             )

[reasonsolo]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #33817 [ run ] triggered by Bot. Commit: f62a4df

[tensorrt-cicd]

PR_Github #33817 [ run ] completed with state SUCCESS. Commit: f62a4df
/LLM/main/L0_MergeRequest_PR pipeline #26080 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

[reasonsolo]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #33857 [ run ] triggered by Bot. Commit: 018f2c1

[tensorrt-cicd]

PR_Github #33857 [ run ] completed with state SUCCESS. Commit: 018f2c1
/LLM/main/L0_MergeRequest_PR pipeline #26107 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

[tensorrt-cicd]

PR_Github #34287 [ run ] completed with state SUCCESS. Commit: 70eae39
/LLM/main/L0_MergeRequest_PR pipeline #26443 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

[reasonsolo]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #34323 [ run ] triggered by Bot. Commit: db4d83f

[tensorrt-cicd]

PR_Github #34323 [ run ] completed with state SUCCESS. Commit: db4d83f
/LLM/main/L0_MergeRequest_PR pipeline #26473 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

[reasonsolo]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #34355 [ run ] triggered by Bot. Commit: 4eb24e9

[tensorrt-cicd]

PR_Github #34355 [ run ] completed with state FAILURE. Commit: 4eb24e9
/LLM/main/L0_MergeRequest_PR pipeline #26503 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

[reasonsolo]

/bot run --disable-fail-fast

[Superjomn]

LGTM on the llmapi changes.

[tensorrt-cicd]

PR_Github #34549 [ run ] triggered by Bot. Commit: a49d49d

[tensorrt-cicd]

PR_Github #34549 [ run ] completed with state FAILURE. Commit: a49d49d
/LLM/main/L0_MergeRequest_PR pipeline #26660 completed with status: 'ABORTED'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

[pcastonguay]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #34678 [ run ] triggered by Bot. Commit: a49d49d

[tensorrt-cicd]

PR_Github #34678 [ run ] completed with state SUCCESS. Commit: a49d49d
/LLM/main/L0_MergeRequest_PR pipeline #26759 completed with status: 'SUCCESS'