fix(snapshot): warm vLLM and TRT-LLM before capture

[galletas1712]

Summary

• Run a small deterministic direct-engine generation through vLLM and TensorRT-LLM before Dynamo Snapshot emits ready-for-snapshot.
• Warm every vLLM data-parallel rank managed by the local AsyncLLM, consuming the rank streams concurrently before capture.
• Skip generation warmup for vLLM embedding, disaggregated encode, and other non-generation workers without preventing the existing snapshot lifecycle.
• Keep the implementation on each backend's supported generation API and avoid a cross-engine warmup abstraction.

This follows the same basic invariant as SGLang's startup warmup: exercise the normal prefill, decode, and sampling path before declaring the engine ready for capture.

Details

• vLLM submits one token-only request per locally managed DP rank with explicit local rank selection, a required unique request ID, deterministic sampling, and two output tokens.
• TensorRT-LLM submits one token-only request with deterministic sampling and two output tokens, fully consumes the asynchronous result, and reports a terminal backend error.
• Two output tokens cover prefill and a subsequent decode step.
• Warmup uses no cache salt, backend-local timeout, or explicit abort machinery. Outer startup/snapshot orchestration owns timeout behavior.
• Workers that cannot use generation warmup are skipped; inability to warm does not make an otherwise existing snapshot path ineligible.
• The optional TensorRT-LLM dependency remains lazily imported through its public tensorrt_llm.llmapi surface.

Where should the reviewer start?

• components/src/dynamo/vllm/snapshot.py
• components/src/dynamo/trtllm/snapshot.py

Validation

• git diff --check — passed.
• Ruff format/check — passed on all four changed files.
• Black format check — passed on all four changed files.
• Python compile checks — passed on all four changed files.
• Source review confirmed the TensorRT-LLM v1.3.0rc20 public SamplingParams export and the pinned vLLM direct-generation/rank-routing API contracts.

No separate local runtime test stage was run for this revision. CI is responsible for the backend test environments. A real GPU checkpoint/capture/restore smoke test was not available locally.

Related Issues

🚫 This PR is NOT linked to an issue:

• Confirmed — no related issue

[devin-ai-integration]

Devin Review found 1 potential issue.

[coderabbitai]

Walkthrough

Snapshot lifecycle configuration now exposes a validated warmup timeout. TensorRT-LLM and vLLM snapshot preparation perform asynchronous generation warmups before snapshot readiness, with timeout, failure, abort, configuration validation, and ordering tests.

Changes

Snapshot warmup configuration

Layer / File(s)	Summary
Warmup timeout contract and validation components/src/dynamo/common/snapshot/constants.py, components/src/dynamo/common/snapshot/lifecycle.py, components/src/dynamo/common/tests/test_snapshot_lifecycle.py	Adds the warmup timeout environment key, default value, parser, finite-positive validation, and unit coverage for unset, valid, and invalid values.

TensorRT-LLM warmup

Layer / File(s)	Summary
TensorRT-LLM generation warmup and readiness gating components/src/dynamo/trtllm/snapshot.py, components/src/dynamo/trtllm/tests/test_trtllm_snapshot.py	Adds deterministic streaming warmup, result validation, timeout abort handling, snapshot lifecycle ordering, and tests for success, timeout, invalid configuration, terminal errors, and generation failures.

vLLM warmup

Layer / File(s)	Summary
vLLM generation warmup and configuration checks components/src/dynamo/vllm/snapshot.py, components/src/dynamo/vllm/tests/test_vllm_snapshot.py	Adds generation warmup before snapshot readiness, timeout and admission abort handling, failure propagation, runner/configuration validation, and asynchronous ordering tests.

Estimated code review effort: 4 (Complex) | ~45 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 3.70% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main change: warming vLLM and TensorRT-LLM before snapshot capture.
Description check	✅ Passed	The description covers overview, implementation details, reviewer start points, validation, and related issues.

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@components/src/dynamo/vllm/snapshot.py`:
- Around line 66-76: Bound warmup abort cleanup with the existing timeout in
vLLM snapshot handling, and avoid swallowing abort failures by handling the
cleanup timeout explicitly while preserving the original generation timeout.
Apply the equivalent cleanup-timeout change to the TRT-LLM snapshot path. Extend
the warmup tests to cover an abort failure, including the expected logging and
timeout behavior at the affected test range.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: acd1967f-350a-4f48-93d2-9fc22c2968a7

📥 Commits

Reviewing files that changed from the base of the PR and between 1590dc6 and e931994.

📒 Files selected for processing (7)

• components/src/dynamo/common/snapshot/constants.py
• components/src/dynamo/common/snapshot/lifecycle.py
• components/src/dynamo/common/tests/test_snapshot_lifecycle.py
• components/src/dynamo/trtllm/snapshot.py
• components/src/dynamo/trtllm/tests/test_trtllm_snapshot.py
• components/src/dynamo/vllm/snapshot.py
• components/src/dynamo/vllm/tests/test_vllm_snapshot.py

[datadog-official]

⚠️ Warnings

🚦 2 Pipeline jobs failed

PR | deploy-operator     

PR | deploy-status-check     

ℹ️ Info

🎯 Code Coverage (details)
• Patch Coverage: 100.00%
• Overall Coverage: 38.06% (-3.65%)

Useful? React with 👍 / 👎

_{This comment will be updated automatically if new data arrives.
🔗 Commit SHA: 08790bb | Docs | Give us feedback!}

[galletas1712]

/ok to test ed6fc53

[galletas1712]

/ok to test 76baa4f

[mohammedabdulwahhab]

Question: Why don't the existing warmup paths on the engine cover this? It would be useful to have a comment on why this is missed on those paths.

[mohammedabdulwahhab]

how were these values chosen? can the selection of these determine which kernels were compiled?