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

components/src/dynamo/trtllm/snapshot.py

@@ -11,6 +11,39 @@
 from dynamo.trtllm.constants import DisaggregationMode, Modality
 
 _EXTERNAL_MODEL_LOAD_FORMATS = {"gms"}
+_WARMUP_INPUT_IDS = (1, 2, 3)
+
+
+def _create_warmup_sampling_params() -> Any:
+    from tensorrt_llm.llmapi import SamplingParams
+
+    return SamplingParams(
+        end_id=-1,
+        pad_id=-1,
+        max_tokens=2,
+        temperature=0.0,
+        ignore_eos=True,
+        detokenize=False,
+    )
+
+
+async def warmup_engine(engine: Any) -> None:
+    """Warm TensorRT-LLM before capture."""
+
+    sampling_params = _create_warmup_sampling_params()
+    logging.info("TensorRT-LLM snapshot warmup starting")
+    generation_result = engine.llm.generate_async(
+        inputs=list(_WARMUP_INPUT_IDS),
+        sampling_params=sampling_params,
+        streaming=True,
+    )
+    async for _ in generation_result:
+        pass
+    if generation_result.error is not None:
+        raise RuntimeError(
+            f"TensorRT-LLM snapshot warmup failed: {generation_result.error}"
+        )
+    logging.info("TensorRT-LLM snapshot warmup complete")
 
 
 def _should_prefetch_model_for_snapshot(config: Any) -> bool:
@@ -109,6 +142,7 @@ async def snapshot_before_endpoint(self, engine: Any, config: Any) -> None:
             "Checkpoint mode enabled: TRT-LLM engine is initialized before "
             "Dynamo runtime creation"
         )
