diff --git a/components/src/dynamo/vllm/instrumented_scheduler.py b/components/src/dynamo/vllm/instrumented_scheduler.py index b6c44213a04e..cab1ac2ca55f 100644 --- a/components/src/dynamo/vllm/instrumented_scheduler.py +++ b/components/src/dynamo/vllm/instrumented_scheduler.py @@ -106,6 +106,7 @@ from vllm.v1.core.sched.async_scheduler import AsyncScheduler from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput from vllm.v1.core.single_type_kv_cache_manager import CrossAttentionManager +from vllm.v1.engine.core import EngineCore from vllm.v1.kv_cache_interface import MambaSpec from vllm.v1.request import Request, RequestStatus @@ -6812,3 +6813,25 @@ def _bench_write_results(self) -> None: dest, len(self._bench_results), ) + + +# TODO(upstream-vllm): remove once vLLM exposes a way to update scheduler +# identity after engine construction. In snapshot mode the engine is built +# before the Dynamo runtime exists, so the FPM worker_id is baked as "". +def _install_fpm_worker_id_utility() -> None: + if hasattr(EngineCore, "set_fpm_worker_id"): + return + + def set_fpm_worker_id(self, new_worker_id: str) -> None: + scheduler = self.scheduler + if not isinstance(scheduler, InstrumentedScheduler): + raise RuntimeError( + f"scheduler is {type(scheduler).__name__}, not InstrumentedScheduler" + ) + scheduler._fpm_worker_id = new_worker_id + scheduler._publisher._worker_id = new_worker_id + + EngineCore.set_fpm_worker_id = set_fpm_worker_id + + +_install_fpm_worker_id_utility() diff --git a/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py b/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py index e56e2825585d..9b2ef140ef23 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py +++ b/components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py @@ -14,6 +14,9 @@ import hashlib import json +import subprocess +import sys +import textwrap import threading import time import uuid @@ -6855,3 +6858,175 @@ def test_random_kda_allows_hybrid_warm_chains_without_expert_parallelism(monkeyp stub._bench_random_kda = True assert stub._kvwarm_warm_eligible() assert stub._kvwarm_meta["skip_reason"] is None + + +# --------------------------------------------------------------------------- +# FPM worker_id propagation into the EngineCore child (snapshot restore) +# --------------------------------------------------------------------------- + +FPM_UTILITY_NAME = "set_fpm_worker_id" + + +def _fpm_utility(): + from vllm.v1.engine.core import EngineCore + + return getattr(EngineCore, FPM_UTILITY_NAME) + + +def _fpm_scheduler_stub(worker_id: str = ""): + """``InstrumentedScheduler`` carrying only the two FPM identity fields.""" + scheduler = object.__new__(InstrumentedScheduler) + scheduler._fpm_worker_id = worker_id + scheduler._publisher = SimpleNamespace(_worker_id=worker_id) + return scheduler + + +def test_fpm_utility_installed_on_engine_core_base_class(): + """Patched on the base class, so every EngineCore variant inherits it.""" + from vllm.v1.engine.core import EngineCore, EngineCoreProc + + assert FPM_UTILITY_NAME in vars(EngineCore) + assert hasattr(EngineCoreProc, FPM_UTILITY_NAME) + + +def test_fpm_utility_install_is_idempotent(): + """A second install must not rebind an already-patched class.""" + before = _fpm_utility() + + instrumented_scheduler_module._install_fpm_worker_id_utility() + + assert _fpm_utility() is before + + +def test_fpm_utility_updates_scheduler_and_publisher(): + """Active samples use the scheduler's id; idle heartbeats use the publisher's.""" + scheduler = _fpm_scheduler_stub() + engine_core = SimpleNamespace(scheduler=scheduler) + + _fpm_utility()(engine_core, "8465209922961459") + + assert scheduler._fpm_worker_id == "8465209922961459" + assert scheduler._publisher._worker_id == "8465209922961459" + + +def test_fpm_utility_overwrites_a_previously_set_id(): + """A pod may be restored more than once; the id must follow the new runtime.""" + scheduler = _fpm_scheduler_stub(worker_id="1111111111111111") + engine_core = SimpleNamespace(scheduler=scheduler) + + _fpm_utility()(engine_core, "2222222222222222") + + assert scheduler._fpm_worker_id == "2222222222222222" + assert scheduler._publisher._worker_id == "2222222222222222" + + +@pytest.mark.parametrize( + "scheduler", [None, SimpleNamespace()], ids=["missing", "foreign"] +) +def test_fpm_utility_rejects_non_instrumented_scheduler(scheduler): + """Raise rather than no-op, so the parent sees the failure.""" + engine_core = SimpleNamespace(scheduler=scheduler) + + with pytest.raises(RuntimeError, match="not InstrumentedScheduler"): + _fpm_utility()(engine_core, "8465209922961459") + + +def test_fpm_utility_argument_is_not_msgspec_converted(): + """vLLM converts msgspec.Struct-annotated args; the id must stay a plain str.""" + from inspect import isclass, signature + + import msgspec + + annotation = signature(_fpm_utility()).parameters["new_worker_id"].annotation + + assert not (isclass(annotation) and issubclass(annotation, msgspec.Struct)) + + +@pytest.mark.slow +@pytest.mark.timeout(300) +def test_scheduler_cls_resolution_installs_the_patch(): + """Resolving ``--scheduler-cls`` is what installs the patch in the child. + + Runs in a fresh interpreter: this module already imported the scheduler. + """ + script = textwrap.dedent( + """ + from vllm.utils.import_utils import resolve_obj_by_qualname + from vllm.v1.engine.core import EngineCore + + assert not hasattr(EngineCore, "set_fpm_worker_id"), ( + "patch present before scheduler_cls resolution" + ) + + resolved = resolve_obj_by_qualname( + "dynamo.vllm.instrumented_scheduler.InstrumentedScheduler" + ) + + assert resolved.__name__ == "InstrumentedScheduler" + assert hasattr(EngineCore, "set_fpm_worker_id"), ( + "resolving scheduler_cls did not install the FPM worker_id utility" + ) + """ + ) + + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=280, + ) + + assert result.returncode == 0, result.stderr + + +def test_fpm_utility_via_vllm_dispatch_retargets_active_and_heartbeat_ids(): + """Through vLLM's own utility dispatch, both FPM payload kinds carry the new id.""" + import queue + + import zmq + from vllm.v1.engine import EngineCoreRequestType + from vllm.v1.engine.core import EngineCoreProc, EngineShutdownState + + from dynamo.common.forward_pass_metrics import decode + + ctx = zmq.Context.instance() + sub = ctx.socket(zmq.SUB) + sub.setsockopt(zmq.SUBSCRIBE, b"") + port = sub.bind_to_random_port("tcp://127.0.0.1") + sub.unbind(sub.getsockopt(zmq.LAST_ENDPOINT)) + publisher = instrumented_scheduler_module._FpmPublisherThread( + f"tcp://127.0.0.1:{port}", worker_id="", dp_rank=0 + ) + sub.connect(f"tcp://127.0.0.1:{port}") + try: + scheduler = object.__new__(InstrumentedScheduler) + scheduler._fpm_worker_id = "" + scheduler._fpm_dp_rank = 0 + scheduler._publisher = publisher + engine = object.__new__(EngineCoreProc) + engine.scheduler = scheduler + engine.shutdown_state = EngineShutdownState.RUNNING + engine.output_queue = queue.Queue() + + engine._handle_client_request( + EngineCoreRequestType.UTILITY, + (0, 7, FPM_UTILITY_NAME, ("8465209922961459",)), + ) + + _client, outputs = engine.output_queue.get_nowait() + assert outputs.utility_output.failure_message is None + active = InstrumentedScheduler._extract_metrics( + scheduler, + None, + None, + 0.0, + scheduled=instrumented_scheduler_module.ScheduledRequestMetrics(), + ) + assert active.worker_id == "8465209922961459" + assert sub.poll(timeout=5000), "no idle heartbeat within 5s" + heartbeat = decode(sub.recv_multipart()[2]) + assert heartbeat is not None + assert heartbeat.worker_id == "8465209922961459" + finally: + publisher.shutdown() + sub.close(linger=0) diff --git a/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py b/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py index 8d7fac4f9cb2..9c6b8e2b2655 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py +++ b/components/src/dynamo/vllm/tests/test_vllm_worker_factory.py @@ -10,10 +10,13 @@ from unittest.mock import AsyncMock, Mock, patch import pytest +from vllm.v1.core.sched.async_scheduler import AsyncScheduler from dynamo.llm import ModelInput, ModelType, WorkerType from dynamo.vllm.constants import DisaggregationMode +from dynamo.vllm.instrumented_scheduler import ENV_FPM_WORKER_ID, InstrumentedScheduler from dynamo.vllm.worker_factory import ( + FPM_SET_WORKER_ID_METHOD_NAME, EngineSetupResult, SnapshotEngineSetupResult, WorkerFactory, @@ -22,6 +25,8 @@ _merge_benchmark_rank_results, _register_request_cache_metrics, _stop_worker_gc_policy, + _sync_fpm_worker_id, + _sync_fpm_worker_id_or_shutdown, _wait_and_load_benchmark, ) @@ -1280,6 +1285,454 @@ def _endpoint(name: str) -> Mock: assert len(shutdown_endpoints) == 3 +# --------------------------------------------------------------------------- +# _sync_fpm_worker_id: parent -> restored EngineCore child +# --------------------------------------------------------------------------- + + +def _snapshot_config(scheduler_cls=InstrumentedScheduler) -> SimpleNamespace: + """A restored ``VllmConfig`` whose snapshot captured ``scheduler_cls``. + + ``shutdown_timeout`` is read by ``_DecodeWorkerLifecycle.cleanup()`` when a + restore fails, so it has to be present on the decode path. + """ + return SimpleNamespace( + scheduler_config=SimpleNamespace(get_scheduler_cls=lambda: scheduler_cls), + shutdown_timeout=5.0, + ) + + +@pytest.mark.asyncio +async def test_sync_fpm_worker_id_invokes_engine_core_utility(): + """The parent reaches the restored child through the EngineCore utility RPC. + + This is the only channel available: the child was created before the + runtime existed, so it inherited an environment with no worker id and + never re-reads it. + """ + call_utility_async = AsyncMock() + engine_client = SimpleNamespace( + engine_core=SimpleNamespace(call_utility_async=call_utility_async) + ) + + await _sync_fpm_worker_id(engine_client, _snapshot_config(), "8465209922961459") + + call_utility_async.assert_awaited_once_with( + FPM_SET_WORKER_ID_METHOD_NAME, "8465209922961459" + ) + + +@pytest.mark.asyncio +async def test_sync_fpm_worker_id_fails_the_restore_when_the_child_rejects_it(): + """A rejected sync fails the restore instead of degrading to a warning. + + Swallowing it leaves both child-side identity fields empty, the FPM + subscriber drops every sample from the replica, and the planner's + ``_reconcile_fpm_worker_count`` abandons each scaling decision — the + ``worker_count_mismatch`` bug this fix exists to remove, arriving through a + second door. Better to fail before the worker registers than to advertise a + replica the planner cannot see. + """ + engine_client = SimpleNamespace( + engine_core=SimpleNamespace( + call_utility_async=AsyncMock( + side_effect=Exception( + "Call to set_fpm_worker_id method failed: 'EngineCoreProc' " + "object has no attribute 'set_fpm_worker_id'" + ) + ) + ) + ) + + with pytest.raises(RuntimeError, match="invisible to the planner"): + await _sync_fpm_worker_id(engine_client, _snapshot_config(), "8465209922961459") + + +@pytest.mark.asyncio +async def test_sync_fpm_worker_id_is_bounded_when_the_rpc_never_answers(monkeypatch): + """``call_utility_async()`` awaits its future with no deadline of its own. + + An EngineCore rank that stops answering would otherwise hang the restore + here — before endpoint registration, with nothing in the logs to point at. + Expiry takes the same fail-restore path as any other sync failure. + """ + monkeypatch.setattr( + "dynamo.vllm.worker_factory.FPM_SET_WORKER_ID_TIMEOUT_SECONDS", 0.01 + ) + rpc_was_cancelled = asyncio.Event() + + async def never_answers(*_args, **_kwargs) -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + rpc_was_cancelled.set() + raise + + engine_client = SimpleNamespace( + engine_core=SimpleNamespace(call_utility_async=never_answers) + ) + + with pytest.raises(RuntimeError, match="Timed out"): + await _sync_fpm_worker_id(engine_client, _snapshot_config(), "8465209922961459") + + # wait_for() cancels the pending RPC on expiry; a leaked task would keep the + # ZMQ future alive for the life of the process. + assert rpc_was_cancelled.is_set() + + +@pytest.mark.asyncio +async def test_sync_fpm_worker_id_skips_a_snapshot_without_the_instrumented_scheduler(): + """Snapshot mode is independent of FPM mode. + + A snapshot captured with vLLM's own scheduler has no FPM publisher to + retarget, and its child raises on the utility. Now that a rejection is + fatal, such a restore must never make the call at all. + """ + call_utility_async = AsyncMock() + engine_client = SimpleNamespace( + engine_core=SimpleNamespace(call_utility_async=call_utility_async) + ) + + await _sync_fpm_worker_id( + engine_client, _snapshot_config(AsyncScheduler), "8465209922961459" + ) + + call_utility_async.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_sync_fpm_worker_id_syncs_a_subclass_of_the_instrumented_scheduler(): + """``--scheduler-cls`` may name a subclass; it still carries the publisher.""" + + class _Derived(InstrumentedScheduler): + pass + + call_utility_async = AsyncMock() + engine_client = SimpleNamespace( + engine_core=SimpleNamespace(call_utility_async=call_utility_async) + ) + + await _sync_fpm_worker_id( + engine_client, _snapshot_config(_Derived), "8465209922961459" + ) + + call_utility_async.assert_awaited_once_with( + FPM_SET_WORKER_ID_METHOD_NAME, "8465209922961459" + ) + + +@pytest.mark.asyncio +async def test_sync_fpm_worker_id_fails_on_an_unresolvable_scheduler_class(): + """The EngineCore resolved this class at build time, so failing here is fatal.""" + + def raises() -> type: + raise ImportError("No module named 'user.scheduler'") + + call_utility_async = AsyncMock() + engine_client = SimpleNamespace( + engine_core=SimpleNamespace(call_utility_async=call_utility_async) + ) + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(get_scheduler_cls=raises) + ) + + with pytest.raises(ImportError, match="user.scheduler"): + await _sync_fpm_worker_id(engine_client, vllm_config, "8465209922961459") + + call_utility_async.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# The production restore call sites: realtime, decode, prefill +# --------------------------------------------------------------------------- + +# Each entry is the worker-creation method and the disaggregation mode that +# reaches it. Every one of them unpacks a snapshot and must sync the identity. +_RESTORE_PATHS = { + "decode": ("_create_decode_worker", DisaggregationMode.DECODE), + "prefill": ("_create_prefill_worker", DisaggregationMode.PREFILL), + "realtime": ("_create_realtime_worker", DisaggregationMode.AGGREGATED), +} + +_RESTORED_ENDPOINT_ID = "cid-4242" + + +class _ReachedRegistration(Exception): + """Raised by the stubbed ``register_vllm_model``: proof a path got that far.""" + + +def _restore_snapshot_engine( + scheduler_cls: type = InstrumentedScheduler, + call_utility_async=None, +) -> tuple[SnapshotEngineSetupResult, Mock]: + """A snapshot engine whose utility RPC and shutdown are recorded.""" + engine_client = Mock() + engine_client.engine_core.call_utility_async = ( + call_utility_async if call_utility_async is not None else AsyncMock() + ) + vllm_config = Mock(shutdown_timeout=5.0, additional_config={}) + vllm_config.scheduler_config.get_scheduler_cls.return_value = scheduler_cls + vllm_config.cache_config.block_size = 16 + vllm_config.cache_config.num_gpu_blocks = 1024 + vllm_config.parallel_config.data_parallel_size = 1 + vllm_config.parallel_config.data_parallel_rank = 0 + vllm_config.parallel_config.data_parallel_size_local = 1 + vllm_config.model_config.max_model_len = 4096 + engine_setup: EngineSetupResult = ( + engine_client, + vllm_config, + Mock(), + None, + Mock(), + ) + return (engine_setup, Mock()), engine_client + + +async def _enter_restore_path( + path: str, + snapshot_engine: SnapshotEngineSetupResult | None, + register: AsyncMock, + endpoint: Mock, + setup_vllm_engine: Mock | None = None, +) -> None: + """Run one worker-creation path; it ends in ``register`` or the error under test.""" + method_name, mode = _RESTORE_PATHS[path] + factory = _make_factory( + register_vllm_model_fn=register, + setup_vllm_engine_fn=setup_vllm_engine or Mock(), + ) + factory._maybe_create_failover_metrics = Mock(return_value=None) # type: ignore[method-assign] + runtime = Mock() + runtime.endpoint.return_value = endpoint + config = _make_config( + disaggregation_mode=mode, + namespace="dyn", + component="backend", + endpoint="generate", + model="m", + served_model_name="m", + enable_rl=False, + enable_multimodal=False, + frontend_decoding=False, + use_vllm_tokenizer=False, + engine_args=SimpleNamespace(enable_lora=False), + endpoint_types="chat,completions", + ) + await getattr(factory, method_name)( + runtime, + config, + asyncio.Event(), + [], + snapshot_engine=snapshot_engine, + ) + + +@pytest.fixture +def restore_path_stubs(monkeypatch: pytest.MonkeyPatch) -> None: + """Stub what each path calls between the sync and registration.""" + monkeypatch.setenv(ENV_FPM_WORKER_ID, "") + + async def _noop(*_args, **_kwargs) -> None: + return None + + monkeypatch.setattr( + "dynamo.vllm.worker_factory.configure_kv_event_block_size", _noop + ) + monkeypatch.setattr( + "dynamo.vllm.worker_factory.get_dp_range_for_worker", lambda _config: (0, 1) + ) + monkeypatch.setattr("dynamo.vllm.worker_factory.StatLoggerFactory", Mock()) + for handler_cls in ( + "PrefillWorkerHandler", + "DecodeWorkerHandler", + "RealtimeHandler", + "RealtimeTranscriptionHandler", + ): + monkeypatch.setattr( + f"dynamo.vllm.worker_factory.{handler_cls}", + Mock(return_value=Mock(embedding_cache_manager=None)), + ) + for payload_cls in ("VllmHealthCheckPayload", "VllmPrefillHealthCheckPayload"): + monkeypatch.setattr( + f"dynamo.vllm.worker_factory.{payload_cls}", + Mock(return_value=Mock(to_dict=Mock(return_value={}))), + ) + monkeypatch.setattr( + WorkerFactory, "_maybe_get_encode_worker_client", AsyncMock(return_value=None) + ) + monkeypatch.setattr(WorkerFactory, "_maybe_wait_for_failover_lock", AsyncMock()) + monkeypatch.setattr(WorkerFactory, "register_engine_routes", Mock()) + + +def _restored_endpoint() -> Mock: + return Mock( + connection_id=Mock(return_value=_RESTORED_ENDPOINT_ID), + first_token_source=AsyncMock(), + ) + + +def _assert_not_served(endpoint: Mock) -> None: + endpoint.serve_endpoint.assert_not_called() + endpoint.serve_bidirectional_endpoint.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", sorted(_RESTORE_PATHS)) +async def test_every_restore_path_syncs_the_endpoint_id_before_registration( + path: str, restore_path_stubs, monkeypatch +): + """Each restore path syncs the endpoint id, and does so before registering.""" + order: list[str] = [] + call_utility_async = AsyncMock(side_effect=lambda *_: order.append("sync")) + snapshot, _engine_client = _restore_snapshot_engine( + call_utility_async=call_utility_async + ) + + def reach_registration(*_args, **_kwargs): + order.append("register") + raise _ReachedRegistration + + with pytest.raises(_ReachedRegistration): + await _enter_restore_path( + path, + snapshot, + AsyncMock(side_effect=reach_registration), + _restored_endpoint(), + ) + + call_utility_async.assert_awaited_once_with( + FPM_SET_WORKER_ID_METHOD_NAME, _RESTORED_ENDPOINT_ID + ) + assert order == ["sync", "register"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", sorted(_RESTORE_PATHS)) +async def test_restore_without_instrumented_scheduler_registers_without_sync( + path: str, restore_path_stubs, monkeypatch +): + """A snapshot without InstrumentedScheduler skips the sync and still registers.""" + snapshot, engine_client = _restore_snapshot_engine(scheduler_cls=AsyncScheduler) + register = AsyncMock(side_effect=_ReachedRegistration) + + with pytest.raises(_ReachedRegistration): + await _enter_restore_path(path, snapshot, register, _restored_endpoint()) + + engine_client.engine_core.call_utility_async.assert_not_awaited() + register.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", sorted(_RESTORE_PATHS)) +async def test_cold_start_registers_without_syncing_the_fpm_worker_id( + path: str, restore_path_stubs, monkeypatch +): + """Cold start reads the id from the environment, so it registers without a sync.""" + (engine_setup, _stat_logger), engine_client = _restore_snapshot_engine() + register = AsyncMock(side_effect=_ReachedRegistration) + + with pytest.raises(_ReachedRegistration): + await _enter_restore_path( + path, + None, + register, + _restored_endpoint(), + setup_vllm_engine=Mock(return_value=engine_setup), + ) + + engine_client.engine_core.call_utility_async.assert_not_awaited() + register.assert_awaited_once() + + +async def _never_answers(*_args, **_kwargs) -> None: + await asyncio.Event().wait() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", sorted(_RESTORE_PATHS)) +@pytest.mark.parametrize( + ("failure", "match"), + [ + ("rejected", "invisible to the planner"), + ("timeout", "Timed out"), + ("unresolvable", "user.scheduler"), + ], +) +async def test_failed_sync_fails_the_restore_before_registration_and_shuts_down_the_engine( + path: str, failure: str, match: str, restore_path_stubs, monkeypatch +): + """A failed sync blocks registration and shuts the engine down once.""" + if failure == "timeout": + monkeypatch.setattr( + "dynamo.vllm.worker_factory.FPM_SET_WORKER_ID_TIMEOUT_SECONDS", 0.01 + ) + call_utility_async = _never_answers + else: + call_utility_async = AsyncMock( + side_effect=Exception( + "Call to set_fpm_worker_id method failed: scheduler is " + "Scheduler, not InstrumentedScheduler" + ) + ) + snapshot, engine_client = _restore_snapshot_engine( + call_utility_async=call_utility_async + ) + expected_error: type[Exception] = RuntimeError + if failure == "unresolvable": + engine_setup, _stat_logger = snapshot + vllm_config = engine_setup[1] + vllm_config.scheduler_config.get_scheduler_cls.side_effect = ImportError( + "No module named 'user.scheduler'" + ) + expected_error = ImportError + register = AsyncMock(side_effect=_ReachedRegistration) + endpoint = _restored_endpoint() + + with pytest.raises(expected_error, match=match): + await _enter_restore_path(path, snapshot, register, endpoint) + + register.assert_not_awaited() + _assert_not_served(endpoint) + engine_client.shutdown.assert_called_once_with(timeout=5.0) + + +@pytest.mark.asyncio +async def test_sync_or_shutdown_keeps_the_sync_error_when_shutdown_also_fails( + caplog, +): + """A failing shutdown is logged; the sync error is what propagates.""" + engine_client = SimpleNamespace( + engine_core=SimpleNamespace( + call_utility_async=AsyncMock(side_effect=Exception("rejected")) + ), + shutdown=Mock(side_effect=RuntimeError("engine already dead")), + ) + + with caplog.at_level(logging.ERROR): + with pytest.raises(RuntimeError, match="invisible to the planner"): + await _sync_fpm_worker_id_or_shutdown( + engine_client, _snapshot_config(), "8465209922961459" + ) + + engine_client.shutdown.assert_called_once_with(timeout=5.0) + assert "Failed to shut down the restored engine" in caplog.text + + +def test_fpm_utility_name_matches_the_installed_method(): + """The caller's method name must match what the scheduler module installs. + + The two sides live in different modules and are linked only by this + string. A rename on either side would otherwise fail silently at runtime, + surfacing as an ``AttributeError`` folded into the utility's + ``failure_message`` long after deploy. + """ + from vllm.v1.engine.core import EngineCore + + import dynamo.vllm.instrumented_scheduler # noqa: F401 installs the patch + + assert hasattr(EngineCore, FPM_SET_WORKER_ID_METHOD_NAME) + + @pytest.mark.asyncio @pytest.mark.parametrize( ("embedding_frontend_tokenization", "expected_model_input"), diff --git a/components/src/dynamo/vllm/worker_factory.py b/components/src/dynamo/vllm/worker_factory.py index 83b1910c1935..38256f38e858 100644 --- a/components/src/dynamo/vllm/worker_factory.py +++ b/components/src/dynamo/vllm/worker_factory.py @@ -48,7 +48,11 @@ VllmHealthCheckPayload, VllmPrefillHealthCheckPayload, ) -from .instrumented_scheduler import ENV_FPM_BENCHMARK_OUTPUT_PATH, ENV_FPM_WORKER_ID +from .instrumented_scheduler import ( + ENV_FPM_BENCHMARK_OUTPUT_PATH, + ENV_FPM_WORKER_ID, + InstrumentedScheduler, +) from .multimodal_handlers import EncodeWorkerHandler from .pooling_handlers import ClassifyWorkerHandler from .publisher import StatLoggerFactory @@ -106,6 +110,78 @@ def _register_request_cache_metrics( EngineSetupResult = tuple[AsyncLLM, VllmConfig, Any, Any, Optional[LLMBackendMetrics]] SnapshotEngineSetupResult = tuple[EngineSetupResult, StatLoggerFactory] +# Utility installed on EngineCore by instrumented_scheduler; retargets a +# snapshot-restored child whose FPM worker id was baked as "". +FPM_SET_WORKER_ID_METHOD_NAME = "set_fpm_worker_id" + +# call_utility_async() has no deadline of its own (as of vLLM 0.30.0); +# under DP it fans out to every rank, so this bounds the slowest one. +FPM_SET_WORKER_ID_TIMEOUT_SECONDS = 30.0 + + +def _snapshot_uses_instrumented_scheduler(vllm_config: VllmConfig) -> bool: + """Whether the snapshot's EngineCore runs an InstrumentedScheduler. + + False only for a user-supplied ``--scheduler-cls``, which has no FPM publisher. + Resolution errors propagate: the EngineCore already resolved the same class. + """ + scheduler_cls = vllm_config.scheduler_config.get_scheduler_cls() + return isinstance(scheduler_cls, type) and issubclass( + scheduler_cls, InstrumentedScheduler + ) + + +async def _sync_fpm_worker_id( + engine_client: AsyncLLM, vllm_config: VllmConfig, new_worker_id: str +) -> None: + """Push ``new_worker_id`` into a restored EngineCore's FPM publisher. + + Raises on failure: a replica publishing an empty id is invisible to the planner. + """ + if not _snapshot_uses_instrumented_scheduler(vllm_config): + return + + try: + await asyncio.wait_for( + engine_client.engine_core.call_utility_async( + FPM_SET_WORKER_ID_METHOD_NAME, new_worker_id + ), + timeout=FPM_SET_WORKER_ID_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError as exc: + raise RuntimeError( + f"Timed out after {FPM_SET_WORKER_ID_TIMEOUT_SECONDS}s setting FPM " + f"worker_id {new_worker_id} on the restored engine; an EngineCore rank " + f"is not answering utility RPCs" + ) from exc + except Exception as exc: + # vLLM surfaces a failed utility as a bare Exception(failure_message). + raise RuntimeError( + f"Failed to set FPM worker_id {new_worker_id} on the restored engine; " + f"this replica would publish an empty worker id and stay invisible to " + f"the planner" + ) from exc + + +async def _sync_fpm_worker_id_or_shutdown( + engine_client: AsyncLLM, vllm_config: VllmConfig, new_worker_id: str +) -> None: + """``_sync_fpm_worker_id``, shutting the engine down on failure. + + For realtime and prefill, which have no lifecycle owning the engine yet. + """ + try: + await _sync_fpm_worker_id(engine_client, vllm_config, new_worker_id) + except BaseException: + try: + engine_client.shutdown(timeout=vllm_config.shutdown_timeout) + except Exception: + logger.exception( + "Failed to shut down the restored engine after the FPM worker_id " + "sync failed" + ) + raise + def _benchmark_rank_path(base_path: Path, dp_rank: int) -> Path: if dp_rank == 0: @@ -842,6 +918,10 @@ async def _create_realtime_worker( ) = engine_setup os.environ[ENV_FPM_WORKER_ID] = fpm_worker_id factory.bind_endpoint(generate_endpoint) + + await _sync_fpm_worker_id_or_shutdown( + engine_client, vllm_config, fpm_worker_id + ) else: factory = StatLoggerFactory(endpoint=generate_endpoint) ( @@ -1316,8 +1396,11 @@ async def _run_decode_worker( prometheus_temp_dir, _component_gauges, ) = self.setup_vllm_engine(config, factory, fpm_worker_id=fpm_worker_id) + # Sync after the lifecycle owns the engine, so a failure shuts it down. lifecycle.engine_client = engine_client lifecycle.vllm_config = vllm_config + if snapshot_engine is not None: + await _sync_fpm_worker_id(engine_client, vllm_config, fpm_worker_id) await configure_kv_event_block_size(engine_client, vllm_config) # TODO Hack to get data, move this to registering in TBD @@ -1378,8 +1461,8 @@ async def _run_decode_worker( handler.kv_publishers = kv_publishers # Set up forward pass metrics relay (child ZMQ -> event plane). - # In checkpoint mode the engine was created before the runtime, so - # ForwardPassMetrics.worker_id will be empty (relay still works). + # In checkpoint mode _sync_fpm_worker_id() has already set the child's + # ForwardPassMetrics.worker_id. fpm_relays = self.setup_fpm_relay(config, generate_endpoint, vllm_config) if fpm_relays: handler.fpm_relays = fpm_relays @@ -1606,11 +1689,11 @@ async def _run_prefill_worker( _component_gauges, ) = engine_setup factory.bind_endpoint(generate_endpoint) - # TODO: The scheduler in the child process still has worker_id="" - # because the engine was forked before the runtime existed. - # Propagating the new ID to the child requires shared memory or - # a restart of the EngineCore process. os.environ[ENV_FPM_WORKER_ID] = fpm_worker_id + + await _sync_fpm_worker_id_or_shutdown( + engine_client, vllm_config, fpm_worker_id + ) else: factory = StatLoggerFactory(endpoint=generate_endpoint) ( @@ -1675,8 +1758,8 @@ async def _run_prefill_worker( handler.kv_publishers = kv_publishers # Set up forward pass metrics relay (child ZMQ -> event plane). - # In checkpoint mode the engine was created before the runtime, so - # ForwardPassMetrics.worker_id will be empty (relay still works). + # In checkpoint mode _sync_fpm_worker_id() has already set the child's + # ForwardPassMetrics.worker_id. fpm_relays = self.setup_fpm_relay(config, generate_endpoint, vllm_config) if fpm_relays: handler.fpm_relays = fpm_relays diff --git a/tests/report_pytest_markers.py b/tests/report_pytest_markers.py index b8b22cfeb1d4..04158b8d16c0 100755 --- a/tests/report_pytest_markers.py +++ b/tests/report_pytest_markers.py @@ -245,6 +245,7 @@ "vllm.v1.core.sched.output", "vllm.v1.engine", "vllm.v1.engine.async_llm", + "vllm.v1.engine.core", "vllm.v1.engine.exceptions", "vllm.v1.engine.input_processor", "vllm.v1.engine.output_processor",