[Core] Add container snapshot configuration and Worker lifecycle hooks (1/3)

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

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

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

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()

tests/utils_/test_network_utils.py

@@ -2,19 +2,22 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import socket
 import threading
+from unittest.mock import MagicMock, call, patch
 
 import pytest
 import zmq
 
 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"

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

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",

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."""

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)

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)

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,