+        await warmup_engine(engine)
         pause_controller = _NoOpSnapshotPauseController()
         snapshot_controller = _create_engine_snapshot_controller(
             engine=engine,

components/src/dynamo/trtllm/tests/test_trtllm_snapshot.py

@@ -1,10 +1,13 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from types import SimpleNamespace
+import sys
+from types import ModuleType, SimpleNamespace
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 
+from dynamo.trtllm import snapshot as snapshot_mod
 from dynamo.trtllm.constants import DisaggregationMode, Modality
 from dynamo.trtllm.snapshot import (
     _should_prefetch_model_for_snapshot,
@@ -15,6 +18,7 @@
 pytestmark = [
     pytest.mark.unit,
     pytest.mark.trtllm,
+    pytest.mark.core,
     pytest.mark.gpu_0,
     pytest.mark.pre_merge,
 ]
@@ -31,6 +35,34 @@ def shutdown(self) -> None:
         self.shutdown_called = True
 
 
+class _GenerationResult:
+    def __init__(self, events=None, *, error=None, iteration_error=None):
+        self._events = events if events is not None else []
+        self.error = error
+        self._iteration_error = iteration_error
+
+    async def __aiter__(self):
+        self._events.append("warmup-chunk")
+        yield self
+        if self._iteration_error is not None:
+            raise self._iteration_error
+        self._events.append("warmup-final")
+        yield self
+
+
+def _engine(result):
+    llm = SimpleNamespace(generate_async=Mock(return_value=result))
+    return SimpleNamespace(llm=llm), llm
+
+
+@pytest.fixture
+def warmup_setup(monkeypatch):
+    controller = Mock()
+    monkeypatch.setattr(snapshot_mod, "_create_warmup_sampling_params", object)
+    monkeypatch.setattr(snapshot_mod, "_create_engine_snapshot_controller", controller)
+    return controller
+
+
 def _snapshot_config(**overrides):
     values = {
         "modality": Modality.TEXT,
@@ -93,6 +125,27 @@ def test_snapshot_prefetch_skips_external_model_loader():
     )
 
 
+def test_create_warmup_sampling_params_uses_lazy_trtllm_import(monkeypatch):
+    constructor = Mock()
+    trtllm_module = ModuleType("tensorrt_llm")
+    trtllm_module.__path__ = []  # type: ignore[attr-defined]
+    llmapi_module = ModuleType("tensorrt_llm.llmapi")
+    llmapi_module.SamplingParams = constructor  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "tensorrt_llm", trtllm_module)
+    monkeypatch.setitem(sys.modules, "tensorrt_llm.llmapi", llmapi_module)
+
+    snapshot_mod._create_warmup_sampling_params()
+
+    constructor.assert_called_once_with(
+        end_id=-1,
+        pad_id=-1,
+        max_tokens=2,
+        temperature=0.0,
+        ignore_eos=True,
+        detokenize=False,
+    )
+
+
 @pytest.mark.parametrize(
     ("override", "expected"),
     [
@@ -117,17 +170,18 @@ def test_snapshot_config_rejects_paths_that_can_create_pre_restore_state(
 
 
 @pytest.mark.asyncio
-async def test_snapshot_runtime_proxy_materializes_runtime_after_restore(monkeypatch):
-    import dynamo.trtllm.snapshot as snapshot_mod
-
+async def test_snapshot_runtime_proxy_materializes_runtime_after_restore(
+    monkeypatch, warmup_setup
+):
     created_runtime = _Runtime()
     lifecycle_calls = []
+    result = _GenerationResult(lifecycle_calls)
+    engine, llm = _engine(result)
 
     class FakeSnapshotController:
         def __init__(self, engine, pause_controller, snapshot_config):
             self.engine = engine
             self.pause_controller = pause_controller
-            self.snapshot_config = snapshot_config
 
         async def wait_for_restore(self):
             lifecycle_calls.append("pause")
@@ -171,9 +225,17 @@ async def fake_refresh_restore_runtime_config(config, argv):
     with pytest.raises(RuntimeError, match="not available until"):
         proxy.endpoint("ns.component.generate")
 
-    await proxy.snapshot_before_endpoint(engine=object(), config=config)
-
-    assert lifecycle_calls == ["pause", "resume"]
+    await proxy.snapshot_before_endpoint(engine=engine, config=config)
+
+    assert lifecycle_calls == [
+        "warmup-chunk",
+        "warmup-final",
+        "pause",
+        "resume",
+    ]
+    call = llm.generate_async.call_args
+    assert call.kwargs["inputs"] == [1, 2, 3]
+    assert call.kwargs["streaming"] is True
     assert config.namespace == "restored-ns"
     assert config.discovery_backend == "kubernetes"
     assert proxy.endpoint("ns.component.generate") == "endpoint:ns.component.generate"
@@ -184,8 +246,6 @@ async def fake_refresh_restore_runtime_config(config, argv):
 
 @pytest.mark.asyncio
 async def test_snapshot_runtime_proxy_exits_without_runtime_after_capture(monkeypatch):
-    import dynamo.trtllm.snapshot as snapshot_mod
-
     class SnapshotCaptured(Exception):
         pass
 
@@ -203,6 +263,7 @@ def fake_exit(code):
         assert code == 0
         raise SnapshotCaptured
 
+    monkeypatch.setattr(snapshot_mod, "warmup_engine", AsyncMock())
     monkeypatch.setattr(
         snapshot_mod,
         "_create_engine_snapshot_controller",
@@ -215,3 +276,31 @@ def fake_exit(code):
 
     with pytest.raises(SnapshotCaptured):
         await proxy.snapshot_before_endpoint(engine=object(), config=_runtime_config())
+
+
+@pytest.mark.asyncio
+async def test_snapshot_warmup_generation_error_prevents_readiness(
+    warmup_setup,
+):
+    result = _GenerationResult(iteration_error=RuntimeError("generation failed"))
+    engine, _ = _engine(result)
+
+    proxy = _SnapshotRuntimeProxy(snapshot_config=object())
+    with pytest.raises(RuntimeError, match="generation failed"):
+        await proxy.snapshot_before_endpoint(engine=engine, config=_runtime_config())
+
+    warmup_setup.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_snapshot_warmup_terminal_error_prevents_readiness(
+    warmup_setup,
+):
+    result = _GenerationResult(error="executor failed")
+    engine, _ = _engine(result)
+
+    proxy = _SnapshotRuntimeProxy(snapshot_config=object())
+    with pytest.raises(RuntimeError, match="executor failed"):
+        await proxy.snapshot_before_endpoint(engine=engine, config=_runtime_config())
+
+    warmup_setup.assert_not_called()

components/src/dynamo/vllm/snapshot.py

@@ -1,22 +1,60 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import asyncio
 import gc
 import logging
+import uuid
 from collections.abc import Callable
 
+from vllm.inputs import TokensPrompt
+from vllm.sampling_params import SamplingParams
+
 from dynamo.common.snapshot.lifecycle import (
     EngineSnapshotController,
     SnapshotConfig,
     configure_snapshot_capture_env,
 )
 
 from .args import Config
-from .handlers import VllmEnginePauseController
+from .constants import DisaggregationMode
+from .handlers import VllmEnginePauseController, get_dp_range_for_worker
 from .worker_factory import EngineSetupResult
 
 logger = logging.getLogger(__name__)
 
+_WARMUP_INPUT_IDS = (1, 2, 3)
+
+
+async def warmup_engine(engine_setup: EngineSetupResult) -> None:
+    """Warm the direct vLLM generation path before snapshot capture."""
+    engine, vllm_config, *_ = engine_setup
+    runner_type = vllm_config.model_config.runner_type
+    if runner_type != "generate":
+        logger.info("Skipping vLLM snapshot warmup for non-generation model")
+        return
+
+    sampling_params = SamplingParams(
+        max_tokens=2,
+        temperature=0.0,
+        ignore_eos=True,
+        detokenize=False,
+    )
+    _, managed_dp_size = get_dp_range_for_worker(vllm_config)
+
+    async def consume_generation(local_dp_rank: int) -> None:
+        async for _ in engine.generate(
+            TokensPrompt(prompt_token_ids=list(_WARMUP_INPUT_IDS)),
+            sampling_params,
+            str(uuid.uuid4()),
+            data_parallel_rank=local_dp_rank,
+        ):
+            pass
+
+    logger.info("vLLM snapshot warmup starting")
+    await asyncio.gather(*(consume_generation(rank) for rank in range(managed_dp_size)))
+    logger.info("vLLM snapshot warmup complete")
+
 
 async def prepare_snapshot_engine(
     config: Config,
@@ -38,6 +76,12 @@ async def prepare_snapshot_engine(
     config.engine_args.enable_sleep_mode = True
 
     engine = setup_vllm_engine(config)
+    # Embedding and encode workers do not serve generation through this engine.
+    if (
+        not config.embedding_worker
+        and config.disaggregation_mode != DisaggregationMode.ENCODE
+    ):
+        await warmup_engine(engine)
     gc.collect()
     snapshot_controller = EngineSnapshotController(
         engine=engine,