diff --git a/tests/entrypoints/launchers/test_cli_args.py b/tests/entrypoints/launchers/test_cli_args.py index e84abe5d2617..f6a10d5e952b 100644 --- a/tests/entrypoints/launchers/test_cli_args.py +++ b/tests/entrypoints/launchers/test_cli_args.py @@ -69,6 +69,27 @@ def test_config_arg_parsing(serve_parser, cli_config_file): assert args.port == 9000 +def test_snapshot_config_arg_parsing(serve_parser): + args = serve_parser.parse_args( + [ + "--snapshot-config", + json.dumps( + { + "snapshot_metadata": "/snapshot/snapshot_metadata.json", + } + ), + ] + ) + + assert args.snapshot_config.snapshot_metadata == "/snapshot/snapshot_metadata.json" + + +def test_snapshot_config_disabled_by_default(serve_parser): + args = serve_parser.parse_args([]) + + assert args.snapshot_config is None + + ### Tests for LoRA module parsing def test_valid_key_value_format(serve_parser): # Test old format: name=path diff --git a/tests/snapshot/test_config.py b/tests/snapshot/test_config.py new file mode 100644 index 000000000000..7217ee130c67 --- /dev/null +++ b/tests/snapshot/test_config.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import patch + +from vllm.config.parallel import ParallelConfig + + +def test_reserve_snapshot_ports(): + parallel_config = SimpleNamespace( + _snapshot_data_parallel_port_list=None, + ) + + with patch( + "vllm.config.parallel.get_open_ports_list", return_value=[2000, 2001] + ) as get_ports: + ParallelConfig.reserve_snapshot_ports(parallel_config) + + get_ports.assert_called_once_with(2) + assert parallel_config._snapshot_data_parallel_port_list == [2000, 2001] + + +def test_restore_uses_snapshot_port_queue(): + parallel_config = SimpleNamespace( + _snapshot_data_parallel_port_list=[2000, 2001], + _data_parallel_master_port_list=[1000], + data_parallel_master_port=999, + ) + + with patch("vllm.snapshot.utils.is_restore", return_value=True): + port = ParallelConfig.get_next_dp_init_port(parallel_config) + + assert port == 2001 + assert parallel_config._snapshot_data_parallel_port_list == [2000] + assert parallel_config._data_parallel_master_port_list == [1000] + + +def test_restore_without_snapshot_ports_uses_normal_queue(): + parallel_config = SimpleNamespace( + _snapshot_data_parallel_port_list=None, + _data_parallel_master_port_list=[1000], + data_parallel_master_port=999, + ) + + with patch("vllm.snapshot.utils.is_restore", return_value=True): + port = ParallelConfig.get_next_dp_init_port(parallel_config) + + assert port == 1000 diff --git a/tests/snapshot/test_model_state.py b/tests/snapshot/test_model_state.py new file mode 100644 index 000000000000..a4e694a22749 --- /dev/null +++ b/tests/snapshot/test_model_state.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.config import CompilationConfig, SnapshotConfig, set_current_vllm_config +from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbeddingBase + + +@pytest.mark.parametrize("snapshot_enabled", [False, True]) +def test_rotary_cache_is_persistent_only_for_snapshot(snapshot_enabled): + vllm_config = SimpleNamespace( + snapshot_config=SnapshotConfig() if snapshot_enabled else None, + compilation_config=CompilationConfig(custom_ops=["none"]), + ) + with set_current_vllm_config(vllm_config): + layer = RotaryEmbeddingBase( + head_size=8, + rotary_dim=8, + max_position_embeddings=16, + base=10000, + is_neox_style=True, + dtype=torch.float32, + ) + + assert ("cos_sin_cache" in layer.state_dict()) is snapshot_enabled diff --git a/tests/snapshot/test_utils.py b/tests/snapshot/test_utils.py new file mode 100644 index 000000000000..0739925ec792 --- /dev/null +++ b/tests/snapshot/test_utils.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from unittest.mock import patch + +from vllm.snapshot.utils import is_restore + + +def test_is_restore_delegates_to_current_platform(): + with patch( + "vllm.snapshot.utils.current_platform.is_restore", return_value=True + ) as platform_is_restore: + assert is_restore() + + platform_is_restore.assert_called_once_with() diff --git a/tests/utils_/test_network_utils.py b/tests/utils_/test_network_utils.py index 02889c20f66a..95b36a7506b9 100644 --- a/tests/utils_/test_network_utils.py +++ b/tests/utils_/test_network_utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import socket import threading +from unittest.mock import MagicMock, call, patch import pytest import zmq @@ -9,12 +10,14 @@ from vllm.utils import network_utils from vllm.utils.network_utils import ( get_file_store_init_method, + get_ip, get_open_port, get_open_ports_list, get_tcp_uri, join_host_port, make_zmq_path, make_zmq_socket, + replace_zmq_tcp_host, split_host_port, split_zmq_path, ) @@ -47,6 +50,53 @@ def target(): return result["value"] +def _udp_socket(ip: str | None = None, error: Exception | None = None): + context = MagicMock() + sock = context.__enter__.return_value + sock.connect.side_effect = error + if ip is not None: + sock.getsockname.return_value = (ip, 0) + return context + + +def test_get_ip_force_probes_ipv4_first(): + ipv4_socket = _udp_socket("10.0.0.2") + + with patch( + "vllm.utils.network_utils.socket.socket", return_value=ipv4_socket + ) as factory: + assert get_ip(force=True) == "10.0.0.2" + + factory.assert_called_once_with(socket.AF_INET, socket.SOCK_DGRAM) + + +def test_get_ip_force_falls_back_to_ipv6(): + ipv4_socket = _udp_socket(error=OSError()) + ipv6_socket = _udp_socket("2001:db8::2") + + with patch( + "vllm.utils.network_utils.socket.socket", + side_effect=[ipv4_socket, ipv6_socket], + ) as factory: + assert get_ip(force=True) == "2001:db8::2" + + assert factory.call_args_list == [ + call(socket.AF_INET, socket.SOCK_DGRAM), + call(socket.AF_INET6, socket.SOCK_DGRAM), + ] + + +def test_get_ip_force_uses_default_when_probe_fails(): + with ( + patch( + "vllm.utils.network_utils.socket.socket", + side_effect=[OSError(), OSError()], + ), + pytest.warns(UserWarning, match="using 0.0.0.0 by default"), + ): + assert get_ip(force=True) == "0.0.0.0" + + def test_get_open_port(monkeypatch: pytest.MonkeyPatch): with monkeypatch.context() as m: m.setenv("VLLM_PORT", "5678") @@ -210,6 +260,18 @@ def test_make_zmq_path(): assert make_zmq_path("tcp", "::1", "5555") == "tcp://[::1]:5555" +@pytest.mark.parametrize( + "path,host,expected", + [ + ("tcp://10.0.0.1:5555", "10.0.0.2", "tcp://10.0.0.2:5555"), + ("tcp://10.0.0.1:5555", "::1", "tcp://[::1]:5555"), + ("ipc://some_path", "10.0.0.2", "ipc://some_path"), + ], +) +def test_replace_zmq_tcp_host(path, host, expected): + assert replace_zmq_tcp_host(path, host) == expected + + def test_get_tcp_uri(): assert get_tcp_uri("127.0.0.1", 5555) == "tcp://127.0.0.1:5555" assert get_tcp_uri("::1", 5555) == "tcp://[::1]:5555" diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index ff2ba0493d8d..0bfbc5eef102 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -40,6 +40,7 @@ from vllm.config.profiler import ProfilerConfig from vllm.config.reasoning import ReasoningConfig from vllm.config.scheduler import SchedulerConfig +from vllm.config.snapshot import SnapshotConfig from vllm.config.speculative import SpeculativeConfig from vllm.config.speech_to_text import SpeechToTextConfig, SpeechToTextParams from vllm.config.structured_outputs import StructuredOutputsConfig @@ -121,6 +122,8 @@ "ReasoningConfig", # From vllm.config.scheduler "SchedulerConfig", + # From vllm.config.snapshot + "SnapshotConfig", # From vllm.config.speculative "SpeculativeConfig", # From vllm.config.speech_to_text diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 0213e06283df..a564527a764a 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -346,6 +346,9 @@ class is dynamically inherited by the worker class. This is used to inject Set to be private as it's not intended to be configured by users. """ + _snapshot_data_parallel_port_list: list[int] | None = None + """Ports reserved for rebuilding communication groups after restore.""" + _coord_store_port: int = 0 """Port of the coordination TCPStore. Can be set by the API server; workers connect as clients to exchange self-picked group ports at runtime.""" @@ -617,6 +620,13 @@ def get_next_dp_init_port(self) -> int: pop a new port from the prepared port list each time we need to initialize a new process group related to data parallelism. """ + from vllm.snapshot.utils import is_restore + + if is_restore() and self._snapshot_data_parallel_port_list is not None: + if not self._snapshot_data_parallel_port_list: + raise RuntimeError("No port reserved for snapshot restore") + return self._snapshot_data_parallel_port_list.pop() + if self._data_parallel_master_port_list: answer = self._data_parallel_master_port_list.pop() else: @@ -625,6 +635,12 @@ def get_next_dp_init_port(self) -> int: return answer + def reserve_snapshot_ports(self) -> None: + if self._snapshot_data_parallel_port_list is not None: + return + + self._snapshot_data_parallel_port_list = get_open_ports_list(2) + def _pick_stateless_dp_port(self) -> tuple[int, socket.socket | None]: """Return ``(port, listen_socket)`` for DP group init. @@ -838,6 +854,7 @@ def compute_hash(self): "data_parallel_master_port", "_data_parallel_master_port_list", "_coord_store_port", + "_snapshot_data_parallel_port_list", "data_parallel_rpc_port", "rank", "master_addr", diff --git a/vllm/config/snapshot.py b/vllm/config/snapshot.py new file mode 100644 index 000000000000..6e5afe281c8e --- /dev/null +++ b/vllm/config/snapshot.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .utils import config + + +@config +class SnapshotConfig: + """Configuration for container snapshot lifecycle management.""" + + snapshot_metadata: str | None = None + """Snapshot metadata file used for checkpoint and restore.""" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ad8f077148c3..7d4dfdc17efa 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -48,6 +48,7 @@ from .profiler import ProfilerConfig from .reasoning import ReasoningConfig from .scheduler import SchedulerConfig +from .snapshot import SnapshotConfig from .speculative import EagleModelTypes, NgramGPUTypes, SpeculativeConfig from .structured_outputs import StructuredOutputsConfig from .utils import SupportsHash, config, get_field, replace @@ -428,6 +429,8 @@ class VllmConfig: """The configurations for custom encoder cache manager.""" reasoning_config: ReasoningConfig | None = None """The configurations for reasoning model.""" + snapshot_config: SnapshotConfig | None = None + """Container snapshot configuration.""" # some opaque config, only used to provide additional information # for the hash computation, mainly used for testing, debugging or out of # tree config registration. @@ -1270,6 +1273,9 @@ def __post_init__(self): # unset falls back to the stock ones. self.parallel_config.set_dcp_defaults() + if self.snapshot_config is not None: + self.parallel_config.reserve_snapshot_ports() + if self.model_config is not None: self.model_config.verify_with_parallel_config(self.parallel_config) self.model_config.verify_dual_chunk_attention_config(self.load_config) diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index dc3022a16fec..5e96c6f2f475 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -139,6 +139,12 @@ def _split_tensor_dict( _group_name_counter: dict[str, int] = {} +def reset_group_name_registry() -> None: + """Reset the group name counter and stale group references.""" + _group_name_counter.clear() + _groups.clear() + + def _get_unique_name(name: str) -> str: """Get a unique name for the group. Example: @@ -1808,8 +1814,14 @@ def init_distributed_environment( # adjust the world size to take into account data parallelism world_size = parallel_config.world_size_across_dp - # Use appropriate IP and port based on configuration - if parallel_config.nnodes > 1: + # Use appropriate IP and port based on configuration. + from vllm.snapshot.utils import is_restore + + if config.snapshot_config is not None and is_restore(): + ip = parallel_config.data_parallel_master_ip + port = parallel_config.get_next_dp_init_port() + distributed_init_method = get_distributed_init_method(ip, port) + elif parallel_config.nnodes > 1: ip = parallel_config.master_addr port = parallel_config.master_port distributed_init_method = get_distributed_init_method(ip, port) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index b420e84f542b..5e6531f996cd 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -61,6 +61,7 @@ ProfilerConfig, ReasoningConfig, SchedulerConfig, + SnapshotConfig, SpeculativeConfig, StructuredOutputsConfig, UVAOffloadConfig, @@ -659,6 +660,7 @@ class EngineArgs: structured_outputs_config: StructuredOutputsConfig = get_field( VllmConfig, "structured_outputs_config" ) + snapshot_config: SnapshotConfig | None = get_field(VllmConfig, "snapshot_config") reasoning_parser: str = StructuredOutputsConfig.reasoning_parser reasoning_parser_plugin: str | None = None @@ -1734,6 +1736,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: vllm_group.add_argument( "--structured-outputs-config", **vllm_kwargs["structured_outputs_config"] ) + vllm_group.add_argument("--snapshot-config", **vllm_kwargs["snapshot_config"]) vllm_group.add_argument("--profiler-config", **vllm_kwargs["profiler_config"]) vllm_group.add_argument( "--optimization-level", **vllm_kwargs["optimization_level"] @@ -2668,6 +2671,7 @@ def create_engine_config( ec_transfer_config=self.ec_transfer_config, ec_manager_config=self.ec_manager_config, reasoning_config=self.reasoning_config, + snapshot_config=self.snapshot_config, profiler_config=self.profiler_config, additional_config=self.additional_config, optimization_level=self.optimization_level, diff --git a/vllm/model_executor/layers/rotary_embedding/base.py b/vllm/model_executor/layers/rotary_embedding/base.py index 41291a9b56d8..0b0f24e5c9ca 100644 --- a/vllm/model_executor/layers/rotary_embedding/base.py +++ b/vllm/model_executor/layers/rotary_embedding/base.py @@ -5,6 +5,7 @@ import torch from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import get_current_vllm_config_or_none from vllm.model_executor.custom_op import CustomOp from .common import ApplyRotaryEmb @@ -60,7 +61,14 @@ def __init__( if not self.use_flashinfer: cache = cache.to(dtype) self.cos_sin_cache: torch.Tensor - self.register_buffer("cos_sin_cache", cache, persistent=False) + vllm_config = get_current_vllm_config_or_none() + self.register_buffer( + "cos_sin_cache", + cache, + persistent=( + vllm_config is not None and vllm_config.snapshot_config is not None + ), + ) # Reuse a precomputed bf16 cache for the AITER compile path. if self.use_aiter and cache.dtype != torch.bfloat16: diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 7ad0b42da774..55029a95179b 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -228,6 +228,10 @@ def is_sleep_mode_available(self) -> bool: # all ROCm platforms for now. return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM, PlatformEnum.XPU) + def is_restore(self) -> bool: + """Return whether the process was restored from a container snapshot.""" + return False + def is_cumem_allocator_available(self) -> bool: try: from vllm.device_allocator.cumem import cumem_available diff --git a/vllm/snapshot/__init__.py b/vllm/snapshot/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/snapshot/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/snapshot/utils.py b/vllm/snapshot/utils.py new file mode 100644 index 000000000000..efc1ed9378be --- /dev/null +++ b/vllm/snapshot/utils.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import os + +from vllm.platforms import current_platform + +RETRY_INTERVAL = 1.0 +RETRY_LOG_FREQUENCY = 60 + + +def is_restore() -> bool: + return current_platform.is_restore() + + +def load_snapshot_metadata(file_path: str, field_name: str) -> str: + if not os.path.exists(file_path): + raise FileNotFoundError(f"Snapshot metadata file not found: {file_path}") + + with open(file_path, encoding="utf-8") as file: + try: + data = json.load(file) + except Exception as exc: + raise ValueError( + f"Snapshot metadata is not valid JSON: {file_path}: {exc}" + ) from exc + + if not isinstance(data, dict): + raise ValueError( + "Snapshot metadata JSON root must be an object, not an array or " + f"scalar: {file_path}" + ) + + field_value = data.get(field_name) + if not isinstance(field_value, str): + raise ValueError( + "Snapshot metadata requires string field: " + f"{field_name}, but got {type(field_value)}" + ) + return field_value diff --git a/vllm/utils/network_utils.py b/vllm/utils/network_utils.py index 356d913879a1..0d3119c30a72 100644 --- a/vllm/utils/network_utils.py +++ b/vllm/utils/network_utils.py @@ -31,18 +31,19 @@ def close_sockets(sockets: Sequence[zmq.Socket | zmq.asyncio.Socket]): sock.close(linger=0) -def get_ip() -> str: - host_ip = envs.VLLM_HOST_IP - if "HOST_IP" in os.environ and "VLLM_HOST_IP" not in os.environ: - logger.warning( - "The environment variable HOST_IP is deprecated and ignored, as" - " it is often used by Docker and other software to" - " interact with the container's network stack. Please " - "use VLLM_HOST_IP instead to set the IP address for vLLM processes" - " to communicate with each other." - ) - if host_ip: - return host_ip +def get_ip(*, force: bool = False) -> str: + if not force: + host_ip = envs.VLLM_HOST_IP + if "HOST_IP" in os.environ and "VLLM_HOST_IP" not in os.environ: + logger.warning( + "The environment variable HOST_IP is deprecated and ignored, as" + " it is often used by Docker and other software to" + " interact with the container's network stack. Please " + "use VLLM_HOST_IP instead to set the IP address for vLLM processes" + " to communicate with each other." + ) + if host_ip: + return host_ip # IP is not set, try to get it from the network interface @@ -306,6 +307,14 @@ def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str: return f"{scheme}://{host}:{port}" +def replace_zmq_tcp_host(path: str, host: str) -> str: + """Replace the host of a TCP ZMQ endpoint, preserving its port.""" + scheme, _, port = split_zmq_path(path) + if scheme != "tcp": + return path + return make_zmq_path(scheme, host, int(port)) + + # Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501 def make_zmq_socket( ctx: zmq.asyncio.Context | zmq.Context, # type: ignore[name-defined] diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 9bce2b7687ef..424bc6bcdb5a 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -1268,6 +1268,15 @@ def wait_for_engine_startup( and not parallel_config.data_parallel_external_lb ) + parallel_config_keys: tuple[str, ...] = ( + "data_parallel_master_ip", + "data_parallel_master_port", + "_data_parallel_master_port_list", + "data_parallel_size", + ) + if parallel_config._snapshot_data_parallel_port_list is not None: + parallel_config_keys += ("_snapshot_data_parallel_port_list",) + # 1. Engine processes if isinstance(launch.engine_manager, CoreEngineProcManager): for sentinel in launch.engine_manager.sentinels(): @@ -1367,13 +1376,7 @@ def wait_for_engine_startup( EngineHandshakeMetadata( addresses=launch.addresses, parallel_config={ - k: getattr(parallel_config, k) - for k in ( - "data_parallel_master_ip", - "data_parallel_master_port", - "_data_parallel_master_port_list", - "data_parallel_size", - ) + k: getattr(parallel_config, k) for k in parallel_config_keys } if coordinated_dp else {}, diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 964f16ed2421..76c830c50beb 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -159,6 +159,24 @@ def set_kv_cache_layout(self, layout_name: str) -> None: """Publish the resolved KV cache layout to the workers.""" self.collective_rpc("set_kv_cache_layout", args=(layout_name,)) + def suspend(self, model_save_path: str | None = None) -> None: + self.collective_rpc("suspend", args=(model_save_path,)) + + def device_unlock(self) -> None: + self.collective_rpc("device_unlock") + + def resume( + self, + local_ip: str, + data_parallel_master_ip: str, + model_path: str | None = None, + new_engine_id: str | None = None, + ) -> None: + self.collective_rpc( + "resume", + args=(local_ip, data_parallel_master_ip, model_path, new_engine_id), + ) + @overload def collective_rpc( self, diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index cf5b0e7ce334..1edab17f307e 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -131,6 +131,24 @@ def synchronize_device(self) -> None: if torch.accelerator.is_available(): torch.accelerator.synchronize() + def suspend(self, model_save_path: str | None = None) -> None: + """Suspend device state for a process snapshot.""" + raise NotImplementedError + + def device_unlock(self) -> None: + """Release device snapshot resources after checkpoint.""" + raise NotImplementedError + + def resume( + self, + local_ip: str, + data_parallel_master_ip: str, + model_path: str | None = None, + new_engine_id: str | None = None, + ) -> None: + """Restore device state after a process snapshot.""" + raise NotImplementedError + def init_device(self) -> None: """Initialize device state, such as loading the model or other on-device memory allocations.