diff --git a/components/src/dynamo/trtllm/snapshot.py b/components/src/dynamo/trtllm/snapshot.py index 05bd9b98b74a..b758393a9f16 100644 --- a/components/src/dynamo/trtllm/snapshot.py +++ b/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, diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_snapshot.py b/components/src/dynamo/trtllm/tests/test_trtllm_snapshot.py index 17021b9ebe6f..2796a356a347 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_snapshot.py +++ b/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() diff --git a/components/src/dynamo/vllm/snapshot.py b/components/src/dynamo/vllm/snapshot.py index 84bd2d6e9c1e..66a458c5d290 100644 --- a/components/src/dynamo/vllm/snapshot.py +++ b/components/src/dynamo/vllm/snapshot.py @@ -1,10 +1,15 @@ # 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, @@ -12,11 +17,44 @@ ) 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, diff --git a/components/src/dynamo/vllm/tests/test_vllm_snapshot.py b/components/src/dynamo/vllm/tests/test_vllm_snapshot.py new file mode 100644 index 000000000000..05b6c36875bd --- /dev/null +++ b/components/src/dynamo/vllm/tests/test_vllm_snapshot.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import logging +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +from dynamo.vllm import snapshot as snapshot_mod +from dynamo.vllm.constants import DisaggregationMode +from dynamo.vllm.snapshot import prepare_snapshot_engine, warmup_engine + +pytestmark = [ + pytest.mark.unit, + pytest.mark.vllm, + pytest.mark.core, + pytest.mark.gpu_0, + pytest.mark.pre_merge, +] + + +def _engine_setup(engine, runner_type="generate"): + config = SimpleNamespace(model_config=SimpleNamespace(runner_type=runner_type)) + return engine, config + + +def _config(**overrides): + values = { + "headless": False, + "embedding_worker": False, + "disaggregation_mode": DisaggregationMode.AGGREGATED, + "engine_args": SimpleNamespace(enable_sleep_mode=False), + } + values.update(overrides) + return SimpleNamespace(**values) + + +async def _prepare(engine): + return await prepare_snapshot_engine(_config(), lambda _: _engine_setup(engine)) + + +@pytest.fixture +def snapshot_enabled(monkeypatch): + controller = Mock() + monkeypatch.setattr( + snapshot_mod.SnapshotConfig, "from_env", Mock(return_value=object()) + ) + monkeypatch.setattr(snapshot_mod, "configure_snapshot_capture_env", lambda: None) + monkeypatch.setattr(snapshot_mod, "EngineSnapshotController", controller) + monkeypatch.setattr( + snapshot_mod, "get_dp_range_for_worker", Mock(return_value=(4, 2)) + ) + return controller + + +@pytest.mark.asyncio +async def test_prepare_snapshot_consumes_all_dp_warmups_before_readiness( + snapshot_enabled, +): + events = [] + + async def generate(*args, data_parallel_rank): + events.append(("first", data_parallel_rank)) + yield object() + await asyncio.sleep(0) + events.append(("final", data_parallel_rank)) + yield object() + + engine = SimpleNamespace(generate=Mock(side_effect=generate)) + snapshot_enabled.return_value.wait_for_restore = AsyncMock( + side_effect=lambda: events.append("ready") or True + ) + + await _prepare(engine) + + assert events == [ + ("first", 0), + ("first", 1), + ("final", 0), + ("final", 1), + "ready", + ] + calls = engine.generate.call_args_list + assert {call.kwargs["data_parallel_rank"] for call in calls} == {0, 1} + identifiers = [call.args[2] for call in calls] + assert len(set(identifiers)) == 2 + for call in calls: + prompt, sampling_params, _ = call.args + assert prompt["prompt_token_ids"] == [1, 2, 3] + assert ( + sampling_params.max_tokens, + sampling_params.temperature, + sampling_params.ignore_eos, + sampling_params.detokenize, + ) == (2, 0.0, True, False) + + +@pytest.mark.asyncio +async def test_warmup_generation_error_propagates_and_prevents_readiness( + snapshot_enabled, +): + async def generate(*args, data_parallel_rank): + if data_parallel_rank == 0: + raise RuntimeError("generation failed") + yield object() + + engine = SimpleNamespace(generate=Mock(side_effect=generate)) + + with pytest.raises(RuntimeError, match="generation failed"): + await _prepare(engine) + + snapshot_enabled.assert_not_called() + + +@pytest.mark.parametrize( + "override", + [ + {"embedding_worker": True}, + {"disaggregation_mode": DisaggregationMode.ENCODE}, + ], +) +@pytest.mark.asyncio +async def test_prepare_snapshot_skips_warmup_for_unwarmable_worker_shapes( + snapshot_enabled, override +): + engine = SimpleNamespace(generate=Mock()) + engine_setup = _engine_setup(engine, runner_type="generate") + setup_engine = Mock(return_value=engine_setup) + snapshot_enabled.return_value.wait_for_restore = AsyncMock(return_value=True) + + result = await prepare_snapshot_engine(_config(**override), setup_engine) + + setup_engine.assert_called_once() + snapshot_enabled.assert_called_once() + engine.generate.assert_not_called() + assert result is snapshot_enabled.return_value + + +@pytest.mark.asyncio +async def test_warmup_skips_non_generation_engine(caplog): + engine = SimpleNamespace(generate=Mock()) + + with caplog.at_level(logging.INFO, logger=snapshot_mod.__name__): + await warmup_engine(_engine_setup(engine, runner_type="pooling")) + + engine.generate.assert_not_called() + assert "Skipping vLLM snapshot warmup for non-generation model" in caplog.text