From 4c793b234fff0b909893dfdbcccd78b98d6fc4e0 Mon Sep 17 00:00:00 2001 From: halt95 Date: Wed, 23 Sep 2026 00:01:58 -0400 Subject: [PATCH 031/103] ple_offload: carry the host-mapped pull transport onto v0.30.0 Restores the fork's PLE offload package (not upstream) and its pull transport: the offload process holds no CUDA context; the captured op waits on a device-mapped host flag and copies host->GPU on the model stream; publication is fenced; releases are synchronous across dummy/capture work; failures close the rank. Includes the test-only fault-injection cuts for the transport. --- probe/faultcut_gate.py | 46 + tests/v1/ple_offload/test_correspondence.py | 197 +++ tests/v1/ple_offload/test_fatal_rpc.py | 145 ++ .../test_fault_bootstrap_disarmed.py | 80 + .../ple_offload/test_fault_injector_basics.py | 260 ++++ .../ple_offload/test_fault_injector_gates.py | 468 ++++++ .../v1/ple_offload/test_fault_wire_frames.py | 96 ++ .../v1/ple_offload/test_ipc_socket_unlink.py | 133 ++ tests/v1/ple_offload/test_pending_rpc.py | 81 + .../v1/ple_offload/test_progress_deadline.py | 216 +++ tests/v1/ple_offload/test_shm_unlink.py | 397 +++++ tests/v1/ple_offload/test_shutdown_order.py | 171 +++ .../ple_offload/test_spawn_bootstrap_cuda.py | 352 +++++ tests/v1/ple_offload/test_test_fault.py | 293 ++++ .../ple_offload/test_transport_bounded_io.py | 343 +++++ .../test_transport_fatal_helper.py | 387 +++++ .../layers/ple_offload_layer.py | 433 ++++++ vllm/v1/executor/multiproc_executor.py | 31 +- vllm/v1/ple_offload/__init__.py | 2 + vllm/v1/ple_offload/connector.py | 855 +++++++++++ vllm/v1/ple_offload/fatal.py | 153 ++ vllm/v1/ple_offload/progress.py | 314 ++++ vllm/v1/ple_offload/protocol.py | 50 + vllm/v1/ple_offload/worker.py | 1298 +++++++++++++++++ vllm/v1/worker/gpu/cudagraph_utils.py | 140 +- vllm/v1/worker/gpu/model_runner.py | 140 ++ vllm/v1/worker/gpu_worker.py | 145 ++ vllm_ple_offload_bootstrap/__init__.py | 210 +++ vllm_ple_offload_bootstrap/__main__.py | 7 + 29 files changed, 7429 insertions(+), 14 deletions(-) create mode 100755 probe/faultcut_gate.py create mode 100644 tests/v1/ple_offload/test_correspondence.py create mode 100644 tests/v1/ple_offload/test_fatal_rpc.py create mode 100644 tests/v1/ple_offload/test_fault_bootstrap_disarmed.py create mode 100644 tests/v1/ple_offload/test_fault_injector_basics.py create mode 100644 tests/v1/ple_offload/test_fault_injector_gates.py create mode 100644 tests/v1/ple_offload/test_fault_wire_frames.py create mode 100644 tests/v1/ple_offload/test_ipc_socket_unlink.py create mode 100644 tests/v1/ple_offload/test_pending_rpc.py create mode 100644 tests/v1/ple_offload/test_progress_deadline.py create mode 100644 tests/v1/ple_offload/test_shm_unlink.py create mode 100644 tests/v1/ple_offload/test_shutdown_order.py create mode 100644 tests/v1/ple_offload/test_spawn_bootstrap_cuda.py create mode 100644 tests/v1/ple_offload/test_test_fault.py create mode 100644 tests/v1/ple_offload/test_transport_bounded_io.py create mode 100644 tests/v1/ple_offload/test_transport_fatal_helper.py create mode 100644 vllm/model_executor/layers/ple_offload_layer.py create mode 100644 vllm/v1/ple_offload/__init__.py create mode 100644 vllm/v1/ple_offload/connector.py create mode 100644 vllm/v1/ple_offload/fatal.py create mode 100644 vllm/v1/ple_offload/progress.py create mode 100644 vllm/v1/ple_offload/protocol.py create mode 100644 vllm/v1/ple_offload/worker.py create mode 100644 vllm_ple_offload_bootstrap/__init__.py create mode 100644 vllm_ple_offload_bootstrap/__main__.py diff --git a/probe/faultcut_gate.py b/probe/faultcut_gate.py new file mode 100755 index 0000000..3eb74b5 --- /dev/null +++ b/probe/faultcut_gate.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python3 +"""CPU-testable validation for external gates excluded from GPU campaigns.""" + +from collections.abc import Callable + + +def validate_partial_send_gate( + *, + pid: int, + expected_start: int, + observed_start: int, + alive: bool, + bytes_sent: int, + frame_size: int, + completed: bool, + child_not_reading: bool, +) -> tuple[bool, str]: + if pid <= 1 or not alive: + return False, "target is stale or not alive" + if expected_start <= 0 or observed_start != expected_start: + return False, "target PID was reused" + if frame_size <= 0 or not 0 < bytes_sent < frame_size: + return False, "send is zero-byte or complete" + if completed: + return False, "send already completed" + if not child_not_reading: + return False, "child non-reading state was not acknowledged" + return True, "valid partial send" + + +def deliver_checked_signal( + *, + pid: int, + expected_start: int, + read_start: Callable[[int], int], + alive: Callable[[int], bool], + send: Callable[[int, int], bool], + signal_number: int, +) -> tuple[bool, str]: + if pid <= 1 or not alive(pid): + return False, "target is stale or not alive" + if read_start(pid) != expected_start: + return False, "target PID was reused" + if not send(pid, signal_number): + return False, "signal delivery failed" + return True, "signal delivered" diff --git a/tests/v1/ple_offload/test_correspondence.py b/tests/v1/ple_offload/test_correspondence.py new file mode 100644 index 0000000..e8d8b9e --- /dev/null +++ b/tests/v1/ple_offload/test_correspondence.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from vllm.config import CUDAGraphMode +from vllm.model_executor.layers.ple_offload_layer import ( + _record_ple_copy_geometry, +) +from vllm.v1.ple_offload.connector import PleOffloadConnector +from vllm.v1.ple_offload.protocol import PleOffloadRequest +from vllm.v1.ple_offload.worker import ( + PleOffloadInputBuffers, + PleOffloadRunner, +) +from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + CudaGraphManager, +) +from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import ( + SpeculatorCudaGraphManager, +) +from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + +class _ConnectorForValidation: + _input_ids_buf = torch.empty(4, dtype=torch.int32) + _query_start_loc_buf = torch.empty(3, dtype=torch.int32) + _rowdump_n = None + + def _signal_dummy_outputs_impl(self, num_tokens: int) -> None: + raise AssertionError("must fail before dummy publication") + + def _launch(self, num_reqs: int, num_tokens: int) -> None: + raise AssertionError("must fail before request publication") + + +def test_request_consumer_count_mismatch_precedes_publication() -> None: + connector = _ConnectorForValidation() + with pytest.raises(RuntimeError, match="request/consumer"): + PleOffloadConnector._prepare_forward_impl( + connector, + 1, + 3, + False, + consumer_num_tokens=4, + ) + + +def _runner(query_start_loc: list[int]) -> PleOffloadRunner: + runner = PleOffloadRunner.__new__(PleOffloadRunner) + flag = torch.zeros(1, dtype=torch.int32) + target = SimpleNamespace( + host_output_buffer=torch.zeros(4, 2), + host_flag=flag, + ) + runner._worker_targets = {0: {"ple": [target]}} + runner._scratch_bufs = {0: {"ple": torch.empty(4, 2)}} + query = torch.zeros(3, dtype=torch.int32) + query[: len(query_start_loc)] = torch.tensor( + query_start_loc, dtype=torch.int32 + ) + runner._input_bufs = { + 0: PleOffloadInputBuffers( + input_ids_buf=torch.zeros(4, dtype=torch.int32), + query_start_loc_buf=query, + ngram_context_buf=torch.zeros(2, 1, dtype=torch.int32), + ) + } + return runner + + +@pytest.mark.parametrize( + ("ple_request", "query_start_loc", "match"), + [ + (PleOffloadRequest(dp_rank=0, num_tokens=3, num_reqs=2), [0, 2, 1], + "query_start_loc"), + (PleOffloadRequest(dp_rank=0, num_tokens=5, num_reqs=1), [0, 1], + "num_tokens"), + (PleOffloadRequest(dp_rank=0, num_tokens=1, num_reqs=3), [0, 1, 1], + "num_reqs"), + (PleOffloadRequest(dp_rank=9, num_tokens=1, num_reqs=1), [0, 1], + "dp_rank=9"), + ], +) +def test_invalid_request_fails_before_publication( + ple_request: PleOffloadRequest, + query_start_loc: list[int], + match: str, +) -> None: + runner = _runner(query_start_loc) + flags = [ + target.host_flag + for layers in runner._worker_targets.values() + for targets in layers.values() + for target in targets + ] + with pytest.raises(RuntimeError, match=match): + runner._validate_requests_before_publication([ple_request]) + assert all(int(flag[0]) == 0 for flag in flags) + + +class _SelectedGraphManager: + def ple_copy_rows(self, desc: BatchExecutionDescriptor) -> int: + return 12 + + +class _PublicationGuard: + called = False + + def prepare_forward(self, *args: Any, **kwargs: Any) -> None: + self.called = True + raise AssertionError("must fail before connector publication") + + +def test_selected_full_graph_copy_count_precedes_publication() -> None: + runner = GPUModelRunner.__new__(GPUModelRunner) + connector = _PublicationGuard() + runner._ple_offload_connector = connector + runner.cudagraph_manager = _SelectedGraphManager() + input_batch = SimpleNamespace( + num_reqs=1, + num_tokens_after_padding=8, + input_ids=torch.empty(8, dtype=torch.int32), + ) + desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, + num_tokens=8, + num_reqs=1, + ) + + with pytest.raises(RuntimeError, match="selected-forward"): + GPUModelRunner._prepare_ple_forward( + runner, input_batch, desc, dummy_run=False + ) + assert connector.called is False + + +def test_production_capture_callback_rejects_actual_ple_copy_mismatch() -> None: + manager = CudaGraphManager.__new__(CudaGraphManager) + manager._ple_full_copy_geometry = {} + manager._expected_ple_copy_layers = ("ple",) + manager._allow_ple_full_graph = True + desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, + num_tokens=8, + num_reqs=1, + ) + + with pytest.raises(RuntimeError, match="actual PLE copy geometry"): + # This is the production capture scope. The injected observation uses + # the same hook called by the real wait-copy custom-op implementation. + with manager.capture_ple_copy_geometry(desc): + _record_ple_copy_geometry("ple", 12, 16) + assert desc not in manager._ple_full_copy_geometry + + +def test_speculator_full_capture_refuses_observed_ple_op() -> None: + manager = SpeculatorCudaGraphManager.__new__(SpeculatorCudaGraphManager) + manager._ple_full_copy_geometry = {} + manager._expected_ple_copy_layers = () + manager._allow_ple_full_graph = False + desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, + num_tokens=8, + num_reqs=1, + ) + + with pytest.raises(RuntimeError, match="speculator/draft"): + with manager.capture_ple_copy_geometry(desc): + _record_ple_copy_geometry("unexpected-draft-ple", 8, 16) + assert desc not in manager._ple_full_copy_geometry + + +def test_eager_correspondence_uses_actual_padded_input_view() -> None: + runner = GPUModelRunner.__new__(GPUModelRunner) + connector = _PublicationGuard() + runner._ple_offload_connector = connector + runner.cudagraph_manager = None + input_batch = SimpleNamespace( + num_reqs=1, + num_tokens_after_padding=8, + input_ids=torch.empty(12, dtype=torch.int32), + ) + desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=8, + num_reqs=1, + ) + with pytest.raises(RuntimeError, match="selected-forward"): + GPUModelRunner._prepare_ple_forward( + runner, input_batch, desc, dummy_run=False + ) + assert connector.called is False diff --git a/tests/v1/ple_offload/test_fatal_rpc.py b/tests/v1/ple_offload/test_fatal_rpc.py new file mode 100644 index 0000000..1c86f9b --- /dev/null +++ b/tests/v1/ple_offload/test_fatal_rpc.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 + +import multiprocessing +from typing import Any + +import pytest + +from vllm.v1.ple_offload.connector import PleOffloadConnector +from vllm.v1.ple_offload.fatal import ( + fail_ple_execute_model, + ple_fail_closed, +) + + +class _Progress: + sequence = 7 + + def begin(self, **_: Any) -> int: + return self.sequence + + def complete(self) -> None: + return + + +class _RaisingProgress(_Progress): + def begin(self, **_: Any) -> int: + raise RuntimeError("begin injection") + + +class _InjectedConnector: + rank = 3 + tp_rank = 1 + _layers = {"ple": object()} + _progress = _Progress() + _input_ids_buf = None + _query_start_loc_buf = None + + def __init__(self, where: str) -> None: + self.where = where + if where == "begin": + self._progress = _RaisingProgress() + + def _fail_closed(self, where: str, error: BaseException) -> None: + ple_fail_closed( + where=where, + rank=self.rank, + layer="ple", + sequence=self._progress.sequence, + error=error, + ) + + def _prepare_forward_impl(self, *args: Any, **kwargs: Any) -> None: + if self.where == "prepare": + raise RuntimeError("prepare injection") + + def _release_outputs_impl(self) -> None: + if self.where == "release": + raise RuntimeError("release injection") + + +class _Runner: + def __init__(self, where: str) -> None: + self._ple_offload_connector = _InjectedConnector(where) + self.where = where + + @fail_ple_execute_model + def execute_model(self) -> None: + connector = self._ple_offload_connector + PleOffloadConnector.prepare_forward( + connector, + 1, + 1, + False, + consumer_num_tokens=1, + ) + if self.where == "forward": + raise RuntimeError("forward injection") + PleOffloadConnector.release_outputs(connector) + + +def _rpc_loop(where: str, successor_count: Any) -> None: + runner = _Runner(where) + try: + runner.execute_model() + except Exception: + # Mirrors WorkerProc._execute_worker_rpc: this must never regain control. + pass + with successor_count.get_lock(): + successor_count.value += 1 + + +@pytest.mark.parametrize("where", ["prepare", "forward", "release"]) +def test_ple_exception_kills_workerproc_style_rpc( + where: str, +) -> None: + context = multiprocessing.get_context("spawn") + successor_count = context.Value("i", 0) + process = context.Process( + target=_rpc_loop, + args=(where, successor_count), + ) + process.start() + # Spawned children import vLLM cold (measured 4-8 s on the reference host); the bound is + # about start-up, not the fatal path, which exits immediately. + process.join(timeout=120) + assert not process.is_alive() + assert process.exitcode not in (None, 0) + assert successor_count.value == 0 + + +def _rpc_loop_connector_only(where: str, successor_count: Any) -> None: + # No execute_model decorator: only the connector's own fatal wrappers can + # stop the RPC loop from regaining control. (The decorated test above + # passes on the decorator alone, so it cannot vouch for these wrappers.) + connector = _InjectedConnector(where) + try: + if where == "begin": + PleOffloadConnector.begin_forward( + connector, dummy_run=False, num_tokens=1 + ) + elif where == "prepare": + PleOffloadConnector.prepare_forward( + connector, 1, 1, False, consumer_num_tokens=1 + ) + else: + PleOffloadConnector.release_outputs(connector) + except Exception: + pass + with successor_count.get_lock(): + successor_count.value += 1 + + +@pytest.mark.parametrize("where", ["begin", "prepare", "release"]) +def test_connector_wrappers_are_fatal_without_decorator(where: str) -> None: + context = multiprocessing.get_context("spawn") + successor_count = context.Value("i", 0) + process = context.Process( + target=_rpc_loop_connector_only, + args=(where, successor_count), + ) + process.start() + process.join(timeout=120) + assert not process.is_alive() + assert process.exitcode not in (None, 0) + assert successor_count.value == 0 diff --git a/tests/v1/ple_offload/test_fault_bootstrap_disarmed.py b/tests/v1/ple_offload/test_fault_bootstrap_disarmed.py new file mode 100644 index 0000000..900851a --- /dev/null +++ b/tests/v1/ple_offload/test_fault_bootstrap_disarmed.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Disarmed-path check: a disarmed PLE child writes nothing but the real READY.""" + +import os +import subprocess +import sys +from multiprocessing.connection import Connection +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + + +def test_partial_ready_writers_are_inert_unless_armed_for_that_cut(monkeypatch): + import vllm_ple_offload_bootstrap as bootstrap + + sender = Mock(return_value=(8, 16)) + monkeypatch.setattr(bootstrap, "_FAULT", None) + for cut in ("launch_partial_ready", "launch_partial_ready_header"): + assert bootstrap._write_and_ack_partial(99, cut, sender) is False + sender.assert_not_called() + + other = SimpleNamespace(cut="exc_before_prepare", fires=Mock(return_value=True)) + monkeypatch.setattr(bootstrap, "_FAULT", other) + assert bootstrap._write_and_ack_partial(99, "launch_partial_ready", sender) is False + sender.assert_not_called() + other.fires.assert_not_called() + + +# The launch payload's config field is the protocol-0 encoding of None (b"N."). +_CHILD = r""" +import os, sys, types +write_fd = int(sys.argv[1]) +os.environ.pop("VLLM_PLE_TEST_FAULT", None) +import vllm_ple_offload_bootstrap as bootstrap +assert bootstrap._FAULT is None +bootstrap._recv_launch_payload = lambda fd: { + "vllm_config_payload": b"N.", + "num_workers": 1, + "ipc_addr": "ipc:///unused", +} + +def proc_main(*, vllm_config, num_workers, ipc_addr, ready_pipe): + assert vllm_config is None + ready_pipe.send({"status": "READY", "layer_names": ["ple"]}) + os._exit(0) + +fake = types.ModuleType("vllm.v1.ple_offload.worker") +fake.PleOffloadWorker = types.SimpleNamespace(proc_main=proc_main) +sys.modules["vllm.v1.ple_offload.worker"] = fake +death_read, death_write = os.pipe() +sys.argv = ["bootstrap", str(write_fd), str(death_read)] +bootstrap.main() +""" + + +@pytest.mark.skipif(sys.platform == "win32", reason="inherited fds are POSIX-only") +def test_disarmed_bootstrap_control_stream_is_exactly_one_ready_frame(): + read_fd, write_fd = os.pipe() + env = dict(os.environ) + env.pop("VLLM_PLE_TEST_FAULT", None) + proc = subprocess.Popen( + [sys.executable, "-c", _CHILD, str(write_fd)], + pass_fds=(write_fd,), + env=env, + ) + os.close(write_fd) + reader = Connection(read_fd) + try: + assert proc.wait(timeout=120) == 0 + assert reader.poll(5), "no READY frame on the control stream" + assert reader.recv() == {"status": "READY", "layer_names": ["ple"]} + # Nothing may follow (or precede, which would have broken the frame). + with pytest.raises(EOFError): + reader.recv_bytes() + finally: + reader.close() + if proc.poll() is None: + proc.kill() + proc.wait(timeout=120) diff --git a/tests/v1/ple_offload/test_fault_injector_basics.py b/tests/v1/ple_offload/test_fault_injector_basics.py new file mode 100644 index 0000000..0a9fd23 --- /dev/null +++ b/tests/v1/ple_offload/test_fault_injector_basics.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import queue +import subprocess +import sys +import threading +from collections import deque +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +import vllm_ple_test_fault as test_fault + + +def _spec(monkeypatch, tmp_path, value): + monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log")) + spec = test_fault._parse_fault_spec(value) + assert spec is not None + return spec + + +@pytest.mark.parametrize( + "value", + [ + "cut=park_before_publish,ordinal=1,tp=0", + "cut=launch_silent_ready,ordinal=1,pp=0", + "cut=drop_request,ordinal=1,tp=1,pp=0", + "cut=exc_after_prepare,ordinal=1,tp=0,pp=1", + ], +) +def test_impossible_role_selectors_are_rejected(monkeypatch, tmp_path, value): + monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log")) + with pytest.raises(ValueError): + test_fault._parse_fault_spec(value) + + +def test_armed_import_requires_dedicated_log(tmp_path): + env = dict(os.environ) + env["VLLM_PLE_TEST_FAULT"] = ( + "cut=exc_before_prepare,ordinal=1,rank=any" + ) + env.pop("VLLM_PLE_TEST_FAULT_LOG", None) + proc = subprocess.run( + [sys.executable, "-c", "import vllm_ple_test_fault"], + env=env, + text=True, + capture_output=True, + timeout=120, + ) + assert proc.returncode != 0 + assert "VLLM_PLE_TEST_FAULT_LOG is required" in proc.stderr + + +def test_failed_marker_write_suppresses_action(monkeypatch, tmp_path): + spec = _spec( + monkeypatch, + tmp_path, + "cut=exc_after_prepare,ordinal=1,tp=0,pp=0", + ) + monkeypatch.setattr(test_fault.os, "write", Mock(side_effect=BlockingIOError)) + assert not spec.fires( + "exc_after_prepare", + role="gpu", + unit="runner-forward", + eligible=True, + tp_rank=0, + pp_rank=0, + dp_rank=0, + transaction_id="sequence:1", + ) + assert spec._state.diagnostic_failed + assert not spec._state.fired + + +@pytest.mark.parametrize("order", ["bootstrap-first", "package-first"]) +@pytest.mark.parametrize("armed", [False, True]) +def test_fresh_interpreter_import_identity(tmp_path, order, armed): + env = dict(os.environ) + env.pop("VLLM_PLE_TEST_FAULT", None) + env.pop("VLLM_PLE_TEST_FAULT_LOG", None) + if armed: + env["VLLM_PLE_TEST_FAULT"] = ( + "cut=launch_silent_ready,ordinal=1,rank=any" + ) + env["VLLM_PLE_TEST_FAULT_LOG"] = str(tmp_path / f"{order}.log") + script = f""" +import sys +if {order!r} == "bootstrap-first": + import vllm_ple_offload_bootstrap as bootstrap + absent_before_package = "vllm_ple_test_fault" not in sys.modules + import vllm_ple_test_fault as package +else: + import vllm_ple_test_fault as package + import vllm_ple_offload_bootstrap as bootstrap + absent_before_package = False +assert bootstrap._FAULT is package._FAULT +assert getattr(sys.modules["vllm_ple_test_fault"], "_FAULT") is package._FAULT +if {armed!r}: + assert package._FAULT is not None +else: + assert package._FAULT is None + if {order!r} == "bootstrap-first": + assert absent_before_package +""" + subprocess.run( + [sys.executable, "-c", script], + env=env, + check=True, + timeout=120, + ) + + +class _IncompleteEvent: + def query(self): + return False + + +class _FatalObserved(BaseException): + pass + + +@pytest.mark.parametrize( + ("record_completion", "complete_submission", "expected"), + [ + (False, False, "host-submission"), + (True, False, "gpu-completion"), + (True, True, "gpu-completion"), + ], +) +def test_real_progress_monitor_observed_state( + monkeypatch, + record_completion, + complete_submission, + expected, +): + import vllm.v1.ple_offload.progress as progress_module + + observed = [] + reached = threading.Event() + + def fatal(**kwargs): + observed.append(kwargs) + reached.set() + raise _FatalObserved + + monkeypatch.setattr(progress_module, "ple_fail_closed", fatal) + monitor = progress_module.PleProgressMonitor( + rank=0, + layer_names=("ple",), + serving_deadline_s=0.02, + startup_deadline_s=0.02, + ) + monitor.begin(startup=False, num_tokens=1) + if record_completion: + monitor.track_completion(_IncompleteEvent()) + if complete_submission: + monitor.complete() + assert reached.wait(2) + assert expected in observed[0]["where"] + monitor.close() + + +def test_drop_retains_event_and_next_pool_one_launch_fails_closed( + monkeypatch, + tmp_path, +): + import vllm.v1.ple_offload.connector as connector_module + + spec = _spec( + monkeypatch, + tmp_path, + "cut=drop_request,ordinal=1,tp=0,pp=0", + ) + monkeypatch.setattr(connector_module, "_FAULT", spec) + + connector = object.__new__(connector_module.PleOffloadConnector) + connector.tp_rank = 0 + connector.pp_rank = 0 + connector.dp_rank = 0 + connector._progress = SimpleNamespace(sequence=7) + connector._dropped_events = [] + connector._d2h_event_pool = queue.Queue(maxsize=1) + event = Mock() + connector._d2h_event_pool.put_nowait(event) + connector._request_queue = Mock() + connector._enqueue_cuda_inputs = Mock() + + connector._launch( + num_reqs=1, + num_tokens=2, + serving=True, + transaction_id="sequence:7@begin:1.000000000", + correlation_id="sequence:7", + ) + assert connector._dropped_events == [event] + connector._request_queue.put_nowait.assert_not_called() + + with pytest.raises( + RuntimeError, + match="more requests than configured concurrent batches", + ): + connector._launch( + num_reqs=1, + num_tokens=2, + serving=True, + transaction_id="sequence:8@begin:2.000000000", + correlation_id="sequence:8", + ) + + +def test_async_output_cut_uses_real_sample_boundary(monkeypatch, tmp_path): + import vllm.v1.worker.gpu.model_runner as model_runner_module + + spec = _spec( + monkeypatch, + tmp_path, + "cut=exc_async_output,ordinal=1,tp=0,pp=0", + ) + monkeypatch.setattr(model_runner_module, "_FAULT", spec) + + class FatalReached(Exception): + pass + + connector = SimpleNamespace( + tp_rank=0, + pp_rank=0, + dp_rank=0, + completion_is_pending=Mock(return_value=True), + _fail_closed=Mock(side_effect=FatalReached), + ) + runner = object.__new__(model_runner_module.GPUModelRunner) + runner._ple_offload_connector = connector + runner.is_last_pp_rank = False + runner.pp_handler = Mock() + runner.execute_model_state = SimpleNamespace( + input_batch=object(), + attn_metadata=None, + slot_mappings_by_layer=None, + hidden_states=None, + aux_hidden_states=None, + dp_sync=None, + finished_req_ids=[], + ec_connector_output=None, + routed_experts=None, + cudagraph_stats=None, + ple_serving=True, + ple_sequence=11, + ple_transaction="sequence:11@begin:1.000000000", + ) + + with pytest.raises(FatalReached): + runner.sample_tokens(None) + assert connector._fail_closed.call_args.args[0] == "sample_tokens" + error = connector._fail_closed.call_args.args[1] + assert isinstance(error, RuntimeError) + assert "exc_async_output" in str(error) diff --git a/tests/v1/ple_offload/test_fault_injector_gates.py b/tests/v1/ple_offload/test_fault_injector_gates.py new file mode 100644 index 0000000..70e0d3e --- /dev/null +++ b/tests/v1/ple_offload/test_fault_injector_gates.py @@ -0,0 +1,468 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import ast +import contextlib +import importlib.util +import os +import pickle +import socket +import struct +import subprocess +import sys +import threading +import time +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +import vllm_ple_test_fault as test_fault + + +class _SequenceMustNotBeRead: + @property + def sequence(self): + raise AssertionError("disarmed path read progress.sequence") + + def track_completion(self, _event): + pass + + def complete(self): + pass + + +def _spec(monkeypatch, tmp_path, text): + monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log")) + value = test_fault._parse_fault_spec(text) + assert value is not None + return value + + +def test_disarmed_sender_deadline_does_not_enter_logger(monkeypatch): + import vllm.v1.ple_offload.worker as worker + + monkeypatch.setattr(worker, "_FAULT", None) + monkeypatch.setattr( + worker.logger, + "info", + Mock(side_effect=AssertionError("logger entered")), + ) + left, right = socket.socketpair() + try: + with pytest.raises(TimeoutError): + worker._send_framed_payload( + left.fileno(), + b"payload", + time.monotonic() - 1, + ) + worker.logger.info.assert_not_called() + finally: + left.close() + right.close() + + +def test_disarmed_launch_does_not_read_progress_sequence(monkeypatch): + import queue + import vllm.v1.ple_offload.connector as connector_module + + monkeypatch.setattr(connector_module, "_FAULT", None) + connector = object.__new__(connector_module.PleOffloadConnector) + connector.tp_rank = 0 + connector.pp_rank = 0 + connector.dp_rank = 0 + connector._progress = _SequenceMustNotBeRead() + connector._d2h_event_pool = queue.Queue() + connector._d2h_event_pool.put_nowait(Mock()) + connector._enqueue_cuda_inputs = Mock() + connector._request_queue = Mock() + + connector._launch( + 1, + 2, + serving=True, + transaction_id=None, + correlation_id=None, + ) + pending = connector._request_queue.put_nowait.call_args.args[0] + assert pending.transaction_id is None + assert pending.correlation_id is None + + +def test_disarmed_release_does_not_read_progress_sequence(monkeypatch): + import vllm.v1.ple_offload.connector as connector_module + + monkeypatch.setattr(connector_module, "_FAULT", None) + monkeypatch.setattr(connector_module._rowdump, "enabled", lambda: False) + event = Mock() + stream = Mock() + monkeypatch.setattr(connector_module.torch.cuda, "Event", lambda: event) + monkeypatch.setattr( + connector_module.torch.cuda, + "current_stream", + lambda _device: stream, + ) + + connector = object.__new__(connector_module.PleOffloadConnector) + connector.tp_rank = 0 + connector.pp_rank = 0 + connector.dp_rank = 0 + connector.device = SimpleNamespace(index=0) + connector._layers = {} + connector._after_dummy = False + connector._progress = _SequenceMustNotBeRead() + + connector._release_outputs_impl(serving=False) + event.record.assert_called_once_with(stream) + + +def test_sample_tokens_connector_rank_fatalizes(): + from vllm.v1.ple_offload.fatal import fail_ple_execute_model + + class Reached(Exception): + pass + + class Runner: + _ple_offload_connector = SimpleNamespace( + _fail_closed=Mock(side_effect=Reached) + ) + + @fail_ple_execute_model + def sample_tokens(self): + raise ValueError("sample failure") + + with pytest.raises(Reached): + Runner().sample_tokens() + + +def test_sample_tokens_non_connector_rank_rethrows(): + from vllm.v1.ple_offload.fatal import fail_ple_execute_model + + class Runner: + _ple_offload_connector = None + + @fail_ple_execute_model + def sample_tokens(self): + raise ValueError("sample failure") + + with pytest.raises(ValueError, match="sample failure"): + Runner().sample_tokens() + + +def test_profile_is_never_eligible_even_with_phase_any( + monkeypatch, + tmp_path, +): + spec = _spec( + monkeypatch, + tmp_path, + "cut=exc_after_prepare,ordinal=1,tp=0,pp=0,phase=any", + ) + assert not spec.fires( + "exc_after_prepare", + role="gpu", + unit="runner-forward", + eligible=False, + tp_rank=0, + pp_rank=0, + dp_rank=0, + transaction_id="sequence:1@begin:1.000000000", + correlation_id="sequence:1", + ) + assert spec._state.count == 0 + + +def test_offload_counts_only_serving_and_preserves_correlation( + monkeypatch, + tmp_path, +): + import vllm.v1.ple_offload.worker as worker + + spec = _spec( + monkeypatch, + tmp_path, + "cut=park_before_publish,ordinal=1,rank=any", + ) + monkeypatch.setattr(worker, "_FAULT", spec) + runner = object.__new__(worker.PleOffloadRunner) + runner._validate_requests_before_publication = Mock() + runner._worker_targets = {0: {}} + runner._input_bufs = {} + runner._scratch_bufs = {} + runner._layers = {} + runner._clamp_input_ids = False + runner._step = 0 + runner._request_index = 0 + + request = worker.PleOffloadRequest(dp_rank=0, num_tokens=2, num_reqs=1) + for _ in range(3): + runner._handle_requests_impl( + [worker._ReceivedPleOffloadRequest(request, False, None, None)] + ) + assert runner._request_index == 0 + + reached = threading.Event() + monkeypatch.setattr( + worker.threading, + "Event", + lambda: SimpleNamespace(wait=lambda: reached.set()), + ) + runner._handle_requests_impl( + [ + worker._ReceivedPleOffloadRequest( + request, + True, + "sequence:17@begin:4.000000000", + "sequence:17", + ) + ] + ) + assert reached.is_set() + line = (tmp_path / "fault.log").read_text() + assert "transaction=request:1@sequence:17@begin:4.000000000" in line + assert "correlation=sequence:17" in line + + +def test_cut6_requires_same_incomplete_completion(monkeypatch): + import vllm.v1.ple_offload.progress as progress + + class Event: + complete = False + + def query(self): + return self.complete + + monitor = progress.PleProgressMonitor( + rank=0, + layer_names=("ple",), + serving_deadline_s=60, + startup_deadline_s=60, + ) + event = Event() + try: + sequence = monitor.begin(startup=False, num_tokens=1) + monitor.track_completion(event) + monitor.complete() + assert monitor.completion_is_pending(sequence) + event.complete = True + assert not monitor.completion_is_pending(sequence) + assert not monitor.completion_is_pending(sequence + 1) + finally: + monitor.close() + + +def test_partial_ack_requires_complete_required_prefix(monkeypatch): + import vllm_ple_offload_bootstrap as bootstrap + + fault = SimpleNamespace( + cut="launch_partial_ready_header", + fires=Mock(return_value=True), + ) + monkeypatch.setattr(bootstrap, "_FAULT", fault) + assert not bootstrap._write_and_ack_partial( + 1, + "launch_partial_ready_header", + lambda _fd: (1, 4), + ) + fault.fires.assert_not_called() + + assert bootstrap._write_and_ack_partial( + 1, + "launch_partial_ready_header", + lambda _fd: (2, 4), + ) + assert fault.fires.call_args.kwargs["bytes_written"] == 2 + assert fault.fires.call_args.kwargs["frame_size"] == 4 + + +@pytest.mark.timeout(180) +@pytest.mark.skipif(sys.platform == "win32", reason="requires POSIX pass_fds") +@pytest.mark.parametrize( + ("cut", "expected_transaction"), + [ + ("launch_block_reconstruction", "reducer:configuration"), + ( + "launch_block_worker_import", + "module:vllm.v1.ple_offload.worker", + ), + ], +) +def test_gate_is_inside_reducer_or_loader_and_parent_death_wins( + tmp_path, + cut, + expected_transaction, +): + gate = tmp_path / "gate" + os.mkfifo(gate) + gate_keeper = os.open(gate, os.O_RDWR | os.O_NONBLOCK) + control_parent, control_child = socket.socketpair() + death_parent, death_child = socket.socketpair() + fault_log = tmp_path / "fault.log" + env = dict(os.environ) + env["VLLM_PLE_TEST_FAULT"] = f"cut={cut},ordinal=1,rank=any" + env["VLLM_PLE_TEST_FAULT_LOG"] = str(fault_log) + env["VLLM_PLE_TEST_FAULT_GATE"] = str(gate) + + config_payload = pickle.dumps(None, protocol=pickle.HIGHEST_PROTOCOL) + if cut == "launch_block_reconstruction": + config_payload = test_fault._FaultReconstructionPayload(config_payload) + launch = pickle.dumps( + { + "vllm_config_payload": config_payload, + "num_workers": 1, + "ipc_addr": "ipc:///unused-fault", + }, + protocol=pickle.HIGHEST_PROTOCOL, + ) + + process = subprocess.Popen( + [ + sys.executable, + "-m", + "vllm_ple_offload_bootstrap", + str(control_child.fileno()), + str(death_child.fileno()), + ], + env=env, + pass_fds=(control_child.fileno(), death_child.fileno()), + ) + control_child.close() + death_child.close() + try: + control_parent.sendall(struct.pack("!Q", len(launch)) + launch) + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + if fault_log.exists() and expected_transaction in fault_log.read_text(): + break + assert process.poll() is None + time.sleep(0.05) + else: + pytest.fail("operation-entry acknowledgement was not emitted") + + assert process.poll() is None + death_parent.close() + assert process.wait(timeout=120) != 0 + finally: + with contextlib.suppress(ProcessLookupError): + process.kill() + process.wait(timeout=120) + control_parent.close() + with contextlib.suppress(OSError): + death_parent.close() + os.close(gate_keeper) + + +def _load_gate_helper(): + root = Path(__file__).resolve().parents[3] + path = root / "probe" / "faultcut_gate.py" + spec = importlib.util.spec_from_file_location("faultcut_gate", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize( + "change", + [ + {"alive": False}, + {"observed_start": 88}, + {"bytes_sent": 0}, + {"bytes_sent": 100}, + {"completed": True}, + {"child_not_reading": False}, + ], +) +def test_external_partial_send_gate_rejections(change): + gate = _load_gate_helper() + values = { + "pid": 12, + "expected_start": 77, + "observed_start": 77, + "alive": True, + "bytes_sent": 4, + "frame_size": 100, + "completed": False, + "child_not_reading": True, + } + values.update(change) + ok, _reason = gate.validate_partial_send_gate(**values) + assert not ok + + +def test_external_gate_rejects_failed_signal(): + gate = _load_gate_helper() + ok, reason = gate.deliver_checked_signal( + pid=12, + expected_start=77, + read_start=lambda _pid: 77, + alive=lambda _pid: True, + send=lambda _pid, _sig: False, + signal_number=9, + ) + assert not ok + assert "delivery failed" in reason + + +def test_every_fault_call_has_one_boolop_guard_and_no_prior_state_read(): + root = Path(__file__).resolve().parents[3] + expected = { + "vllm/v1/ple_offload/connector.py": 3, + "vllm/v1/ple_offload/worker.py": 1, + "vllm/v1/worker/gpu/model_runner.py": 4, + "vllm_ple_offload_bootstrap/__init__.py": 1, + } + + for relative, wanted in expected.items(): + tree = ast.parse((root / relative).read_text(encoding="utf-8")) + parents = {} + for parent in ast.walk(tree): + for child in ast.iter_child_nodes(parent): + parents[child] = parent + + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "fires" + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "_FAULT" + ] + assert len(calls) == wanted, relative + for call in calls: + owner = call + while owner in parents and not isinstance(owner, ast.BoolOp): + owner = parents[owner] + assert isinstance(owner, ast.BoolOp), (relative, call.lineno) + assert isinstance(owner.op, ast.And), (relative, call.lineno) + + guard_index = None + call_index = None + for index, value in enumerate(owner.values): + if ( + isinstance(value, ast.Compare) + and isinstance(value.left, ast.Name) + and value.left.id == "_FAULT" + and len(value.ops) == 1 + and isinstance(value.ops[0], ast.IsNot) + and len(value.comparators) == 1 + and isinstance(value.comparators[0], ast.Constant) + and value.comparators[0].value is None + ): + guard_index = index + if call in ast.walk(value): + call_index = index + assert guard_index is not None + assert call_index is not None + assert guard_index < call_index + + forbidden = {"sequence", "has_pending_completions", + "completion_is_pending", "fault_transaction"} + before_guard = owner.values[:guard_index] + assert not any( + isinstance(node, ast.Attribute) and node.attr in forbidden + for value in before_guard + for node in ast.walk(value) + ) diff --git a/tests/v1/ple_offload/test_fault_wire_frames.py b/tests/v1/ple_offload/test_fault_wire_frames.py new file mode 100644 index 0000000..b864355 --- /dev/null +++ b/tests/v1/ple_offload/test_fault_wire_frames.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Disarmed-path checks: the disarmed PLE request wire format stays single-frame.""" + +import contextlib +import queue +from types import SimpleNamespace +from unittest.mock import Mock + +import msgspec +import pytest + + +def _request(): + from vllm.v1.ple_offload.protocol import PleOffloadRequest + + return PleOffloadRequest(dp_rank=0, num_tokens=2, num_reqs=1) + + +def test_disarmed_offload_receive_uses_single_frame(monkeypatch): + import vllm.v1.ple_offload.worker as worker_module + + monkeypatch.setattr(worker_module, "_FAULT", None) + request = _request() + socket = Mock() + socket.recv.return_value = msgspec.msgpack.encode(request) + socket.recv_multipart.side_effect = AssertionError("multipart read while disarmed") + + assert worker_module._receive_offload_request(socket) == request + socket.recv.assert_called_once_with(0) + + +def test_armed_offload_receive_requires_metadata_frame(monkeypatch): + import vllm.v1.ple_offload.worker as worker_module + + monkeypatch.setattr(worker_module, "_FAULT", object()) + request = _request() + metadata = {"serving": True, "transaction": "request:1", "correlation": "sequence:3"} + socket = Mock() + socket.recv.side_effect = AssertionError("single-frame read while armed") + socket.recv_multipart.return_value = [ + msgspec.msgpack.encode(request), + msgspec.msgpack.encode(metadata), + ] + + received = worker_module._receive_offload_request(socket, 1) + assert received.request == request + assert received.serving is True + assert received.transaction_id == "request:1" + assert received.correlation_id == "sequence:3" + socket.recv_multipart.assert_called_once_with(1) + + socket.recv_multipart.return_value = [msgspec.msgpack.encode(request)] + with pytest.raises(RuntimeError): + worker_module._receive_offload_request(socket) + + +@pytest.mark.parametrize("armed", [False, True]) +def test_connector_publication_frames_follow_arming(monkeypatch, armed): + import vllm.v1.ple_offload.connector as connector_module + + fault = SimpleNamespace(fires=lambda *args, **kwargs: False) if armed else None + monkeypatch.setattr(connector_module, "_FAULT", fault) + monkeypatch.setattr( + connector_module.torch.accelerator, + "device_index", + lambda _index: contextlib.nullcontext(), + ) + connector = object.__new__(connector_module.PleOffloadConnector) + connector.tp_rank = 0 + connector.pp_rank = 0 + connector.dp_rank = 0 + connector.device = SimpleNamespace(index=0) + connector._d2h_event_pool = queue.Queue() + event = Mock() + request = _request() + pending = connector_module._PendingPleOffloadRequest( + request, event, True, "sequence:4@begin:1.000000000", "sequence:4" + ) + socket = Mock() + + connector._process_request(pending, socket) + + event.synchronize.assert_called_once() + if armed: + socket.send.assert_not_called() + payload, metadata = socket.send_multipart.call_args.args[0] + assert msgspec.msgpack.decode(metadata) == { + "serving": True, + "transaction": "sequence:4@begin:1.000000000", + "correlation": "sequence:4", + } + else: + socket.send_multipart.assert_not_called() + (payload,) = socket.send.call_args.args + assert payload == msgspec.msgpack.encode(request) + assert connector._d2h_event_pool.get_nowait() is event diff --git a/tests/v1/ple_offload/test_ipc_socket_unlink.py b/tests/v1/ple_offload/test_ipc_socket_unlink.py new file mode 100644 index 0000000..92fbc42 --- /dev/null +++ b/tests/v1/ple_offload/test_ipc_socket_unlink.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 + +import contextlib +import os +from pathlib import Path +import signal +import subprocess +import sys +import tempfile +import time + +import pytest + +from vllm.v1.ple_offload.worker import PleOffloadWorkerHandle + + +pytestmark = pytest.mark.skipif( + sys.platform != "linux", reason="requires Linux filesystem ipc sockets" +) + + +_STUB_CHILD = r""" +import pathlib +import signal +import sys +import threading + +import zmq + +from vllm.v1.ple_offload.worker import _unlink_ipc_socket + +ipc_addr, mode, marker = sys.argv[1:] +stopping = threading.Event() +if mode == "graceful": + signal.signal(signal.SIGTERM, lambda _signum, _frame: stopping.set()) +else: + signal.signal(signal.SIGTERM, signal.SIG_IGN) + +context = zmq.Context() +socket = context.socket(zmq.PULL) +socket.bind(ipc_addr) +try: + stopping.wait() +finally: + socket.close(linger=0) + context.term() + if mode == "graceful": + _unlink_ipc_socket(ipc_addr) + pathlib.Path(marker).touch() +""" + + +def _start_stub(ipc_addr: str, mode: str, marker: Path) -> subprocess.Popen: + return subprocess.Popen( + [sys.executable, "-c", _STUB_CHILD, ipc_addr, mode, str(marker)], + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + ) + + +def _wait_for_socket(proc: subprocess.Popen, path: Path) -> None: + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + if path.exists(): + return + if proc.poll() is not None: + stderr = proc.stderr.read() if proc.stderr is not None else "" + pytest.fail(f"stub child exited before bind: {stderr}") + time.sleep(0.01) + pytest.fail("stub child did not bind its ipc socket") + + +def _cleanup_process(proc: subprocess.Popen, socket_path: Path) -> None: + if proc.poll() is None: + proc.kill() + proc.wait(timeout=5) + with contextlib.suppress(FileNotFoundError): + os.unlink(socket_path) + if proc.stderr is not None: + proc.stderr.close() + + +def test_handle_graceful_close_removes_child_bound_socket() -> None: + with tempfile.TemporaryDirectory(prefix="ple-ipc-", dir="/tmp") as temp_dir: + socket_path = Path(temp_dir) / "owned.sock" + marker = Path(temp_dir) / "child-cleaned" + ipc_addr = f"ipc://{socket_path}" + proc = _start_stub(ipc_addr, "graceful", marker) + try: + handle = PleOffloadWorkerHandle( + proc=proc, + death_writer=None, + ready_pipe_reader=None, + launch_deadline=time.monotonic() + 30, + ipc_addr=ipc_addr, + ) + _wait_for_socket(proc, socket_path) + + handle.close() + + assert proc.returncode == 0 + assert marker.exists(), "the child ran its post-context unlink" + assert not socket_path.exists() + finally: + _cleanup_process(proc, socket_path) + + +def test_handle_close_removes_socket_after_child_sigkill() -> None: + with tempfile.TemporaryDirectory(prefix="ple-ipc-", dir="/tmp") as temp_dir: + socket_path = Path(temp_dir) / "owned.sock" + marker = Path(temp_dir) / "unused" + ipc_addr = f"ipc://{socket_path}" + proc = _start_stub(ipc_addr, "park", marker) + try: + handle = PleOffloadWorkerHandle( + proc=proc, + death_writer=None, + ready_pipe_reader=None, + launch_deadline=time.monotonic() + 30, + ipc_addr=ipc_addr, + ) + _wait_for_socket(proc, socket_path) + proc.kill() + proc.wait(timeout=5) + assert proc.returncode == -signal.SIGKILL + assert socket_path.exists() + + handle.close() + + assert not socket_path.exists() + finally: + _cleanup_process(proc, socket_path) diff --git a/tests/v1/ple_offload/test_pending_rpc.py b/tests/v1/ple_offload/test_pending_rpc.py new file mode 100644 index 0000000..ef3c0fd --- /dev/null +++ b/tests/v1/ple_offload/test_pending_rpc.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 + +import multiprocessing +from types import SimpleNamespace +from typing import Any + +import pytest + +import vllm.v1.executor.multiproc_executor as executor_module +from vllm.v1.executor.multiproc_executor import WorkerProc +from vllm.v1.ple_offload.fatal import ple_fail_closed + + +class _BrokenAsyncOutput: + def get_output(self) -> Any: + raise RuntimeError("injected asynchronous output failure") + + +class _Progress: + sequence = 9 + + def __init__(self, pending: bool) -> None: + self.has_pending_completions = pending + + +class _Connector: + rank = 2 + _layers = {"ple": object()} + + def __init__(self, pending: bool) -> None: + self._progress = _Progress(pending) + + @property + def has_pending_completions(self) -> bool: + return self._progress.has_pending_completions + + def _fail_closed(self, where: str, error: BaseException) -> None: + ple_fail_closed( + where=where, + rank=self.rank, + layer="ple", + sequence=self._progress.sequence, + error=error, + ) + + +def _async_output_failure_case(pending: bool) -> None: + connector = _Connector(pending) + worker = SimpleNamespace( + model_runner=SimpleNamespace(_ple_offload_connector=connector) + ) + proc = WorkerProc.__new__(WorkerProc) + proc.worker = SimpleNamespace(worker=worker) + proc.worker_response_mq = None + + # Exercise the real async-output exception recovery branch without needing + # to construct the production dataclass's CUDA-backed fields. + executor_module.AsyncModelRunnerOutput = _BrokenAsyncOutput + proc.enqueue_output(_BrokenAsyncOutput()) + + +@pytest.mark.parametrize( + ("pending", "must_die"), + [(True, True), (False, False)], +) +def test_async_output_failure_obeys_pending_ple_ownership( + pending: bool, + must_die: bool, +) -> None: + context = multiprocessing.get_context("spawn") + process = context.Process( + target=_async_output_failure_case, + args=(pending,), + ) + process.start() + process.join(timeout=120) + assert not process.is_alive() + if must_die: + assert process.exitcode not in (None, 0) + else: + assert process.exitcode == 0 diff --git a/tests/v1/ple_offload/test_progress_deadline.py b/tests/v1/ple_offload/test_progress_deadline.py new file mode 100644 index 0000000..5e7ca19 --- /dev/null +++ b/tests/v1/ple_offload/test_progress_deadline.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: Apache-2.0 + +import multiprocessing +import os +import threading +from pathlib import Path +from typing import Any + +import pytest +import torch + +import vllm.v1.ple_offload.connector as connector_module +from vllm.v1.ple_offload.connector import PleOffloadConnector +from vllm.v1.ple_offload.progress import PleProgressMonitor + + +class _PendingEvent: + def record(self, stream: Any) -> None: + stream.recorded = True + + def query(self) -> bool: + return False + + +class _BrokenEvent(_PendingEvent): + def __init__(self, query_entered: Any) -> None: + self.query_entered = query_entered + + def query(self) -> bool: + self.query_entered.set() + raise RuntimeError("injected event query failure") + + +class _FakeStream: + recorded = False + + def synchronize(self) -> None: + return + + +class _FakeLayer: + def release_offloaded_output(self, stream: Any) -> None: + stream.reset_written = True + + +def _silent_producer_case(block_input: bool, started: Any) -> None: + monitor = PleProgressMonitor( + rank=2, + layer_names=("ple",), + serving_deadline_s=0.2, + startup_deadline_s=1.0, + ) + monitor.begin(startup=False, num_tokens=32) + started.set() + + if block_input: + # Models the request thread stuck in d2h_done_event.synchronize(). + threading.Thread( + target=threading.Event().wait, + daemon=True, + ).start() + + # Neither the fake producer nor the test terminates the process. The + # independent progress monitor must do so. + threading.Event().wait() + + +@pytest.mark.parametrize("block_input", [False, True]) +def test_progress_deadline_terminates_alive_silent_worker( + block_input: bool, +) -> None: + context = multiprocessing.get_context("spawn") + started = context.Event() + process = context.Process( + target=_silent_producer_case, + args=(block_input, started), + ) + process.start() + # Cold vLLM import in the spawned child takes 4-8 s on the reference host; the deadline + # itself (0.2 s) is only charged after `started` is set. + assert started.wait(timeout=120) + process.join(timeout=10) + assert not process.is_alive() + assert process.exitcode not in (None, 0) + + +def _pending_completion_case(use_release_wrapper: bool, started: Any) -> None: + monitor = PleProgressMonitor( + rank=2, + layer_names=("ple",), + serving_deadline_s=0.2, + startup_deadline_s=1.0, + ) + if use_release_wrapper: + stream = _FakeStream() + connector_module.torch.cuda.current_stream = lambda device=None: stream + connector_module.torch.cuda.Event = _PendingEvent + connector_module.torch.cuda.is_current_stream_capturing = lambda: False + connector_module._rowdump.enabled = lambda: False + + connector = PleOffloadConnector.__new__(PleOffloadConnector) + connector.device = torch.device("cpu") + connector.rank = 2 + connector.tp_rank = 0 + connector._layers = {"ple": _FakeLayer()} + connector._progress = monitor + connector._rowdump_n = None + connector._after_dummy = False + monitor.begin(startup=False, num_tokens=8) + connector.release_outputs() + assert stream.recorded + assert stream.reset_written + else: + monitor.begin(startup=False, num_tokens=8) + monitor.track_completion(_PendingEvent()) + monitor.complete() + started.set() + threading.Event().wait() + + +@pytest.mark.parametrize("use_release_wrapper", [False, True]) +def test_pending_gpu_completion_expires_after_release_returns( + use_release_wrapper: bool, +) -> None: + context = multiprocessing.get_context("spawn") + started = context.Event() + process = context.Process( + target=_pending_completion_case, + args=(use_release_wrapper, started), + ) + process.start() + assert started.wait(timeout=120) + process.join(timeout=10) + assert not process.is_alive() + assert process.exitcode not in (None, 0) + + +@pytest.mark.parametrize("value", [float("inf"), float("nan"), -1.0, 1e12]) +def test_constructor_rejects_invalid_deadlines(value: float) -> None: + with pytest.raises(ValueError, match="finite"): + PleProgressMonitor( + rank=0, + layer_names=("ple",), + serving_deadline_s=value, + startup_deadline_s=1.0, + ) + with pytest.raises(ValueError, match="finite"): + PleProgressMonitor( + rank=0, + layer_names=("ple",), + serving_deadline_s=1.0, + startup_deadline_s=value, + ) + + +@pytest.mark.parametrize( + "name", + ["VLLM_PLE_PROGRESS_DEADLINE_S", "VLLM_PLE_STARTUP_DEADLINE_S"], +) +@pytest.mark.parametrize("value", ["inf", "nan", "-1", "1e12"]) +def test_environment_rejects_invalid_deadlines( + monkeypatch: pytest.MonkeyPatch, + name: str, + value: str, +) -> None: + monkeypatch.setenv("VLLM_PLE_PROGRESS_DEADLINE_S", "1") + monkeypatch.setenv("VLLM_PLE_STARTUP_DEADLINE_S", "1") + monkeypatch.setenv(name, value) + with pytest.raises(ValueError, match="finite"): + PleProgressMonitor.from_environment(rank=0, layer_names=("ple",)) + + +def _monitor_exception_case( + started: Any, + query_entered: Any, + diagnostic_path: str, +) -> None: + diagnostic_fd = os.open(diagnostic_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + os.dup2(diagnostic_fd, 2) + os.close(diagnostic_fd) + monitor = PleProgressMonitor( + rank=4, + layer_names=("ple",), + serving_deadline_s=10.0, + startup_deadline_s=10.0, + ) + monitor.begin(startup=False, num_tokens=8) + # Set before arming the failing event: the monitor thread's os._exit must not + # race a multiprocessing Event.set() that holds the shared condition lock. + started.set() + monitor.track_completion(_BrokenEvent(query_entered)) + monitor.complete() + threading.Event().wait() + + +def test_monitor_thread_exception_is_fatal(tmp_path: Path) -> None: + context = multiprocessing.get_context("spawn") + started = context.Event() + query_entered = context.Event() + diagnostic_path = tmp_path / "monitor-fatal.txt" + process = context.Process( + target=_monitor_exception_case, + args=(started, query_entered, str(diagnostic_path)), + ) + process.start() + assert started.wait(timeout=120) + assert query_entered.wait(timeout=10) + process.join(timeout=10) + assert not process.is_alive() + assert process.exitcode not in (None, 0) + diagnostic = diagnostic_path.read_text(encoding="utf-8") + assert "where=progress-monitor-thread" in diagnostic + assert ( + "error=RuntimeError:injected event query failure" + in diagnostic + ) diff --git a/tests/v1/ple_offload/test_shm_unlink.py b/tests/v1/ple_offload/test_shm_unlink.py new file mode 100644 index 0000000..600a7bb --- /dev/null +++ b/tests/v1/ple_offload/test_shm_unlink.py @@ -0,0 +1,397 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +from multiprocessing.reduction import ForkingPickler +from pathlib import Path +import subprocess +import sys + +import pytest +import torch +import torch.multiprocessing as torch_mp + +from vllm.v1.ple_offload import connector as connector_module +from vllm.v1.ple_offload.connector import PleOffloadConnector +from vllm.v1.ple_offload.protocol import PleOffloadRegistration + + +def _require_file_system_shm() -> None: + if ( + os.name != "posix" + or not os.path.isdir("/dev/shm") + or not os.access("/dev/shm", os.W_OK) + ): + pytest.skip("POSIX /dev/shm is unavailable") + if "file_system" not in torch_mp.get_all_sharing_strategies(): + pytest.skip("torch file_system sharing is unavailable") + + +def _shm_path(tensor: torch.Tensor) -> Path: + _, filename, _ = tensor.untyped_storage()._share_filename_cpu_() + return Path("/dev/shm") / os.fsdecode(filename).lstrip("/") + + +def _make_connector_storage() -> tuple[PleOffloadConnector, list[torch.Tensor]]: + original_strategy = torch_mp.get_sharing_strategy() + torch_mp.set_sharing_strategy("file_system") + try: + host_outputs = { + "layer.0": torch.arange(12, dtype=torch.float32) + .reshape(3, 4) + .share_memory_(), + "layer.1": torch.arange(8, dtype=torch.float32) + .reshape(2, 4) + .share_memory_(), + } + host_flags = { + name: torch.zeros(16, dtype=torch.int32).share_memory_() + for name in host_outputs + } + input_ids = torch.arange(16, dtype=torch.int32).share_memory_() + query_start_loc = torch.arange(5, dtype=torch.int32).share_memory_() + ngram_context = torch.arange(12, dtype=torch.int32).reshape(4, 3) + ngram_context.share_memory_() + + registration = PleOffloadRegistration( + worker_id=0, + tp_rank=0, + dp_rank=0, + host_output_buffers=host_outputs, + host_flag_tensors=host_flags, + input_ids_buf=input_ids, + query_start_loc_buf=query_start_loc, + ngram_context_buf=ngram_context, + ) + # Match producer serialization and worker reconstruction. loads() + # consumes the transfer reference added by dumps(), leaving storage + # destruction able to reach the allocator's zero-counter branch. + payload = ForkingPickler.dumps(registration) + rebuilt_registration = ForkingPickler.loads(payload) + assert isinstance(rebuilt_registration, PleOffloadRegistration) + del rebuilt_registration + finally: + torch_mp.set_sharing_strategy(original_strategy) + + connector = PleOffloadConnector.__new__(PleOffloadConnector) + connector._host_output_bufs = host_outputs + connector._host_flags = host_flags + connector._reserved_host_storage_filenames = set() + connector._input_ids_buf = input_ids + connector._query_start_loc_buf = query_start_loc + connector._ngram_context_buf = ngram_context + tensors = [ + *host_outputs.values(), + *host_flags.values(), + input_ids, + query_start_loc, + ngram_context, + ] + assert len(tensors) == 7 + return connector, tensors + + +def test_release_host_storage_unlinks_names_but_keeps_mappings() -> None: + _require_file_system_shm() + connector, tensors = _make_connector_storage() + paths = [_shm_path(tensor) for tensor in tensors] + before = [tensor.clone() for tensor in tensors] + assert len(set(paths)) == 7 + assert all(path.exists() for path in paths) + + connector.release_host_storage() + connector.release_host_storage() + + assert all(not path.exists() for path in paths) + for tensor, expected in zip(tensors, before): + assert torch.equal(tensor, expected) + tensor.add_(1) + assert torch.equal(tensor, expected + 1) + + +def test_without_release_helper_names_remain(monkeypatch) -> None: + _require_file_system_shm() + connector, tensors = _make_connector_storage() + paths = [_shm_path(tensor) for tensor in tensors] + + with monkeypatch.context() as patch: + patch.setattr( + connector_module, + "_unlink_filename_shared_storage", + lambda _tensor, _reservations: None, + ) + connector.release_host_storage() + + # Do not manually unlink these while their unreserved storages are live. + # Normal storage destruction removes them after the test frame is released. + assert all(path.exists() for path in paths) + + +def test_release_host_storage_does_not_convert_fd_shared_storage() -> None: + if "file_descriptor" not in torch_mp.get_all_sharing_strategies(): + pytest.skip("torch file_descriptor sharing is unavailable") + original_strategy = torch_mp.get_sharing_strategy() + torch_mp.set_sharing_strategy("file_descriptor") + try: + fd_tensor = torch.arange(8, dtype=torch.int32).share_memory_() + finally: + torch_mp.set_sharing_strategy(original_strategy) + + connector = PleOffloadConnector.__new__(PleOffloadConnector) + connector._host_output_bufs = {"layer.0": fd_tensor} + connector._host_flags = {} + connector._reserved_host_storage_filenames = set() + connector._input_ids_buf = torch.arange(2, dtype=torch.int32) + connector._query_start_loc_buf = torch.arange(2, dtype=torch.int32) + connector._ngram_context_buf = None + expected = fd_tensor.clone() + + connector.release_host_storage() + + assert fd_tensor.untyped_storage()._get_shared_fd() >= 0 + assert torch.equal(fd_tensor, expected) + + +class _FakeFilenameStorage: + def __init__(self, filename: bytes) -> None: + self.device = torch.device("cpu") + self.filename = filename + self.increfs = 0 + self.decrefs = 0 + + def is_shared(self) -> bool: + return True + + def _get_shared_fd(self) -> int: + raise RuntimeError("filename-backed storage has no shared fd") + + def _share_filename_cpu_(self) -> tuple[bytes, bytes, int]: + return b"manager", self.filename, 4 + + def _shared_incref(self) -> None: + self.increfs += 1 + + def _shared_decref(self) -> "_FakeFilenameStorage": + self.decrefs += 1 + return self + + +class _FakeTensor: + def __init__(self, storage: _FakeFilenameStorage) -> None: + self.storage = storage + + def untyped_storage(self) -> _FakeFilenameStorage: + return self.storage + + +def _make_fake_connector(tensor: _FakeTensor) -> PleOffloadConnector: + connector = PleOffloadConnector.__new__(PleOffloadConnector) + # Repeat the same storage in several connector positions to cover aliases. + connector._host_output_bufs = {"layer.0": tensor} + connector._host_flags = {"layer.0": tensor} + connector._reserved_host_storage_filenames = set() + connector._input_ids_buf = tensor + connector._query_start_loc_buf = tensor + connector._ngram_context_buf = None + return connector + + +def test_release_host_storage_reserves_each_name_once(monkeypatch) -> None: + storage = _FakeFilenameStorage(b"/ple-test-reservation") + connector = _make_fake_connector(_FakeTensor(storage)) + unlinked: list[str] = [] + monkeypatch.setattr(connector_module.os, "unlink", unlinked.append) + + connector.release_host_storage() + connector.release_host_storage() + + assert storage.increfs == 1 + assert storage.decrefs == 0 + assert unlinked == ["/dev/shm/ple-test-reservation"] + + +def test_release_host_storage_rolls_back_on_unlink_failure(monkeypatch) -> None: + storage = _FakeFilenameStorage(b"/ple-test-rollback") + connector = _make_fake_connector(_FakeTensor(storage)) + + def fail_unlink(_path: str) -> None: + raise PermissionError("test unlink failure") + + monkeypatch.setattr(connector_module.os, "unlink", fail_unlink) + + connector.release_host_storage() + + # The fake connector aliases one storage in four positions. A rolled-back + # reservation is retried by the next alias in the same pass (by design: the + # name is still present), so every reservation must be balanced by a + # rollback and nothing may stay reserved. + assert storage.increfs >= 1 + assert storage.increfs == storage.decrefs + assert connector._reserved_host_storage_filenames == set() + + +_BALANCED_LIFETIME_SCRIPT = r""" +import gc +import os +from multiprocessing.reduction import ForkingPickler +from pathlib import Path + +import torch +import torch.multiprocessing as torch_mp + +from vllm.v1.ple_offload.connector import ( + PleOffloadConnector, + _PROCESS_LIFETIME_REGISTERED_STORAGE, +) +from vllm.v1.ple_offload.protocol import PleOffloadRegistration + + +def shm_path(tensor): + _, filename, _ = tensor.untyped_storage()._share_filename_cpu_() + return Path("/dev/shm") / os.fsdecode(filename).lstrip("/") + + +class Progress: + def close(self): + pass + + +original_strategy = torch_mp.get_sharing_strategy() +torch_mp.set_sharing_strategy("file_system") +try: + host_outputs = { + "layer.0": torch.arange(12, dtype=torch.float32) + .reshape(3, 4) + .share_memory_(), + "layer.1": torch.arange(8, dtype=torch.float32) + .reshape(2, 4) + .share_memory_(), + } + host_flags = { + name: torch.zeros(16, dtype=torch.int32).share_memory_() + for name in host_outputs + } + input_ids = torch.arange(16, dtype=torch.int32).share_memory_() + query_start_loc = torch.arange(5, dtype=torch.int32).share_memory_() + ngram_context = torch.arange(12, dtype=torch.int32).reshape(4, 3) + ngram_context.share_memory_() + + registration = PleOffloadRegistration( + worker_id=1, + tp_rank=1, + dp_rank=0, + host_output_buffers=host_outputs, + host_flag_tensors=host_flags, + input_ids_buf=input_ids, + query_start_loc_buf=query_start_loc, + ngram_context_buf=ngram_context, + ) + payload = ForkingPickler.dumps(registration) + rebuilt_registration = ForkingPickler.loads(payload) +finally: + torch_mp.set_sharing_strategy(original_strategy) + +connector = PleOffloadConnector.__new__(PleOffloadConnector) +connector._host_output_bufs = host_outputs +connector._host_flags = host_flags +connector._reserved_host_storage_filenames = set() +connector._input_ids_buf = input_ids +connector._query_start_loc_buf = query_start_loc +connector._ngram_context_buf = ngram_context + +# Exercise the production close ordering without CUDA or transport setup. +connector._request_thread = None +connector._pinned_input_buffers = [] +connector._mapped_host_buffers = [] +connector._d2h_event_pool = object() +connector._registration_socket = None +connector._zmq_ctx = None +connector._progress = Progress() + +tensors = [ + *host_outputs.values(), + *host_flags.values(), + input_ids, + query_start_loc, + ngram_context, +] +assert len(tensors) == 7 +assert len({tensor.untyped_storage()._weak_ref() for tensor in tensors}) == 7 + +# TP != 0 creates and serializes its input buffers without CUDA registration +# or the process-lifetime reference used by TP0. +unpinned_tp_input = input_ids +assert not connector._pinned_input_buffers +assert all( + unpinned_tp_input is not registered + for registered in _PROCESS_LIFETIME_REGISTERED_STORAGE +) + +paths = [shm_path(tensor) for tensor in tensors] +before = [tensor.clone() for tensor in tensors] +assert all(path.exists() for path in paths) + +connector.close() +connector.close() + +assert all(not path.exists() for path in paths) +for tensor, expected in zip(tensors, before): + assert torch.equal(tensor, expected) + tensor.add_(1) + assert torch.equal(tensor, expected + 1) +del tensor, expected + +# Remove every original and reconstructed tensor reference. With an unreserved +# allocator this reaches counter zero after manual unlink and aborts the child. +connector._host_output_bufs = {} +connector._host_flags = {} +connector._input_ids_buf = None +connector._query_start_loc_buf = None +connector._ngram_context_buf = None +del connector +del rebuilt_registration +del registration +del payload +del unpinned_tp_input +del tensors +del before +del host_outputs +del host_flags +del input_ids +del query_start_loc +del ngram_context +gc.collect() + +assert all(not path.exists() for path in paths) +""" + + +def test_close_balanced_filename_storage_lifetime_exits_cleanly() -> None: + _require_file_system_shm() + + # Resolve from the connector actually imported by this pytest process. This + # lets the red command run this patched test against an unpatched worktree. + import_root = Path(connector_module.__file__).resolve().parents[3] + env = os.environ.copy() + inherited_pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + os.pathsep.join((str(import_root), inherited_pythonpath)) + if inherited_pythonpath + else str(import_root) + ) + + completed = subprocess.run( + [sys.executable, "-c", _BALANCED_LIFETIME_SCRIPT], + cwd=import_root, + env=env, + capture_output=True, + text=True, + timeout=60, + check=False, + ) + + assert completed.returncode == 0, ( + "balanced PLE storage destruction failed " + f"(exit code {completed.returncode})\n" + f"stdout:\n{completed.stdout}\n" + f"stderr:\n{completed.stderr}" + ) diff --git a/tests/v1/ple_offload/test_shutdown_order.py b/tests/v1/ple_offload/test_shutdown_order.py new file mode 100644 index 0000000..101c3da --- /dev/null +++ b/tests/v1/ple_offload/test_shutdown_order.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 + +import contextlib +from types import SimpleNamespace +from typing import Any + +from vllm.v1.ple_offload.connector import PleOffloadConnector +from vllm.v1.ple_offload import worker as ple_worker +from vllm.v1.worker import gpu_worker + + +def test_shutdown_marks_ple_stopping_before_teardown(monkeypatch) -> None: + worker = gpu_worker.Worker.__new__(gpu_worker.Worker) + worker._ple_offload_stopping = False + order: list[tuple[str, bool]] = [] + + def record(name: str): + def call() -> None: + order.append((name, worker._ple_offload_stopping)) + + return call + + monkeypatch.setattr(gpu_worker.gc, "unfreeze", record("gc")) + monkeypatch.setattr( + gpu_worker, + "ensure_kv_transfer_shutdown", + record("kv"), + ) + monkeypatch.setattr( + gpu_worker, + "ensure_ec_transfer_shutdown", + record("ec"), + ) + monkeypatch.setattr( + gpu_worker.current_platform, + "is_cuda_alike", + lambda: False, + ) + + worker.profiler = None + worker.weight_transfer_engine = None + worker.elastic_ep_executor = SimpleNamespace(shutdown=record("elastic")) + worker._ple_offload_worker_handle = SimpleNamespace(close=record("offload")) + worker.model_runner = SimpleNamespace(shutdown=record("runner")) + + worker.shutdown() + + assert order + assert all(stopping for _, stopping in order) + assert [name for name, _ in order][:2] == ["gc", "kv"] + + +class _RequestThread: + def __init__(self, order: list[str], *, stops: bool) -> None: + self.order = order + self.stops = stops + self.alive = True + + def is_alive(self) -> bool: + return self.alive + + def join(self, timeout: float) -> None: + self.order.append("thread.join") + if self.stops: + self.alive = False + + +def _record(order: list[str], name: str): + def call(*_args: Any, **_kwargs: Any) -> None: + order.append(name) + + return call + + +def _make_connector_for_close( + order: list[str], *, thread_stops: bool +) -> PleOffloadConnector: + connector = PleOffloadConnector.__new__(PleOffloadConnector) + connector._request_thread = _RequestThread(order, stops=thread_stops) + connector._request_queue = SimpleNamespace(put=_record(order, "queue.stop")) + connector._pinned_input_buffers = [] + connector._mapped_host_buffers = [] + connector._d2h_event_pool = object() + connector._registration_socket = SimpleNamespace( + close=_record(order, "socket.close") + ) + connector._zmq_ctx = SimpleNamespace(term=_record(order, "context.term")) + connector._progress = SimpleNamespace(close=_record(order, "progress.close")) + connector.release_host_storage = _record(order, "storage.release") + return connector + + +def test_connector_close_releases_storage_last() -> None: + order: list[str] = [] + connector = _make_connector_for_close(order, thread_stops=True) + + connector.close() + + assert order[-1] == "storage.release" + assert order.index("thread.join") < order.index("context.term") + assert order.index("context.term") < order.index("storage.release") + + +def test_connector_close_defers_storage_release_if_thread_does_not_stop() -> None: + order: list[str] = [] + connector = _make_connector_for_close(order, thread_stops=False) + + connector.close() + + assert "storage.release" not in order + assert "socket.close" not in order + assert "context.term" not in order + assert order[-1] == "progress.close" + + +def test_offload_child_unlinks_owned_ipc_after_context_term(monkeypatch) -> None: + order: list[str] = [] + ipc_addr = "ipc:///tmp/ple-owned-test.sock" + pull_socket = SimpleNamespace( + bind=_record(order, "socket.bind"), + close=_record(order, "socket.close"), + ) + zmq_context = SimpleNamespace( + socket=lambda _kind: pull_socket, + term=_record(order, "context.term"), + ) + + class FakeRunner: + layer_names: list[str] = [] + + def __init__(self, _config: object) -> None: + pass + + def accept_registrations(self, _socket: object, _count: int) -> None: + order.append("registrations.accept") + + def busy_loop(self, _socket: object, _event: object) -> None: + order.append("runner.busy_loop") + + ready_pipe = SimpleNamespace( + send_bytes=_record(order, "ready.send"), + close=_record(order, "ready.close"), + ) + monkeypatch.setattr(ple_worker, "decorate_logs", lambda _name: None) + monkeypatch.setattr(ple_worker, "mark_as_offload_worker", lambda: None) + monkeypatch.setattr(ple_worker, "_prepare_cuda_free_process", lambda: None) + monkeypatch.setattr(ple_worker, "_init_offload_distributed", lambda: None) + monkeypatch.setattr( + ple_worker, + "set_current_vllm_config", + lambda _config: contextlib.nullcontext(), + ) + monkeypatch.setattr(ple_worker, "PleOffloadRunner", FakeRunner) + monkeypatch.setattr(ple_worker.zmq, "Context", lambda: zmq_context) + monkeypatch.setattr(ple_worker.signal, "signal", lambda *_args: None) + monkeypatch.setattr(ple_worker.torch.cuda, "is_initialized", lambda: False) + monkeypatch.setattr( + ple_worker, + "_unlink_ipc_socket", + lambda address: order.append(f"ipc.unlink:{address}"), + ) + + ple_worker.PleOffloadWorker.proc_main( + SimpleNamespace(), + num_workers=1, + ipc_addr=ipc_addr, + ready_pipe=ready_pipe, + ) + + assert order.index("socket.close") < order.index("context.term") + assert order[-1] == f"ipc.unlink:{ipc_addr}" diff --git a/tests/v1/ple_offload/test_spawn_bootstrap_cuda.py b/tests/v1/ple_offload/test_spawn_bootstrap_cuda.py new file mode 100644 index 0000000..c709a93 --- /dev/null +++ b/tests/v1/ple_offload/test_spawn_bootstrap_cuda.py @@ -0,0 +1,352 @@ +# SPDX-License-Identifier: Apache-2.0 + +import ctypes +import ctypes.util +import inspect +import json +import multiprocessing +import os +import pickle +import signal +import socket +import subprocess +import sys +import threading +import time +from multiprocessing.connection import Connection +from pathlib import Path + +import pytest + +def _native_cuda_probe(output_path: str) -> None: + """Run during config unpickling, before the bootstrap imports vLLM.""" + library = ctypes.util.find_library("cudart") + if library is None: + os._exit(77) + cudart = ctypes.CDLL(library) + result = int(cudart.cudaFree(ctypes.c_void_p())) + Path(output_path).write_text( + json.dumps( + { + "cuda_visible_devices": os.environ.get( + "CUDA_VISIBLE_DEVICES" + ), + "cuda_result": result, + "torch_imported": "torch" in sys.modules, + "vllm_imported": "vllm" in sys.modules, + } + ), + encoding="utf-8", + ) + os._exit(0) + + +class _ProbePayload: + def __init__(self, output_path: str) -> None: + self.output_path = output_path + + def __reduce__(self): + # The bootstrap interpreter cannot import this test module, so rebuild + # the probe from its source text through a builtin instead of by name. + source = ( + "import ctypes, ctypes.util, json, os, sys\n" + "from pathlib import Path\n" + + inspect.getsource(_native_cuda_probe) + + f"_native_cuda_probe({self.output_path!r})\n" + ) + return exec, (source, {}) + + +def test_spawn_bootstrap_sets_visibility_before_native_cuda( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.v1.ple_offload.worker import PleOffloadWorker + + if ctypes.util.find_library("cudart") is None: + pytest.skip("CUDA runtime is not installed") + + original = os.environ.get("CUDA_VISIBLE_DEVICES") + output_path = tmp_path / "probe.json" + + # Use the production Popen launcher. Config reconstruction executes the + # probe in the fresh bootstrap interpreter before worker imports. + with monkeypatch.context() as child_settings: + child_settings.setenv("CUDA_VISIBLE_DEVICES", "0") + handle = PleOffloadWorker.make_process( + _ProbePayload(str(output_path)), 0, "unused" + ) + try: + returncode = handle.proc.wait(timeout=120) + finally: + handle.close() + + if returncode == 77: + pytest.skip("spawned child could not load the CUDA runtime") + assert returncode == 0 + + result = json.loads(output_path.read_text(encoding="utf-8")) + assert result["cuda_visible_devices"] == "0" + assert result["cuda_result"] == 0 + assert result["torch_imported"] is False + assert result["vllm_imported"] is False + assert os.environ.get("CUDA_VISIBLE_DEVICES") == original + + +class _FreshInterpreterProbe: + def __init__(self, output_path: str) -> None: + self.output_path = output_path + + def __reduce__(self): + expression = ( + "(__import__('pathlib').Path(" + + repr(self.output_path) + + ").write_text(__import__('json').dumps({" + "'cuda_visible_devices': __import__('os').environ.get(" + "'CUDA_VISIBLE_DEVICES')," + "'torch_imported': 'torch' in __import__('sys').modules," + "'vllm_imported': 'vllm' in __import__('sys').modules" + "}), encoding='utf-8'), None)[1]" + ) + return eval, (expression,) + + +@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only") +def test_fresh_interpreter_has_child_environment_before_vllm_import( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.v1.ple_offload.worker import PleOffloadWorker + + # This test module has unconditional top-level imports. The dedicated + # bootstrap interpreter must not import it as __mp_main__. + original = os.environ.get("CUDA_VISIBLE_DEVICES") + output_path = tmp_path / "fresh-interpreter.json" + + with monkeypatch.context() as child_settings: + child_settings.setenv("CUDA_VISIBLE_DEVICES", "child-only") + handle = PleOffloadWorker.make_process( + _FreshInterpreterProbe(str(output_path)), 0, "unused" + ) + try: + handle.proc.wait(timeout=120) + finally: + handle.close() + + assert output_path.exists() + result = json.loads(output_path.read_text(encoding="utf-8")) + assert result["cuda_visible_devices"] == "child-only" + assert result["torch_imported"] is False + assert result["vllm_imported"] is False + assert os.environ.get("CUDA_VISIBLE_DEVICES") == original + + +@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only") +def test_launch_payload_send_has_absolute_deadline() -> None: + reader, sender = socket.socketpair() + sender.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1024) + code = ( + "import sys,time\n" + "from vllm.v1.ple_offload.worker import _send_framed_payload\n" + "payload=b'x'*(64<<20)\n" + "start=time.monotonic()\n" + "try:\n" + " _send_framed_payload(int(sys.argv[1]),payload,start+0.25)\n" + "except TimeoutError:\n" + " print('TIMED_OUT',time.monotonic()-start,flush=True)\n" + "else:\n" + " raise AssertionError('blocked reader unexpectedly accepted payload')\n" + ) + process = subprocess.Popen( + [sys.executable, "-c", code, str(sender.fileno())], + pass_fds=(sender.fileno(),), + close_fds=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + sender.close() + try: + stdout, stderr = process.communicate(timeout=120) + finally: + reader.close() + if process.poll() is None: + process.terminate() + process.wait(timeout=5) + assert process.returncode == 0, stderr + # The injector logs the send entry through vLLM's logger on stdout; read only + # the probe's own result line. + result = [line for line in stdout.splitlines() if line.startswith('TIMED_OUT ')] + assert len(result) == 1, stdout + marker, elapsed = result[0].split() + assert marker == "TIMED_OUT" + assert float(elapsed) < 5 + + +class _BlockingConfigProbe: + def __init__(self, output_path: str) -> None: + self.output_path = output_path + + def __reduce__(self): + expression = ( + "(__import__('pathlib').Path(" + + repr(self.output_path) + + ").write_text(str(__import__('os').getpid()), encoding='utf-8')," + "__import__('threading').Event().wait())[1]" + ) + return eval, (expression,) + + +def _launch_blocked_config(output_path: str, child_pids) -> None: + from vllm.v1.ple_offload.worker import PleOffloadWorker + + handle = PleOffloadWorker.make_process( + _BlockingConfigProbe(output_path), 0, "unused" + ) + child_pids.put(handle.proc.pid) + # The test kills this launcher parent. Do not close the handle here. + threading.Event().wait() + + +def _pid_has_exited(pid: int) -> bool: + stat_path = Path(f"/proc/{pid}/stat") + if not stat_path.exists(): + return True + try: + return stat_path.read_text(encoding="utf-8").split()[2] == "Z" + except FileNotFoundError: + return True + + +@pytest.mark.skipif(sys.platform != "linux", reason="uses Linux process state") +def test_parent_death_during_config_reconstruction_exits_child( + tmp_path: Path, +) -> None: + context = multiprocessing.get_context("spawn") + child_pids = context.Queue() + marker = tmp_path / "blocked-config.pid" + parent = context.Process( + target=_launch_blocked_config, + args=(str(marker), child_pids), + ) + parent.start() + child_pid = child_pids.get(timeout=120) + deadline = time.monotonic() + 120 + while not marker.exists() and time.monotonic() < deadline: + time.sleep(0.05) + assert marker.exists() + + os.kill(parent.pid, signal.SIGKILL) + parent.join(timeout=120) + assert not parent.is_alive() + + deadline = time.monotonic() + 10 + while not _pid_has_exited(child_pid) and time.monotonic() < deadline: + time.sleep(0.05) + if not _pid_has_exited(child_pid): + # Harness cleanup is performed only after the autonomous-exit assertion + # has failed. + os.kill(child_pid, signal.SIGKILL) + pytest.fail("PLE child survived launcher-parent death during unpickle") + + +class _ExecFdProbe: + def __init__( + self, + output_path: str, + inherited: list[tuple[int, int]], + ) -> None: + self.output_path = output_path + self.inherited = inherited + + def __reduce__(self): + helper = ( + "import json,os,sys\n" + "from pathlib import Path\n" + "pairs=json.loads(sys.argv[3])\n" + "retained=[]\n" + "for fd,inode in pairs:\n" + " try:\n" + " retained.append(os.fstat(fd).st_ino==inode)\n" + " except OSError:\n" + " retained.append(False)\n" + "Path(sys.argv[1]).write_text(json.dumps({" + "'noninheritable_before_exec':sys.argv[2]=='True'," + "'closed_after_exec':not any(retained)}),encoding='utf-8')\n" + ) + pairs = json.dumps(self.inherited) + expression = ( + "__import__('os').execv(__import__('sys').executable," + "[__import__('sys').executable,'-c'," + + repr(helper) + + "," + + repr(self.output_path) + + ",str(all(not __import__('os').get_inheritable(int(fd)) " + "for fd in __import__('sys').argv[1:]))," + + repr(pairs) + + "])" + ) + return eval, (expression,) + + +@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only") +def test_actual_bootstrap_restores_cloexec_before_helper_exec( + tmp_path: Path, +) -> None: + from vllm.v1.ple_offload.worker import _send_framed_payload + + control_parent, control_child = socket.socketpair() + death_parent, death_child = socket.socketpair() + inherited = [ + (control_child.fileno(), os.fstat(control_child.fileno()).st_ino), + (death_child.fileno(), os.fstat(death_child.fileno()).st_ino), + ] + output_path = tmp_path / "bootstrap-fds.json" + config_payload = pickle.dumps( + _ExecFdProbe(str(output_path), inherited), + protocol=pickle.HIGHEST_PROTOCOL, + ) + payload = pickle.dumps( + { + "vllm_config_payload": config_payload, + "num_workers": 0, + "ipc_addr": "unused", + }, + protocol=pickle.HIGHEST_PROTOCOL, + ) + control = Connection(control_parent.detach()) + death_writer = Connection(death_parent.detach()) + process = subprocess.Popen( + [ + sys.executable, + "-m", + "vllm_ple_offload_bootstrap", + str(control_child.fileno()), + str(death_child.fileno()), + ], + close_fds=True, + pass_fds=(control_child.fileno(), death_child.fileno()), + ) + control_child.close() + death_child.close() + try: + _send_framed_payload( + control.fileno(), payload, time.monotonic() + 10 + ) + assert process.wait(timeout=120) == 0 + finally: + control.close() + death_writer.close() + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + + result = json.loads(output_path.read_text(encoding="utf-8")) + assert result == { + "noninheritable_before_exec": True, + "closed_after_exec": True, + } diff --git a/tests/v1/ple_offload/test_test_fault.py b/tests/v1/ple_offload/test_test_fault.py new file mode 100644 index 0000000..3dc8470 --- /dev/null +++ b/tests/v1/ple_offload/test_test_fault.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import contextlib +import os +import queue +import sys +from dataclasses import FrozenInstanceError +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + + +@pytest.fixture +def fault_module(monkeypatch, tmp_path): + monkeypatch.delenv("VLLM_PLE_TEST_FAULT", raising=False) + monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log")) + import vllm_ple_test_fault as test_fault + + return test_fault + + +def test_disarmed_import_is_none(fault_module): + assert fault_module._FAULT is None + + +@pytest.mark.parametrize( + "value", + [ + "cut=exc_before_prepare,ordinal=60,rank=any", + "cut=exc_after_prepare,ordinal=2,tp=1", + "cut=exc_in_release,ordinal=3,pp=0", + "cut=park_before_publish,ordinal=4,rank=any,phase=any", + ], +) +def test_valid_specs(fault_module, value): + spec = fault_module._parse_fault_spec(value) + assert spec is not None + assert spec.ordinal > 0 + + +@pytest.mark.parametrize( + "value", + [ + "", + " cut=exc_in_release,ordinal=1,rank=any", + "cut=unknown,ordinal=1,rank=any", + "cut=exc_in_release,ordinal=0,rank=any", + "cut=exc_in_release,ordinal=01,rank=any", + "cut=exc_in_release,ordinal=1", + "cut=exc_in_release,ordinal=1,rank=tp1", + "cut=exc_in_release,ordinal=1,rank=any,tp=0", + "cut=exc_in_release,ordinal=1,tp=-1", + "cut=exc_in_release,ordinal=1,tp=0,phase=startup", + "cut=exc_in_release,ordinal=1,tp=0,extra=x", + "cut=exc_in_release,cut=exc_after_release,ordinal=1,tp=0", + ], +) +def test_malformed_specs_raise(fault_module, value): + with pytest.raises(ValueError): + fault_module._parse_fault_spec(value) + + +def test_one_shot_rank_phase_and_ordinal(fault_module, monkeypatch): + spec = fault_module._parse_fault_spec( + "cut=exc_after_prepare,ordinal=2,tp=1,pp=0" + ) + assert spec is not None + + monkeypatch.setattr(fault_module.os, "getpid", lambda: 123) + monkeypatch.setattr(fault_module, "_pid_start_time", lambda _pid=None: 789) + monkeypatch.setattr(fault_module.time, "monotonic", lambda: 12.5) + monkeypatch.setattr(fault_module.time, "time", lambda: 456.25) + + def fires(cut, tp_rank=1, eligible=True): + return spec.fires( + cut, + role="gpu", + unit="runner-forward", + eligible=eligible, + tp_rank=tp_rank, + pp_rank=0, + dp_rank=0, + transaction_id="sequence:9@begin:12.000000000", + correlation_id="sequence:9", + ) + + assert not fires("other") + assert not fires("exc_after_prepare", tp_rank=0) + assert not fires("exc_after_prepare", eligible=False) + assert not fires("exc_after_prepare") + assert fires("exc_after_prepare") + assert not fires("exc_after_prepare") + line = Path(os.environ["VLLM_PLE_TEST_FAULT_LOG"]).read_text() + assert line == ( + "PLE_TEST_FAULT v=1 cut=exc_after_prepare ordinal=2 count=2 " + "unit=runner-forward role=gpu tp=1 pp=0 dp=0 pid=123 " + "pid_start=789 transaction=sequence:9@begin:12.000000000 " + "correlation=sequence:9 monotonic=12.500000000 " + "wall=456.250000000\n" + ) + with pytest.raises(FrozenInstanceError): + spec.ordinal = 3 + + +def test_disarmed_connector_sites_continue(fault_module, monkeypatch): + import vllm.v1.ple_offload.connector as connector_module + + monkeypatch.setattr(connector_module, "_FAULT", None) + + connector = object.__new__(connector_module.PleOffloadConnector) + connector.tp_rank = 0 + connector.pp_rank = 0 + connector.dp_rank = 0 + connector.device = SimpleNamespace(index=0) + connector._d2h_event_pool = queue.Queue() + event = Mock() + connector._d2h_event_pool.put_nowait(event) + connector._request_queue = Mock() + connector._progress = SimpleNamespace(sequence=0) + connector._enqueue_cuda_inputs = Mock() + + connector._launch( + num_reqs=1, + num_tokens=2, + serving=True, + transaction_id=None, + correlation_id=None, + ) + connector._enqueue_cuda_inputs.assert_called_once() + connector._request_queue.put_nowait.assert_called_once() + + staged = connector_module._PendingPleOffloadRequest( + connector_module.PleOffloadRequest(dp_rank=0, num_tokens=2, num_reqs=1), + event, + True, + "sequence:0", + ) + socket = Mock() + monkeypatch.setattr( + connector_module.torch.accelerator, + "device_index", + lambda _index: contextlib.nullcontext(), + ) + connector._process_request(staged, socket) + event.synchronize.assert_called_once() + socket.send.assert_called_once() + + +def test_disarmed_release_site_continues(fault_module, monkeypatch): + import vllm.v1.ple_offload.connector as connector_module + + monkeypatch.setattr(connector_module, "_FAULT", None) + monkeypatch.setattr(connector_module._rowdump, "enabled", lambda: False) + + stream = Mock() + completion = Mock() + monkeypatch.setattr( + connector_module.torch.cuda, + "current_stream", + lambda _device: stream, + ) + monkeypatch.setattr(connector_module.torch.cuda, "Event", lambda: completion) + + connector = object.__new__(connector_module.PleOffloadConnector) + connector.tp_rank = 0 + connector.pp_rank = 0 + connector.device = SimpleNamespace(index=0) + connector._layers = {} + connector._after_dummy = False + connector._progress = Mock() + + connector._release_outputs_impl() + completion.record.assert_called_once_with(stream) + connector._progress.track_completion.assert_called_once_with(completion) + + +def test_disarmed_worker_park_site_continues(fault_module, monkeypatch): + import vllm.v1.ple_offload.worker as worker_module + + monkeypatch.setattr(worker_module, "_FAULT", None) + runner = object.__new__(worker_module.PleOffloadRunner) + runner._validate_requests_before_publication = Mock() + runner._worker_targets = {0: {}} + runner._input_bufs = {} + runner._scratch_bufs = {} + runner._layers = {} + runner._clamp_input_ids = False + runner._step = 0 + runner._request_index = 0 + request = SimpleNamespace(dp_rank=0) + + runner._handle_requests_impl([request]) + runner._validate_requests_before_publication.assert_called_once_with([request]) + assert runner._step == 1 + + +def test_all_injection_sites_use_the_disarmed_short_circuit(): + root = Path(__file__).resolve().parents[3] + expected = { + "vllm/v1/ple_offload/connector.py": 3, + "vllm/v1/ple_offload/worker.py": 1, + "vllm/v1/worker/gpu/model_runner.py": 4, + "vllm_ple_offload_bootstrap/__init__.py": 1, + } + for relative, count in expected.items(): + lines = (root / relative).read_text(encoding="utf-8").splitlines() + sites = [i for i, line in enumerate(lines) if "_FAULT.fires(" in line] + assert len(sites) == count, relative + for i in sites: + window = "\n".join(lines[max(0, i - 6): i + 1]) + # The None check must precede fires() in the same condition, so a + # disarmed process never calls into the fault module. + assert "_FAULT is not None" in window, (relative, i + 1) + assert window.index("_FAULT is not None") < window.index("_FAULT.fires("), (relative, i + 1) + + +class _FatalReached(Exception): + pass + + +def test_exc_in_release_reaches_real_connector_fatal_wrapper( + fault_module, monkeypatch +): + import vllm.v1.ple_offload.connector as connector_module + + armed = fault_module._parse_fault_spec( + "cut=exc_in_release,ordinal=1,tp=1,pp=0" + ) + assert armed is not None + monkeypatch.setattr(connector_module, "_FAULT", armed) + + fatal = Mock(side_effect=_FatalReached) + monkeypatch.setattr(connector_module, "ple_fail_closed", fatal) + + connector = object.__new__(connector_module.PleOffloadConnector) + connector.rank = 1 + connector.tp_rank = 1 + connector.pp_rank = 0 + connector.dp_rank = 0 + connector._layers = {} + connector._after_dummy = False + connector._progress = SimpleNamespace(sequence=9) + + with pytest.raises(_FatalReached): + connector.release_outputs( + serving=True, + transaction_id="sequence:9@begin:1.000000000", + correlation_id="sequence:9", + ) + + assert fatal.call_args.kwargs["where"] == "release_outputs" + assert fatal.call_args.kwargs["rank"] == 1 + assert isinstance(fatal.call_args.kwargs["error"], RuntimeError) + + +@pytest.mark.timeout(180) +@pytest.mark.skipif( + sys.platform == "win32", + reason="production PLE launcher requires POSIX pass_fds", +) +@pytest.mark.parametrize("cut", ["launch_partial_ready", "launch_partial_ready_header"]) +def test_partial_ready_times_out_and_terminates_production_child( + monkeypatch, + tmp_path, + cut, +): + import vllm.v1.ple_offload.worker as worker_module + + monkeypatch.setenv( + "VLLM_PLE_TEST_FAULT", + f"cut={cut},ordinal=1,rank=any", + ) + monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log")) + monkeypatch.setattr( + worker_module.envs, + "VLLM_PLE_OFFLOAD_READY_TIMEOUT", + 1.0, + ) + + handle = worker_module.PleOffloadWorker.make_process( + object(), # The bootstrap fault fires before config reconstruction. + num_workers=1, + ipc_addr="ipc:///unused-ple-test", + ) + try: + with pytest.raises(TimeoutError): + worker_module.PleOffloadWorker.wait_for_ready(handle) + assert handle.proc.poll() is not None + finally: + handle.close(graceful=False) diff --git a/tests/v1/ple_offload/test_transport_bounded_io.py b/tests/v1/ple_offload/test_transport_bounded_io.py new file mode 100644 index 0000000..fb37122 --- /dev/null +++ b/tests/v1/ple_offload/test_transport_bounded_io.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 + +import multiprocessing +import os +import signal +import struct +import subprocess +import sys +import time +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + + +class _BlockingStrError(Exception): + def __str__(self) -> str: + __import__("threading").Event().wait() + return "unreachable" + + +def _blocking_fatal_formatter_case(marker_path: str) -> None: + from vllm.v1.ple_offload.fatal import ple_fail_closed + + Path(marker_path).write_text("armed", encoding="utf-8") + try: + raise _BlockingStrError() + except BaseException as error: + ple_fail_closed( + where="blocking-formatter", + rank=1, + layer="ple", + sequence=3, + error=error, + ) + + +def test_fatal_helper_does_not_call_blocking_exception_str( + tmp_path: Path, +) -> None: + context = multiprocessing.get_context("spawn") + marker = tmp_path / "fatal-formatter-armed" + process = context.Process( + target=_blocking_fatal_formatter_case, + args=(str(marker),), + ) + process.start() + try: + deadline = time.monotonic() + 120 + while not marker.exists() and time.monotonic() < deadline: + time.sleep(0.05) + assert marker.exists() + + process.join(timeout=5) + assert not process.is_alive() + assert process.exitcode not in (None, 0) + finally: + if process.is_alive(): + process.terminate() + process.join(timeout=120) + + +@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only") +@pytest.mark.parametrize( + ("behavior", "ready_timeout", "expected_exception"), + [ + pytest.param("partial", "1", TimeoutError, id="partial-frame"), + pytest.param("silent", "1", TimeoutError, id="no-response"), + pytest.param("ready", "120", None, id="complete-ready"), + pytest.param("invalid", "120", RuntimeError, id="invalid-layer-names"), + ], +) +def test_wait_for_ready_receives_complete_bounded_frame( + monkeypatch: pytest.MonkeyPatch, + behavior: str, + ready_timeout: str, + expected_exception: type[BaseException] | None, +) -> None: + import vllm.v1.ple_offload.worker as worker_module + + real_popen = subprocess.Popen + spawned: list[subprocess.Popen] = [] + child_code = ( + "import os,signal,struct,sys,time\n" + "from multiprocessing.connection import Connection\n" + "signal.alarm(10)\n" + "fd=int(sys.argv[1])\n" + "behavior=sys.argv[3]\n" + "if behavior=='partial':\n" + " os.write(fd,struct.pack('!i',128)+b'partial')\n" + "elif behavior=='ready':\n" + " Connection(fd).send({'status':'READY','layer_names':['ple']})\n" + "elif behavior=='invalid':\n" + " Connection(fd).send({'status':'READY','layer_names':'ple'})\n" + "time.sleep(120)\n" + ) + + def launch_fixture(argv: list[str], **kwargs: Any) -> subprocess.Popen: + process = real_popen( + [ + sys.executable, + "-c", + child_code, + argv[-2], + argv[-1], + behavior, + ], + **kwargs, + ) + spawned.append(process) + return process + + monkeypatch.setenv("VLLM_PLE_OFFLOAD_READY_TIMEOUT", ready_timeout) + monkeypatch.setattr(worker_module.subprocess, "Popen", launch_fixture) + + handle = worker_module.PleOffloadWorker.make_process(None, 0, "unused") + try: + started = time.monotonic() + if expected_exception is None: + worker_module.PleOffloadWorker.wait_for_ready(handle) + assert handle.ready_pipe_reader is None + else: + with pytest.raises(expected_exception): + worker_module.PleOffloadWorker.wait_for_ready(handle) + if expected_exception is TimeoutError: + assert time.monotonic() - started < 5 + assert spawned[0].poll() is not None + finally: + handle.close(graceful=False) + + +@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only") +def test_make_process_large_payload_send_has_absolute_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import vllm.v1.ple_offload.worker as worker_module + + real_popen = subprocess.Popen + spawned: list[subprocess.Popen] = [] + + def launch_nonreader(argv: list[str], **kwargs: Any) -> subprocess.Popen: + process = real_popen( + [ + sys.executable, + "-c", + "import signal,time; signal.alarm(120); time.sleep(120)", + argv[-2], + argv[-1], + ], + **kwargs, + ) + spawned.append(process) + return process + + monkeypatch.setenv("VLLM_PLE_OFFLOAD_READY_TIMEOUT", "5") + monkeypatch.setattr(worker_module.subprocess, "Popen", launch_nonreader) + + started = time.monotonic() + with pytest.raises(TimeoutError, match="transmission"): + worker_module.PleOffloadWorker.make_process( + b"x" * (64 << 20), + 0, + "unused", + ) + + assert time.monotonic() - started < 20 + assert spawned + assert spawned[0].poll() is not None + + +@pytest.mark.parametrize( + "identities", + [ + pytest.param(("shared", "shared"), id="duplicate"), + pytest.param(("",), id="unbound"), + ], +) +def test_expected_ple_layers_are_bound_and_unique_before_publication( + monkeypatch: pytest.MonkeyPatch, + identities: tuple[str, ...], +) -> None: + from vllm.config import CUDAGraphMode + from vllm.v1.worker.gpu.cudagraph_utils import ( + CudaGraphManager, + ModelCudaGraphManager, + ) + + manager = ModelCudaGraphManager.__new__(ModelCudaGraphManager) + manager.graphs = {} + manager.use_breakable_cg = False + manager.cudagraph_mode = CUDAGraphMode.FULL + + modules = [ + ( + f"module_{index}", + SimpleNamespace( + _is_cpu_offloaded=True, + _host_output_ptr=1, + _ple_offload_name=identity, + ), + ) + for index, identity in enumerate(identities) + ] + model = SimpleNamespace(named_modules=lambda: iter(modules)) + + def publish_graph(self: Any, *_args: Any, **_kwargs: Any) -> None: + self.graphs["published"] = object() + + monkeypatch.setattr(CudaGraphManager, "capture", publish_graph) + + with pytest.raises(RuntimeError, match="distinct, bound"): + ModelCudaGraphManager.capture( + manager, + model=model, + model_state=object(), + input_buffers=object(), + intermediate_tensors=None, + block_tables=object(), + attn_groups=[], + kv_cache_config=object(), + ) + assert manager.graphs == {} + + +def test_real_wait_copy_geometry_rejected_before_graph_publication( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import torch + + import vllm.model_executor.layers.ple_offload_layer as ple_layer_module + from vllm.config import CUDAGraphMode + from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + CudaGraphManager, + ) + + manager = CudaGraphManager.__new__(CudaGraphManager) + manager.graphs = {} + manager._ple_full_copy_geometry = {} + manager._expected_ple_copy_layers = ("ple",) + manager._allow_ple_full_graph = True + desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, + num_tokens=8, + num_reqs=1, + ) + + success = SimpleNamespace(value=0) + monkeypatch.setattr( + ple_layer_module.torch.cuda, + "current_stream", + lambda: SimpleNamespace(cuda_stream=1), + ) + monkeypatch.setattr( + ple_layer_module.cuda_driver, + "CUstream", + lambda value: value, + ) + monkeypatch.setattr( + ple_layer_module.cuda_driver, + "CUdeviceptr", + lambda value: value, + ) + monkeypatch.setattr( + ple_layer_module.cuda_driver, + "cuStreamWaitValue32", + lambda *_args: (success,), + ) + monkeypatch.setattr( + ple_layer_module.cuda_driver, + "cuMemcpyHtoDAsync", + lambda *_args: (success,), + ) + + gpu_output_buffer = torch.empty(12, 4, dtype=torch.float32) + hidden_states = torch.empty(12, 1, dtype=torch.float32) + + def capture_then_publish() -> None: + with manager.capture_ple_copy_geometry(desc): + ple_layer_module._ple_offload_wait_copy_impl( + gpu_output_buffer, + hidden_states, + 1, + 2, + 16, + 12, + "ple", + ) + manager.graphs[desc] = object() + + with pytest.raises(RuntimeError, match="actual PLE copy geometry"): + capture_then_publish() + assert desc not in manager._ple_full_copy_geometry + assert desc not in manager.graphs + + +class _CleanupConfig: + def __init__(self, started_path: str, cleaned_path: str) -> None: + self.started_path = started_path + self.cleaned_path = cleaned_path + + def __reduce__(self): + source = ( + "import signal,threading\n" + "from pathlib import Path\n" + "from vllm.v1.ple_offload.worker import PleOffloadWorker\n" + f"_started=Path({self.started_path!r})\n" + f"_cleaned=Path({self.cleaned_path!r})\n" + "def _cleanup_proc_main(vllm_config,num_workers,ipc_addr,ready_pipe):\n" + " _stop=threading.Event()\n" + " signal.signal(signal.SIGTERM,lambda *_args:_stop.set())\n" + " _started.write_text('started',encoding='utf-8')\n" + " try:\n" + " _stop.wait()\n" + " finally:\n" + " _cleaned.write_text('cleaned',encoding='utf-8')\n" + "PleOffloadWorker.proc_main=staticmethod(_cleanup_proc_main)\n" + ) + return exec, (source, {}) + + +@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only") +def test_normal_close_runs_child_cleanup(tmp_path: Path) -> None: + from vllm.v1.ple_offload.worker import PleOffloadWorker + + started = tmp_path / "cleanup-started" + cleaned = tmp_path / "cleanup-finished" + handle = PleOffloadWorker.make_process( + _CleanupConfig(str(started), str(cleaned)), + 0, + "unused", + ) + try: + deadline = time.monotonic() + 120 + while not started.exists() and time.monotonic() < deadline: + time.sleep(0.05) + assert started.exists() + finally: + handle.close() + + assert handle.proc.returncode == 0 + assert cleaned.read_text(encoding="utf-8") == "cleaned" diff --git a/tests/v1/ple_offload/test_transport_fatal_helper.py b/tests/v1/ple_offload/test_transport_fatal_helper.py new file mode 100644 index 0000000..4e5d5c6 --- /dev/null +++ b/tests/v1/ple_offload/test_transport_fatal_helper.py @@ -0,0 +1,387 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +import signal +import subprocess +import sys +import time +from pathlib import Path + +import pytest + + +def _run_bounded_child( + source: str, + tmp_path: Path, + *arguments: str, + expected_exitcode: int = 0, +) -> None: + marker = tmp_path / f"started-{time.monotonic_ns()}" + process = subprocess.Popen( + [sys.executable, "-c", source, str(marker), *arguments], + stderr=subprocess.DEVNULL, + start_new_session=os.name == "posix", + ) + stayed_alive = False + try: + deadline = time.monotonic() + 120 + while not marker.exists() and time.monotonic() < deadline: + time.sleep(0.05) + assert marker.exists(), ( + "child did not finish cold startup within 120 seconds" + ) + process.wait(timeout=5) + except subprocess.TimeoutExpired: + stayed_alive = True + finally: + if process.poll() is None: + if os.name == "posix": + os.killpg(process.pid, signal.SIGKILL) + else: + process.terminate() + process.wait(timeout=120) + + assert not stayed_alive, "child stayed alive after the bounded operation" + assert process.returncode == expected_exitcode + + +_FATAL_SOURCE = r""" +import os +import sys +import threading +from pathlib import Path + +from vllm.v1.ple_offload.fatal import ple_fail_closed + +scenario = sys.argv[2] + +if scenario == "blocking-name-property": + class Meta(type): + @property + def __name__(cls): + threading.Event().wait() + + class Error(Exception, metaclass=Meta): + pass + + error = Error() +elif scenario == "blocking-str-descriptor": + class Descriptor: + def __get__(self, instance, owner): + threading.Event().wait() + + class Error(Exception): + __str__ = Descriptor() + + error = Error() +elif scenario == "blocking-name-format": + class BlockingName(str): + def __format__(self, format_spec): + threading.Event().wait() + + class Meta(type): + @property + def __name__(cls): + return BlockingName("Error") + + class Error(Exception, metaclass=Meta): + pass + + error = Error() +elif scenario == "blocking-str-method": + class Error(Exception): + def __str__(self): + threading.Event().wait() + + error = Error() +elif scenario == "overridden-getattribute": + class Error(Exception): + def __getattribute__(self, name): + threading.Event().wait() + + error = Error("detail") +elif scenario == "str-subclass-argument": + class BlockingString(str): + def __format__(self, format_spec): + threading.Event().wait() + + def __str__(self): + threading.Event().wait() + + error = RuntimeError(BlockingString("detail")) +elif scenario == "huge-string": + error = RuntimeError("x" * (32 << 20)) +else: + raise AssertionError(scenario) + +Path(sys.argv[1]).write_text("armed", encoding="utf-8") +ple_fail_closed( + where="transport-fatal", + rank=1, + layer="ple", + sequence=3, + error=error, +) +""" + + +@pytest.mark.parametrize( + "scenario", + [ + "blocking-name-property", + "blocking-str-descriptor", + "blocking-name-format", + ], +) +def test_fatal_helper_avoids_exception_controlled_descriptors( + tmp_path: Path, + scenario: str, +) -> None: + _run_bounded_child( + _FATAL_SOURCE, + tmp_path, + scenario, + expected_exitcode=1, + ) + + +@pytest.mark.parametrize( + "scenario", + [ + "blocking-str-method", + "overridden-getattribute", + "str-subclass-argument", + "huge-string", + ], +) +def test_fatal_helper_hard_exit_controls( + tmp_path: Path, + scenario: str, +) -> None: + _run_bounded_child( + _FATAL_SOURCE, + tmp_path, + scenario, + expected_exitcode=1, + ) + + +@pytest.mark.skipif(sys.platform == "win32", reason="process groups are POSIX-only") +def test_wait_for_ready_rejects_non_string_status_without_comparison( + tmp_path: Path, +) -> None: + source = r""" +import subprocess +import sys +import threading +import time +from pathlib import Path + +import vllm.v1.ple_offload.worker as worker_module + +class Status: + def __eq__(self, other): + threading.Event().wait() + + def __ne__(self, other): + threading.Event().wait() + +class Reader: + def close(self): + return None + +sleeper = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(120)"] +) +handle = worker_module.PleOffloadWorkerHandle( + proc=sleeper, + death_writer=None, + ready_pipe_reader=Reader(), + launch_deadline=time.monotonic() + 120, +) +worker_module._recv_bounded_connection_message = ( + lambda reader, deadline: { + "status": Status(), + "layer_names": ["ple"], + } +) +Path(sys.argv[1]).write_text("armed", encoding="utf-8") +rejected = False +closed = False +reader_cleared = False +try: + try: + worker_module.PleOffloadWorker.wait_for_ready(handle) + except RuntimeError: + rejected = True + # Postconditions are read before the emergency cleanup below can act. + # Stay inside the parent's 5 second bounded wait so a miss reports here. + close_deadline = time.monotonic() + 3 + while sleeper.poll() is None and time.monotonic() < close_deadline: + time.sleep(0.05) + closed = sleeper.poll() is not None + reader_cleared = handle.ready_pipe_reader is None +finally: + if sleeper.poll() is None: + sleeper.kill() + sleeper.wait(timeout=120) + +if not rejected: + raise AssertionError("invalid response status was not rejected") +if not closed: + raise AssertionError("invalid response did not close the child") +if not reader_cleared: + raise AssertionError("invalid response did not clear the ready reader") +""" + _run_bounded_child(source, tmp_path) + + +@pytest.mark.parametrize("identity", ["duplicate", "unbound"]) +def test_profile_breakable_initializer_follows_identity_validation( + tmp_path: Path, + identity: str, +) -> None: + source = r""" +import os +import sys +import threading +from pathlib import Path +from types import SimpleNamespace + +from vllm.v1.worker.gpu.cudagraph_utils import ModelCudaGraphManager + +identity = sys.argv[2] +if identity == "duplicate": + identities = ("shared", "shared") +else: + identities = ("",) + +modules = [ + ( + "module_" + str(index), + SimpleNamespace( + _is_cpu_offloaded=True, + _host_output_ptr=1, + _ple_offload_name=value, + ), + ) + for index, value in enumerate(identities) +] +model = SimpleNamespace(named_modules=lambda: iter(modules)) +manager = ModelCudaGraphManager.__new__(ModelCudaGraphManager) +spy = Path(sys.argv[1] + ".initializer") + +def blocking_initializer(model): + spy.write_text("called", encoding="utf-8") + threading.Event().wait() + +manager.init_breakable_cg_runner = blocking_initializer +Path(sys.argv[1]).write_text("armed", encoding="utf-8") + +prepare = getattr(manager, "_prepare_ple_capture", None) +if prepare is None: + # This is the pre-fix profiling order: initialize before validation. + manager.init_breakable_cg_runner(model) +else: + try: + prepare(model, initialize_breakable=True) + except RuntimeError as error: + if "distinct, bound" not in str(error): + raise + else: + raise AssertionError("invalid identities were accepted") + +if spy.exists(): + raise AssertionError("initializer ran before identity rejection") +""" + _run_bounded_child(source, tmp_path, identity) + + +@pytest.mark.skipif(sys.platform == "win32", reason="process groups are POSIX-only") +@pytest.mark.parametrize("status", ["READY", "FAILURE"]) +def test_oversized_startup_response_becomes_bounded_failure( + tmp_path: Path, + status: str, +) -> None: + source = r""" +import socket +import sys +from multiprocessing.connection import Connection +from pathlib import Path + +import vllm.v1.ple_offload.worker as worker_module + +parent_socket, child_socket = socket.socketpair() +raw_writer = Connection(child_socket.detach()) +writer_type = getattr( + worker_module, + "_BoundedStartupResponseWriter", + None, +) +writer = raw_writer if writer_type is None else writer_type(raw_writer) + +if sys.argv[2] == "READY": + message = { + "status": "READY", + "layer_names": ["x" * (2 << 20)], + } +else: + message = { + "status": "FAILURE", + "error": "x" * (2 << 20), + } + +Path(sys.argv[1]).write_text("armed", encoding="utf-8") +writer.send(message) +""" + _run_bounded_child(source, tmp_path, status) + + +def test_startup_response_payload_at_cap_is_accepted(tmp_path: Path) -> None: + source = r""" +import pickle +import sys +import threading +from pathlib import Path + +import vllm.v1.ple_offload.worker as worker_module + +writer_type = getattr( + worker_module, + "_BoundedStartupResponseWriter", + None, +) +Path(sys.argv[1]).write_text("armed", encoding="utf-8") +if writer_type is None: + threading.Event().wait() + +cap = worker_module._PLE_READY_MAX_PAYLOAD_BYTES +name = "x" * cap +for unused in range(8): + message = {"status": "READY", "layer_names": [name]} + payload = pickle.dumps(message, protocol=pickle.HIGHEST_PROTOCOL) + if len(payload) == cap: + break + next_size = len(name) + cap - len(payload) + if next_size < 0: + raise AssertionError("cannot construct exact-cap payload") + name = "x" * next_size +else: + raise AssertionError("exact-cap payload did not converge") + +class Writer: + def __init__(self): + self.payload = None + + def send_bytes(self, payload): + self.payload = payload + + def close(self): + return None + +raw_writer = Writer() +writer_type(raw_writer).send(message) +if raw_writer.payload != payload or len(raw_writer.payload) != cap: + raise AssertionError("exact-cap startup response was not preserved") +""" + _run_bounded_child(source, tmp_path) diff --git a/vllm/model_executor/layers/ple_offload_layer.py b/vllm/model_executor/layers/ple_offload_layer.py new file mode 100644 index 0000000..1b67578 --- /dev/null +++ b/vllm/model_executor/layers/ple_offload_layer.py @@ -0,0 +1,433 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Base class for PLE layers that support cross-process CPU offload. + +Synchronization between the offload process and the GPU main stream is +handled by :class:`CpuGpuSemaphore`, which uses CUDA's +``cuStreamWaitValue32`` / ``cuStreamWriteValue32`` stream memory operations. + +Decode-step timeline, host-mapped PULL design +-------------------------------------------------------- +Offload process (CPU, no CUDA) GPU main stream (forward) +------------------------------ ------------------------------ +forward_impl() on CPU ... other GPU ops ... +wait host flag == 0 WaitValue32(host flag == 1) <- in graph +rows -> shared pinned host slot cuMemcpyHtoDAsync(slot -> GPU) <- in graph +publish fence; host flag = 1 return _gpu_output_buffer[:n] + ... model consumes output ... + WriteValue32(host flag = 0) <- release_outputs + +``release_outputs`` performs the sole reset. Safety requires serialized local +forwards on one ordered model stream: a new transaction cannot begin until the +prior release has been enqueued. If any prepared forward is abandoned or its +correspondence becomes uncertain, the worker must terminate; resetting, +retrying, or continuing could reuse stale rows or overwrite live input. +""" + +import contextlib +import functools +import threading +from abc import ABC, abstractmethod +from collections.abc import Callable, Iterator +from typing import Any, cast + +import torch +from cuda.bindings import driver as cuda_driver +from cuda.bindings.driver import CUstreamWaitValue_flags +from torch import nn + +import vllm.envs as envs +from vllm.utils.torch_utils import direct_register_custom_op + +# Module-level flag set to True inside the offload subprocess. +# Because the offload process and GPU worker processes are separate OS +# processes (spawned via multiprocessing), each has its own memory space. +# A plain module-level bool is sufficient -- no thread-local storage needed. +_offload_worker_flag = False + +# Native FULL capture runs on one Python thread. The custom-op implementation +# appends the exact arguments from every wait-copy invocation while that thread +# is inside the graph manager's capture scope. +_ple_copy_capture = threading.local() + + +@contextlib.contextmanager +def observe_ple_copy_geometry() -> Iterator[list[tuple[str, int, int]]]: + """Collect ``(layer, rows, bytes)`` from actual wait-copy executions.""" + if getattr(_ple_copy_capture, "records", None) is not None: + raise RuntimeError("nested PLE copy-geometry observation is unsupported") + records: list[tuple[str, int, int]] = [] + _ple_copy_capture.records = records + try: + yield records + finally: + del _ple_copy_capture.records + + +def _record_ple_copy_geometry( + layer_name: str, + num_tokens: int, + row_bytes: int, +) -> None: + records = getattr(_ple_copy_capture, "records", None) + if records is not None: + rows = int(num_tokens) + records.append((str(layer_name), rows, rows * int(row_bytes))) + + +def is_offload_process() -> bool: + """Return True inside the dedicated PLE CPU-offload subprocess.""" + return _offload_worker_flag + + +def mark_as_offload_worker() -> None: + """Mark the current process as the dedicated PLE offload worker.""" + global _offload_worker_flag + _offload_worker_flag = True + + +def _cuda_check(result: Any, operation: str) -> Any: + """Check the ``(CUresult, ...)`` tuple returned by cuda-python calls.""" + error = result[0] if isinstance(result, tuple) else result + if error.value != 0: + raise RuntimeError(f"{operation} failed: {error}") + return result + + +class CpuGpuSemaphore: + """Cross-process semaphore backed by a one-element int32 CUDA tensor. + + The flag is stored in a regular GPU tensor that is shared through + PyTorch's CUDA IPC mechanism. Both processes obtain device pointers that + map to the same physical GPU memory, so a stream-memory wait in the GPU + worker observes the write issued by the offload process. + """ + + RESET_VALUE = 0 + DONE_VALUE = 1 + + def __init__(self, device: torch.device) -> None: + self._flag_tensor = torch.zeros(1, dtype=torch.int32, device=device) + self._flag_ptr: int | None = None + + @classmethod + def from_ipc_tensor(cls, flag_tensor: torch.Tensor) -> "CpuGpuSemaphore": + """Construct a semaphore from a CUDA tensor received through IPC.""" + semaphore = cls.__new__(cls) + semaphore._flag_tensor = flag_tensor + semaphore._flag_ptr = None + return semaphore + + @classmethod + def from_host_mapped( + cls, host_flag_tensor: torch.Tensor, device_ptr: int + ) -> "CpuGpuSemaphore": + """Semaphore on a page-locked, device-mapped shared host int32. + + ``host_flag_tensor`` keeps the shared storage alive; ``device_ptr`` is + its device address (``cuMemHostGetDevicePointer``). Stream memory ops + on it are executed by the GPU against host memory, so a CPU process + with no CUDA context can signal it with a plain store. + """ + semaphore = cls.__new__(cls) + semaphore._flag_tensor = host_flag_tensor + semaphore._flag_ptr = int(device_ptr) + return semaphore + + @property + def flag_tensor(self) -> torch.Tensor: + """Return the tensor that backs the semaphore flag.""" + return self._flag_tensor + + @property + def flag_ptr(self) -> int: + """Device address used by the stream memory operations.""" + if self._flag_ptr is not None: + return self._flag_ptr + return self._flag_tensor.data_ptr() + + def reset(self, stream: torch.cuda.Stream | None = None) -> None: + """Enqueue ``WriteValue32(flag=0)`` on ``stream``.""" + if stream is None: + stream = torch.cuda.current_stream() + _cuda_check( + cuda_driver.cuStreamWriteValue32( + cuda_driver.CUstream(stream.cuda_stream), + cuda_driver.CUdeviceptr(self.flag_ptr), + self.RESET_VALUE, + 0, + ), + "CpuGpuSemaphore.reset", + ) + + def signal(self, stream: torch.cuda.Stream | None = None) -> None: + """Enqueue ``WriteValue32(flag=1)`` on ``stream``.""" + if stream is None: + stream = torch.cuda.current_stream() + _cuda_check( + cuda_driver.cuStreamWriteValue32( + cuda_driver.CUstream(stream.cuda_stream), + cuda_driver.CUdeviceptr(self.flag_ptr), + self.DONE_VALUE, + 0, + ), + "CpuGpuSemaphore.signal", + ) + + def wait_reset(self, stream: torch.cuda.Stream | None = None) -> None: + """Enqueue ``WaitValue32(flag==0)`` on ``stream``.""" + if stream is None: + stream = torch.cuda.current_stream() + _cuda_check( + cuda_driver.cuStreamWaitValue32( + cuda_driver.CUstream(stream.cuda_stream), + cuda_driver.CUdeviceptr(self.flag_ptr), + self.RESET_VALUE, + CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ.value, + ), + "CpuGpuSemaphore.wait_reset", + ) + + +# --------------------------------------------------------------------------- +# Custom op: vllm::ple_offload_wait +# --------------------------------------------------------------------------- +# Wraps cuStreamWaitValue32 as a torch custom op so that +# torch.compile(fullgraph=True) treats the CUDA Driver call as an opaque node. +# The hidden_states argument creates a dynamo data-dependency edge that prevents +# the wait from being reordered before preceding GPU work. +# --------------------------------------------------------------------------- + + +def _ple_offload_wait_impl( + sem_flag_tensor: torch.Tensor, + gpu_output_buffer: torch.Tensor, + hidden_states: torch.Tensor, +) -> None: + """Wait for the CPU result without releasing its output buffer.""" + stream = torch.cuda.current_stream() + cuda_stream = cuda_driver.CUstream(stream.cuda_stream) + dev_ptr = cuda_driver.CUdeviceptr(sem_flag_tensor.data_ptr()) + _cuda_check( + cuda_driver.cuStreamWaitValue32( + cuda_stream, + dev_ptr, + CpuGpuSemaphore.DONE_VALUE, + CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ.value, + ), + "cuStreamWaitValue32(done)", + ) + + +def _ple_offload_wait_fake( + sem_flag_tensor: torch.Tensor, + gpu_output_buffer: torch.Tensor, + hidden_states: torch.Tensor, +) -> None: + """Represent the side-effect-only wait during dynamo tracing.""" + pass + + +direct_register_custom_op( + op_name="ple_offload_wait", + op_func=_ple_offload_wait_impl, + mutates_args=["gpu_output_buffer"], + fake_impl=_ple_offload_wait_fake, +) + + +# --------------------------------------------------------------------------- +# Custom op: vllm::ple_offload_wait_copy (host-mapped pull design) +# --------------------------------------------------------------------------- +# The offload process writes the result into a page-locked shared host buffer +# and stores flag=1 in a page-locked shared host int32; it holds no CUDA +# context. On the CURRENT stream (captured into the graph like the wait above): +# cuStreamWaitValue32(host flag == 1) -> GPU polls host memory +# cuMemcpyAsync(gpu buffer <- host buffer, num_tokens * row_bytes) +# The copy engine reads host memory coherently after the flag was observed, +# and the consumers read device memory as before. The flag is NOT reset here: +# ``release_outputs`` resets it after the forward (unchanged semantics), so +# capture warm-ups and dummy runs behave exactly as with ``ple_offload_wait``. +# Pointers are passed as ints so no CPU tensor enters the compiled graph. +# --------------------------------------------------------------------------- + + +def _ple_offload_wait_copy_impl( + gpu_output_buffer: torch.Tensor, + hidden_states: torch.Tensor, + flag_ptr: int, + host_ptr: int, + row_bytes: int, + num_tokens: int, + layer_name: str, +) -> None: + """Wait for the host-side result, then DMA it into the GPU buffer.""" + if not 0 <= int(num_tokens) <= gpu_output_buffer.shape[0] or int(row_bytes) != ( + gpu_output_buffer.shape[1] * gpu_output_buffer.element_size() + ): + raise RuntimeError( + f"PLE wait_copy out of bounds: {num_tokens} rows x {row_bytes} B vs " + f"GPU buffer {tuple(gpu_output_buffer.shape)} {gpu_output_buffer.dtype}" + ) + nbytes = int(num_tokens) * int(row_bytes) + _record_ple_copy_geometry(layer_name, num_tokens, row_bytes) + stream = cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream) + _cuda_check( + cuda_driver.cuStreamWaitValue32( + stream, + cuda_driver.CUdeviceptr(flag_ptr), + CpuGpuSemaphore.DONE_VALUE, + CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ.value, + ), + "cuStreamWaitValue32(host flag done)", + ) + if nbytes > 0: + # host_ptr is this worker's host address of page-locked (registered) + # memory, so the copy is a true async DMA and is capturable. + _cuda_check( + cuda_driver.cuMemcpyHtoDAsync( + cuda_driver.CUdeviceptr(gpu_output_buffer.data_ptr()), + host_ptr, + nbytes, + stream, + ), + "cuMemcpyHtoDAsync(PLE host -> GPU)", + ) + + +def _ple_offload_wait_copy_fake( + gpu_output_buffer: torch.Tensor, + hidden_states: torch.Tensor, + flag_ptr: int, + host_ptr: int, + row_bytes: int, + num_tokens: int, + layer_name: str, +) -> None: + """Represent the side-effect-only wait + copy during dynamo tracing.""" + pass + + +direct_register_custom_op( + op_name="ple_offload_wait_copy", + op_func=_ple_offload_wait_copy_impl, + mutates_args=["gpu_output_buffer"], + fake_impl=_ple_offload_wait_copy_fake, +) + + +class PleOffloadLayer(nn.Module, ABC): + """Base class for embedding-like PLE layers that can run on CPU. + + Subclasses implement :meth:`forward_impl`, which contains the regular + on-device computation. In a GPU worker with PLE offload enabled, subclass + constructors are skipped so their large weights are never allocated. The + CPU process constructs a complete copy and owns those weights instead. + """ + + _is_cpu_offloaded = False + _gpu_output_buffer: torch.Tensor + _sem: CpuGpuSemaphore + # Host-mapped pull design: host address of the shared output buffer and + # the bytes per output row. ``None`` keeps the IPC push behaviour. + _host_output_ptr: int | None = None + _host_row_bytes: int = 0 + _ple_offload_name: str = "" + + def __init_subclass__(cls, **kwargs: object) -> None: + """Skip subclass initialization in cross-process GPU-worker mode.""" + super().__init_subclass__(**kwargs) + original_init_obj = cls.__dict__.get("__init__") + if original_init_obj is None or not callable(original_init_obj): + return + original_init = cast(Callable[..., None], original_init_obj) + + @functools.wraps(original_init) + def guarded_init( + self: "PleOffloadLayer", *args: object, **kwargs: object + ) -> None: + if envs.VLLM_PLE_CPU_OFFLOAD and not is_offload_process(): + nn.Module.__init__(self) + return + original_init(self, *args, **kwargs) + + cls.__init__ = guarded_init # type: ignore[method-assign, assignment] + + @classmethod + def get_target_device(cls) -> torch.device: + """Return CPU for the offload process and the active GPU otherwise.""" + if envs.VLLM_PLE_CPU_OFFLOAD: + return torch.device("cpu") + return torch.device("cuda", torch.accelerator.current_device_index()) + + @abstractmethod + def forward_impl( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + *args: object, + **kwargs: object, + ) -> torch.Tensor: + """Execute the actual embedding computation on the owning device.""" + raise NotImplementedError + + def get_offload_output_dtype(self, default_dtype: torch.dtype) -> torch.dtype: + """Return the dtype used by cross-process output buffers.""" + return default_dtype + + def get_offload_output_dim(self, default_dim: int) -> int: + """Return the last dimension used by cross-process output buffers.""" + return default_dim + + def initialize_dummy_offload_metadata(self, device: torch.device) -> None: + """Initialize metadata needed by dummy offload execution.""" + + def setup_cross_process_offload( + self, + gpu_output_buffer: torch.Tensor, + semaphore: CpuGpuSemaphore, + host_output_ptr: int | None = None, + host_row_bytes: int = 0, + ) -> None: + """Configure the GPU-worker placeholder with its offload resources.""" + self._is_cpu_offloaded = True + self._gpu_output_buffer = gpu_output_buffer + self._sem = semaphore + self._host_output_ptr = host_output_ptr + self._host_row_bytes = int(host_row_bytes) + + def forward( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + *args: object, + **kwargs: object, + ) -> torch.Tensor: + """Wait for an offloaded result or delegate to ``forward_impl``.""" + if self._is_cpu_offloaded: + if self._host_output_ptr is not None: + torch.ops.vllm.ple_offload_wait_copy( + self._gpu_output_buffer, + hidden_states, + self._sem.flag_ptr, + self._host_output_ptr, + self._host_row_bytes, + input_ids.shape[0], + self._ple_offload_name, + ) + else: + torch.ops.vllm.ple_offload_wait( + self._sem.flag_tensor, + self._gpu_output_buffer, + hidden_states, + ) + return self._gpu_output_buffer[: input_ids.shape[0]] + return self.forward_impl(hidden_states, input_ids, *args, **kwargs) + + def release_offloaded_output( + self, + stream: torch.cuda.Stream | None = None, + ) -> None: + """Mark the cross-process output buffer reusable on ``stream``.""" + if self._is_cpu_offloaded: + self._sem.reset(stream) diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 1ccd23a..e27a18b 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -352,9 +352,10 @@ class MultiprocExecutor(Executor): proc = next(h.proc for h in workers if h.proc.sentinel == died[0]) logger.error( "Worker proc %s died unexpectedly (exit code: %s), " - "shutting down executor.", + "shutting down executor. executor_shutdown_monotonic=%.9f", proc.name, proc.exitcode, + time.monotonic(), ) _self.shutdown() callback = _self.failure_callback @@ -1034,6 +1035,28 @@ class WorkerProc: SUCCESS = auto() FAILURE = auto() + def _fail_ple_rpc_if_pending( + self, where: str, error: BaseException + ) -> None: + """Fail closed if RPC recovery would abandon queued PLE GPU work.""" + owners = (self.worker, getattr(self.worker, "worker", None)) + for owner in owners: + if owner is None: + continue + runner = getattr(owner, "model_runner", None) + connector = getattr(runner, "_ple_offload_connector", None) + if connector is None: + continue + try: + pending = connector.has_pending_completions + except BaseException as state_error: + connector._fail_closed( + f"{where}:pending-state", state_error + ) + if pending: + connector._fail_closed(where, error) + return + def enqueue_output(self, output: Any): """Prepares output from the worker and enqueues it to the worker_response_mq. If the output is an Exception, it is @@ -1044,6 +1067,9 @@ class WorkerProc: output = output.get_output() except Exception as e: logger.exception("Error getting async model runner output") + self._fail_ple_rpc_if_pending( + "async-output-copy", e + ) output = e if isinstance(output, Exception): @@ -1104,6 +1130,9 @@ class WorkerProc: if output_rank is None or self.rank == output_rank: self.handle_output(output) except Exception as e: + self._fail_ple_rpc_if_pending( + f"worker-rpc:{method}", e + ) # Notes have been introduced in python 3.11 if hasattr(e, "add_note"): e.add_note(traceback.format_exc()) diff --git a/vllm/v1/ple_offload/__init__.py b/vllm/v1/ple_offload/__init__.py new file mode 100644 index 0000000..208f01a --- /dev/null +++ b/vllm/v1/ple_offload/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/ple_offload/connector.py b/vllm/v1/ple_offload/connector.py new file mode 100644 index 0000000..52bc462 --- /dev/null +++ b/vllm/v1/ple_offload/connector.py @@ -0,0 +1,855 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Exchange PLE data between a GPU worker and the CPU-offload process.""" + +import os +import queue +import threading +from dataclasses import dataclass +from multiprocessing.reduction import ForkingPickler +from typing import Any + +import msgspec +import torch +from vllm.v1.ple_offload import rowdump as _rowdump +from vllm.v1.ple_offload.fatal import ple_fail_closed +from vllm.v1.ple_offload.progress import PleProgressMonitor +import torch.nn as nn +import zmq +from cuda.bindings import driver as cuda_driver + +from vllm.config import VllmConfig +from vllm.distributed.parallel_state import get_dp_group, get_pp_group, get_tp_group +from vllm.logger import init_logger +from vllm.model_executor.layers.ple_offload_layer import ( + CpuGpuSemaphore, + PleOffloadLayer, +) +from vllm.v1.ple_offload.protocol import ( + PleOffloadRegistration, + PleOffloadRequest, +) +from vllm.v1.utils import record_function_or_nullcontext +from vllm_ple_test_fault import _FAULT + +logger = init_logger(__name__) + +# CUDA registration alone does not own Python storage. Captured graphs retain +# raw addresses, so every successfully registered host allocation remains +# strongly referenced until process exit. +_PROCESS_LIFETIME_REGISTERED_STORAGE: list[torch.Tensor] = [] + + +def _unlink_filename_shared_storage( + tensor: torch.Tensor, + reserved_filenames: set[bytes], +) -> None: + """Reserve and unlink one already filename-shared CPU storage.""" + try: + storage = tensor.untyped_storage() + if storage.device.type != "cpu" or not storage.is_shared(): + return + try: + storage._get_shared_fd() + except RuntimeError: + # Filename-shared storages do not expose a shared fd. Only now is + # it safe to query the filename: doing so on ordinary or fd-shared + # storage would convert it to a new filename-backed allocation. + pass + else: + return + _, filename, _ = storage._share_filename_cpu_() + if not isinstance(filename, bytes): + filename = os.fsencode(filename) + path = os.path.join("/dev/shm", os.fsdecode(filename).lstrip("/")) + except Exception: + logger.warning( + "Failed to inspect a PLE shared-memory storage during teardown", + exc_info=True, + ) + return + + # RefcountedMapAllocator unlinks the name when its shared counter reaches + # zero. Reserve one count before manual unlink so later storage destruction + # cannot try to unlink the already-removed name. The filename identifies + # the shared counter and also deduplicates tensor aliases and repeated close. + if filename in reserved_filenames: + return + try: + storage._shared_incref() + except Exception: + logger.warning( + "Failed to reserve the PLE shared-memory counter for %s", + path, + exc_info=True, + ) + return + reserved_filenames.add(filename) + + try: + os.unlink(path) + except FileNotFoundError: + # Another teardown participant already removed the name. Retain the + # reservation because this allocator must still avoid a later unlink. + pass + except OSError as error: + # The name remains available, so restore the allocator's original + # counter and permit a later release attempt to retry. + try: + storage._shared_decref() + except Exception: + # The reservation state is uncertain after a failed private call. + # Keep it marked to avoid adding another reservation on retry. + logger.warning( + "Failed to roll back the PLE shared-memory counter for %s", + path, + exc_info=True, + ) + else: + reserved_filenames.discard(filename) + logger.warning("Failed to unlink PLE shared memory %s: %s", path, error) + + +@dataclass(frozen=True) +class _PendingPleOffloadRequest: + """Bind request metadata to its MRV2 D2H completion event.""" + + request: PleOffloadRequest + d2h_done_event: torch.cuda.Event + eligible: bool = False + transaction_id: str | None = None + correlation_id: str | None = None + + +def _cuda_check(result: Any, operation: str) -> Any: + """Check the ``(CUresult, ...)`` tuple returned by cuda-python calls.""" + error = result[0] if isinstance(result, tuple) else result + if error.value != 0: + raise RuntimeError(f"{operation} failed: {error}") + return result + + +class PleOffloadConnector: + """Connect a model runner V2 GPU worker to the PLE CPU worker.""" + + def __init__( + self, + vllm_config: VllmConfig, + model: nn.Module, + device: torch.device, + ipc_addr: str, + *, + input_ids_source: torch.Tensor, + query_start_loc_source: torch.Tensor, + ngram_context_source: torch.Tensor | None, + ) -> None: + self.device = device + self.rank = int(vllm_config.parallel_config.rank) + self.dp_rank = get_dp_group().rank_in_group + self.tp_rank = get_tp_group().rank_in_group + self.pp_rank = get_pp_group().rank_in_group + # Host-mapped pull design: per layer a shared host output buffer and a + # shared host int32 flag (element 0), filled/stored by the offload + # process with plain CPU writes. Page-locked + device-mapped here only + # after registration, because serialising for registration may move + # the shared storage to a new address. + self._host_output_bufs: dict[str, torch.Tensor] = {} + self._host_flags: dict[str, torch.Tensor] = {} + self._reserved_host_storage_filenames: set[bytes] = set() + self._gpu_output_bufs: dict[str, torch.Tensor] = {} + self._layers = self._setup_layers(vllm_config, model) + self._progress = PleProgressMonitor.from_environment( + rank=self.rank, + layer_names=tuple(self._layers), + ) + + # TP0 registers shared buffers with CUDA for asynchronous D2H copies. + scheduler_config = vllm_config.scheduler_config + self._input_ids_buf = torch.empty( + scheduler_config.max_num_batched_tokens, + dtype=torch.int32, + device="cpu", + ).share_memory_() + self._query_start_loc_buf = torch.empty( + scheduler_config.max_num_seqs + 1, + dtype=torch.int32, + device="cpu", + ).share_memory_() + self._ngram_context_buf = None + config = vllm_config.model_config.hf_text_config + ngram_context_len = int(config.ngram_size) - 1 + if ngram_context_len > 0: + self._ngram_context_buf = torch.empty( + scheduler_config.max_num_seqs, + ngram_context_len, + dtype=torch.int32, + device="cpu", + ).share_memory_() + + # Runner input allocations are address-stable, so bind them once and + # pass only batch sizes through the per-forward request queue. + self._input_ids_source = input_ids_source + self._query_start_loc_source = query_start_loc_source + self._ngram_context_source = ngram_context_source + self._validate_input_sources() + + self._pinned_input_buffers: list[torch.Tensor] = [] + # Device-mapped host outputs/flags referenced by raw address from + # captured graphs: never unregistered while the process lives + # (unregistering under a pending graph wait unmaps it). + self._mapped_host_buffers: list[torch.Tensor] = [] + # True while dummy/capture work (seeds, dummy forwards, capture + # cleanup resets) has run since the last real forward. Releases in + # that state complete synchronously, and the next real forward first + # passes a TP barrier, so no rank can publish a real request while + # another rank still has a dummy/cleanup reset pending. + # In this tree real-path kernel warm-ups run BEFORE + # capture and capture adds seeds/resets after them, so this must hold + # at every dummy->real transition, not only at the first real one. + self._after_dummy = True + self._request_queue: queue.Queue[_PendingPleOffloadRequest | None] = ( + queue.Queue(maxsize=vllm_config.max_concurrent_batches) + ) + self._dropped_events: list[torch.cuda.Event] = [] + self._request_thread: threading.Thread | None = None + self._request_thread_ready = threading.Event() + self._zmq_ctx: zmq.Context | None = None + self._registration_socket: zmq.Socket | None = None + self._d2h_event_pool: queue.Queue[torch.cuda.Event] | None = None + + try: + self._zmq_ctx = zmq.Context() + self._registration_socket = self._zmq_ctx.socket(zmq.PUSH) + self._registration_socket.connect(ipc_addr) + payload = self._serialize_registration(vllm_config) + + # Serialising moved the shared storages to their final addresses. + # Page-lock/map exactly those, THEN publish, so a received + # registration always means "storage usable". + with torch.accelerator.device_index(self.device.index): + self._map_host_outputs(vllm_config) + if self.tp_rank == 0: + self._pin_input_buffers() + self._send_registration(payload, ipc_addr) + + if self.tp_rank == 0: + with torch.accelerator.device_index(self.device.index): + self._d2h_event_pool = queue.Queue( + maxsize=vllm_config.max_concurrent_batches + ) + for _ in range(vllm_config.max_concurrent_batches): + self._d2h_event_pool.put_nowait(torch.cuda.Event()) + self._start_request_thread(ipc_addr) + except Exception: + self.close() + raise + + def _setup_layers( + self, + vllm_config: VllmConfig, + model: nn.Module, + ) -> dict[str, PleOffloadLayer]: + """Attach output buffers and semaphores to GPU PLE placeholders.""" + layers = { + name: module + for name, module in model.named_modules() + if isinstance(module, PleOffloadLayer) + } + if not layers: + raise RuntimeError( + "VLLM_PLE_CPU_OFFLOAD is enabled, but the model has no PleOffloadLayer" + ) + + config = vllm_config.model_config.hf_text_config + max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + for name, layer in layers.items(): + if vllm_config.load_config.load_format == "dummy": + layer.initialize_dummy_offload_metadata(self.device) + out_dim = layer.get_offload_output_dim(int(config.ple_embed_dim)) + out_dtype = layer.get_offload_output_dtype(vllm_config.model_config.dtype) + # The captured forward DMAs the host result into this device + # buffer after the host flag is set; consumers read it as before. + self._gpu_output_bufs[name] = torch.empty( + max_num_tokens, out_dim, dtype=out_dtype, device=self.device + ) + # Zero-initialised so dummy/profile forwards copy finite values. + self._host_output_bufs[name] = torch.zeros( + max_num_tokens, out_dim, dtype=out_dtype, device="cpu" + ).share_memory_() + # 16 x int32 = one 64-byte line; element 0 is the flag. + self._host_flags[name] = torch.zeros( + 16, dtype=torch.int32, device="cpu" + ).share_memory_() + return layers + + def _map_host_outputs(self, vllm_config: VllmConfig) -> None: + """Page-lock + device-map host outputs/flags and arm the placeholders.""" + for name, layer in self._layers.items(): + layer._ple_offload_name = name + host_out = self._host_output_bufs[name] + host_flag = self._host_flags[name] + for buffer in (host_out, host_flag): + self._register_host_buffer(buffer, device_map=True) + flag_dev_ptr = _cuda_check( + cuda_driver.cuMemHostGetDevicePointer(host_flag.data_ptr(), 0), + "cuMemHostGetDevicePointer(PLE host flag)", + )[1] + out_dev_ptr = _cuda_check( + cuda_driver.cuMemHostGetDevicePointer(host_out.data_ptr(), 0), + "cuMemHostGetDevicePointer(PLE host output)", + )[1] + row_bytes = host_out.shape[1] * host_out.element_size() + gpu_buf = self._gpu_output_bufs[name] + if ( + gpu_buf.shape[0] != host_out.shape[0] + or gpu_buf.shape[1] * gpu_buf.element_size() != row_bytes + ): + raise RuntimeError(f"PLE host/GPU buffer geometry mismatch for {name}") + # Stream memops use the flag's DEVICE pointer; the H2D copy uses + # this worker's HOST pointer (no raw addresses cross + # processes; the offload writes through its own mapping). + layer.setup_cross_process_offload( + gpu_buf, + CpuGpuSemaphore.from_host_mapped(host_flag, int(flag_dev_ptr)), + host_output_ptr=host_out.data_ptr(), + host_row_bytes=row_bytes, + ) + logger.info( + "PleOffload: host-mapped pull for %s (rows %d x %d B, flag dev " + "0x%x, out dev 0x%x, host==dev %s)", + name, + host_out.shape[0], + row_bytes, + int(flag_dev_ptr), + int(out_dev_ptr), + int(out_dev_ptr) == host_out.data_ptr(), + ) + + def _pin_input_buffers(self) -> None: + """Page-lock shared input allocations without replacing their storage.""" + buffers = [self._input_ids_buf, self._query_start_loc_buf] + if self._ngram_context_buf is not None: + buffers.append(self._ngram_context_buf) + for buffer in buffers: + self._register_host_buffer(buffer, device_map=False) + + def _register_host_buffer(self, buffer: torch.Tensor, device_map: bool) -> None: + """Page-lock (and optionally device-map) a shared CPU allocation.""" + if buffer.device.type != "cpu" or not buffer.is_shared(): + raise RuntimeError("PLE host buffers must be shared CPU tensors") + if not buffer.is_contiguous(): + raise RuntimeError("PLE host buffers must be contiguous") + flags = cuda_driver.CU_MEMHOSTREGISTER_PORTABLE + if device_map: + flags |= cuda_driver.CU_MEMHOSTREGISTER_DEVICEMAP + _cuda_check( + cuda_driver.cuMemHostRegister( + buffer.data_ptr(), + buffer.numel() * buffer.element_size(), + flags, + ), + "cuMemHostRegister(PLE host buffer)", + ) + _PROCESS_LIFETIME_REGISTERED_STORAGE.append(buffer) + if device_map: + self._mapped_host_buffers.append(buffer) + else: + self._pinned_input_buffers.append(buffer) + if not buffer.is_pinned(): + raise RuntimeError("CUDA did not page-lock a PLE host buffer") + + def _unpin_input_buffers(self) -> None: + """Release CUDA registrations after the request thread has stopped.""" + for buffer in reversed(self._pinned_input_buffers): + try: + _cuda_check( + cuda_driver.cuMemHostUnregister(buffer.data_ptr()), + "cuMemHostUnregister(PLE input buffer)", + ) + except RuntimeError: + logger.exception("Failed to unregister a PLE input buffer") + self._pinned_input_buffers.clear() + + def _serialize_registration(self, vllm_config: VllmConfig) -> Any: + """Build and serialise the registration (may move shared storages).""" + # Each GPU worker owns distinct host output buffers and flags, while + # TP0's shared inputs become the request source for its DP rank. + registration = PleOffloadRegistration( + worker_id=( + self.dp_rank * vllm_config.parallel_config.world_size + + vllm_config.parallel_config.rank + ), + tp_rank=self.tp_rank, + dp_rank=self.dp_rank, + host_output_buffers=dict(self._host_output_bufs), + host_flag_tensors=dict(self._host_flags), + input_ids_buf=self._input_ids_buf, + query_start_loc_buf=self._query_start_loc_buf, + ngram_context_buf=self._ngram_context_buf, + ) + + # ForkingPickler transmits tensors through shared-memory and CUDA IPC. + import torch.multiprocessing as torch_mp + + original_strategy = torch_mp.get_sharing_strategy() + torch_mp.set_sharing_strategy("file_system") + try: + payload = ForkingPickler.dumps(registration) + finally: + torch_mp.set_sharing_strategy(original_strategy) + return payload + + def _send_registration(self, payload: Any, ipc_addr: str) -> None: + """Publish the serialised registration (storages already mapped).""" + assert self._registration_socket is not None + self._registration_socket.send(payload) + + logger.info( + "PleOffload: registered %d PleOffloadLayer(s) " + "(dp_rank=%d, tp_rank=%d, ipc_addr=%s): %s", + len(self._layers), + self.dp_rank, + self.tp_rank, + ipc_addr, + sorted(self._layers), + ) + + def _start_request_thread(self, ipc_addr: str) -> None: + """Start the thread that publishes batches after inputs are ready.""" + self._request_thread = threading.Thread( + target=self._request_loop, + args=(ipc_addr,), + name=f"ple-offload-dp{self.dp_rank}", + daemon=True, + ) + self._request_thread.start() + if not self._request_thread_ready.wait(timeout=10): + raise RuntimeError("Timed out starting the PLE request thread") + + def _request_loop(self, ipc_addr: str) -> None: + """Wait for staged inputs, then notify the CPU worker.""" + socket: zmq.Socket | None = None + try: + if self._zmq_ctx is None: + raise RuntimeError("PLE ZMQ context closed before thread startup") + socket = self._zmq_ctx.socket(zmq.PUSH) + socket.connect(ipc_addr) + self._request_thread_ready.set() + while True: + request = self._request_queue.get() + if request is None: + return + self._process_request(request, socket) + except Exception: + logger.exception("PLE request thread failed") + os._exit(1) + finally: + self._request_thread_ready.set() + if socket is not None: + socket.close(linger=0) + + def _process_request( + self, pending: _PendingPleOffloadRequest, socket: zmq.Socket + ) -> None: + """Wait for one staged batch and publish its request.""" + request = pending.request + event_pool = self._d2h_event_pool + event = pending.d2h_done_event + assert event_pool is not None, "PLE D2H event pool is not initialized" + if _FAULT is not None and _FAULT.fires( + "block_input_completion", + role="gpu", + unit="connector-request", + eligible=pending.eligible, + tp_rank=self.tp_rank, + pp_rank=self.pp_rank, + dp_rank=self.dp_rank, + transaction_id=pending.transaction_id, + correlation_id=pending.correlation_id, + ): + threading.Event().wait() + + with torch.accelerator.device_index(self.device.index): + event.synchronize() + + payload = msgspec.msgpack.encode(request) + if _FAULT is None: + socket.send(payload) + else: + # Armed runs only: correlation metadata rides as a second frame. + metadata = msgspec.msgpack.encode( + { + "serving": pending.eligible, + "transaction": pending.transaction_id, + "correlation": pending.correlation_id, + } + ) + socket.send_multipart([payload, metadata]) + event_pool.put_nowait(event) + + def _validate_input_sources(self) -> None: + """Validate fixed runner sources against shared input buffers.""" + sources = [ + ("input_ids", self._input_ids_source, self._input_ids_buf), + ( + "query_start_loc", + self._query_start_loc_source, + self._query_start_loc_buf, + ), + ] + if (self._ngram_context_source is None) != (self._ngram_context_buf is None): + raise ValueError("PLE ngram_context source and buffer must match") + if self._ngram_context_source is not None: + assert self._ngram_context_buf is not None + sources.append( + ( + "ngram_context", + self._ngram_context_source, + self._ngram_context_buf, + ) + ) + + for name, source, buffer in sources: + if ( + source.device != self.device + or source.dtype != buffer.dtype + or source.ndim != buffer.ndim + or source.shape[0] < buffer.shape[0] + or source.shape[1:] != buffer.shape[1:] + ): + raise ValueError(f"PLE {name} source is incompatible") + + def _enqueue_cuda_inputs( + self, + request: PleOffloadRequest, + d2h_done_event: torch.cuda.Event, + ) -> None: + """Stage inputs on the model stream and record D2H completion.""" + with ( + record_function_or_nullcontext("ple_offload: enqueue_cuda_inputs"), + torch.accelerator.device_index(self.device.index), + ): + stream = torch.cuda.current_stream(self.device) + self._input_ids_buf[: request.num_tokens].copy_( + self._input_ids_source[: request.num_tokens], + non_blocking=True, + ) + self._query_start_loc_buf[: request.num_reqs + 1].copy_( + self._query_start_loc_source[: request.num_reqs + 1], + non_blocking=True, + ) + if self._ngram_context_buf is not None: + assert self._ngram_context_source is not None + self._ngram_context_buf[: request.num_reqs].copy_( + self._ngram_context_source[: request.num_reqs], + non_blocking=True, + ) + d2h_done_event.record(stream) + + def _launch( + self, + num_reqs: int, + num_tokens: int, + *, + serving: bool = False, + transaction_id: str | None = None, + correlation_id: str | None = None, + ) -> None: + """Stage or queue one batch for request publication.""" + # Inputs are replicated across TP ranks. One request per DP rank drives + # the CPU result fan-out to every registered TP output buffer. + if self.tp_rank != 0: + return + + request = PleOffloadRequest( + dp_rank=self.dp_rank, + num_tokens=num_tokens, + num_reqs=num_reqs, + ) + assert self._d2h_event_pool is not None, "PLE D2H event pool is not initialized" + try: + d2h_done_event = self._d2h_event_pool.get_nowait() + except queue.Empty as exc: + raise RuntimeError( + "PLE has more requests than configured concurrent batches" + ) from exc + self._enqueue_cuda_inputs(request, d2h_done_event) + if _FAULT is not None and _FAULT.fires( + "drop_request", + role="gpu", + unit="connector-request", + eligible=serving, + tp_rank=self.tp_rank, + pp_rank=self.pp_rank, + dp_rank=self.dp_rank, + transaction_id=transaction_id, + correlation_id=correlation_id, + ): + # Deliberately retain the event: the fault is loss after D2H, not + # premature event-pool reuse while that D2H may still be in flight. + self._dropped_events.append(d2h_done_event) + return + self._request_queue.put_nowait( + _PendingPleOffloadRequest( + request, + d2h_done_event, + eligible=serving, + transaction_id=transaction_id, + correlation_id=correlation_id, + ) + ) + + def _fail_closed(self, where: str, error: BaseException) -> None: + ple_fail_closed( + where=where, + rank=self.rank, + layer=",".join(self._layers), + sequence=self._progress.sequence, + error=error, + ) + + @property + def has_pending_completions(self) -> bool: + return self._progress.has_pending_completions + + def begin_forward(self, *, dummy_run: bool, num_tokens: int) -> int: + """Begin the bounded transaction before any runner fault boundary.""" + try: + return self._progress.begin( + startup=dummy_run, + num_tokens=num_tokens, + ) + except Exception as error: + self._fail_closed("prepare_forward", error) + + def fault_transaction(self, sequence: int) -> str: + """Read diagnostic state only when the caller has established arming.""" + return self._progress.describe_submission(sequence) + + def completion_is_pending(self, sequence: int) -> bool: + """Return exact-sequence incomplete completion state for cut 6.""" + return self._progress.completion_is_pending(sequence) + + def prepare_forward( + self, + num_reqs: int, + num_tokens: int, + dummy_run: bool, + *, + consumer_num_tokens: int | None = None, + serving: bool = False, + transaction_id: str | None = None, + correlation_id: str | None = None, + ) -> None: + """Prepare one PLE transaction; any ordinary exception is fatal.""" + try: + self._prepare_forward_impl( + num_reqs, + num_tokens, + dummy_run, + consumer_num_tokens=consumer_num_tokens, + serving=serving, + transaction_id=transaction_id, + correlation_id=correlation_id, + ) + except Exception as error: + self._fail_closed("prepare_forward", error) + + def _prepare_forward_impl( + self, + num_reqs: int, + num_tokens: int, + dummy_run: bool, + *, + consumer_num_tokens: int | None = None, + serving: bool = False, + transaction_id: str | None = None, + correlation_id: str | None = None, + ) -> None: + """Submit real inputs or satisfy the PLE wait for a dummy forward.""" + if consumer_num_tokens is None: + consumer_num_tokens = num_tokens + if int(num_tokens) != int(consumer_num_tokens): + raise RuntimeError( + "PLE request/consumer token-count mismatch: " + f"request={num_tokens}, consumer={consumer_num_tokens}" + ) + if not 0 <= int(num_tokens) <= self._input_ids_buf.shape[0]: + raise RuntimeError( + f"PLE token count {num_tokens} exceeds connector capacity " + f"{self._input_ids_buf.shape[0]}" + ) + if not 0 <= int(num_reqs) < self._query_start_loc_buf.shape[0]: + raise RuntimeError( + f"PLE request count {num_reqs} exceeds connector capacity " + f"{self._query_start_loc_buf.shape[0] - 1}" + ) + self._rowdump_n = None if dummy_run else num_tokens + if dummy_run: + self._signal_dummy_outputs_impl(num_tokens) + return + if self._after_dummy: + # Every dummy/cleanup reset on THIS rank already completed + # (synchronous releases); the barrier extends that to every + # stage-0 TP rank before TP0 can publish (CPU group, no GPU work; + # all TP ranks run the same forwards in the same order). + get_tp_group().barrier() + self._after_dummy = False + logger.info( + "PleOffload: dummy->real transition (dp_rank=%d, tp_rank=%d): " + "releases were synchronous, TP barrier passed", + self.dp_rank, + self.tp_rank, + ) + self._launch( + num_reqs, + num_tokens, + serving=serving, + transaction_id=transaction_id, + correlation_id=correlation_id, + ) + + def signal_dummy_outputs(self, num_tokens: int) -> None: + """Locally satisfy PLE waits for dummy and capture forwards.""" + self._progress.begin(startup=True, num_tokens=num_tokens) + self._signal_dummy_outputs_impl(num_tokens) + + def _signal_dummy_outputs_impl(self, num_tokens: int) -> None: + """Signal dummy outputs after the progress transaction is armed.""" + # Dummy and capture forwards do not send CPU requests, but every PLE + # placeholder still waits for a completed output semaphore. + self._after_dummy = True + stream = torch.cuda.current_stream(self.device) + for layer in self._layers.values(): + layer._gpu_output_buffer[:num_tokens].zero_() + layer._sem.signal(stream) + + def release_outputs( + self, + *, + serving: bool = False, + transaction_id: str | None = None, + correlation_id: str | None = None, + ) -> None: + """Release one PLE transaction; any ordinary exception is fatal.""" + try: + self._release_outputs_impl( + serving=serving, + transaction_id=transaction_id, + correlation_id=correlation_id, + ) + self._progress.complete() + except Exception as error: + self._fail_closed("release_outputs", error) + + def _release_outputs_impl( + self, + *, + serving: bool = False, + transaction_id: str | None = None, + correlation_id: str | None = None, + ) -> None: + """Mark GPU output buffers reusable after the model consumes them.""" + if _FAULT is not None and _FAULT.fires( + "exc_in_release", + role="gpu", + unit="connector-release", + eligible=serving, + tp_rank=self.tp_rank, + pp_rank=self.pp_rank, + dp_rank=self.dp_rank, + transaction_id=transaction_id, + correlation_id=correlation_id, + ): + raise RuntimeError("PLE test fault: exc_in_release") + n = getattr(self, "_rowdump_n", None) + self._rowdump_n = None + if _rowdump.enabled() and n is not None and not torch.cuda.is_current_stream_capturing(): + torch.cuda.current_stream(self.device).synchronize() + seq = getattr(self, "_rowdump_seq", 0) + self._rowdump_seq = seq + 1 + for layer in self._layers.values(): + _rowdump.write(f"gpu-tp{self.tp_rank}", 0, self.tp_rank, seq, + self._input_ids_source[:n], layer._gpu_output_buffer[:n]) + # Reset only after the consumer forward so the CPU worker cannot + # overwrite an output that a GPU PLE placeholder may still read. + stream = torch.cuda.current_stream(self.device) + for layer in self._layers.values(): + layer.release_offloaded_output(stream) + # Retain the serving deadline until the model stream has actually + # executed every reset above. The monitor thread is the sole query site. + completion_event = torch.cuda.Event() + completion_event.record(stream) + self._progress.track_completion(completion_event) + # After dummy/capture work (profile, warm-up sizes, capture seed, + # adaptive-verification dummy replays, capture cleanup): complete the + # reset before returning, so it cannot still be pending when a real + # request lets the offload publish. Real-forward releases stay async. + # Never inside a stream capture. + if self._after_dummy and not torch.cuda.is_current_stream_capturing(): + stream.synchronize() + + def release_host_storage(self) -> None: + """Unlink host-storage names while retaining mappings and tensors.""" + buffers = [ + *self._host_output_bufs.values(), + *self._host_flags.values(), + self._input_ids_buf, + self._query_start_loc_buf, + ] + if self._ngram_context_buf is not None: + buffers.append(self._ngram_context_buf) + for buffer in buffers: + _unlink_filename_shared_storage( + buffer, + self._reserved_host_storage_filenames, + ) + + def close(self) -> None: + """Stop transport and unlink names, retaining host mappings.""" + request_thread = self._request_thread + if request_thread is not None and request_thread.is_alive(): + try: + self._request_queue.put(None, timeout=5) + except queue.Full: + logger.error("Timed out stopping the PLE request thread") + request_thread.join(timeout=5) + if request_thread is not None and request_thread.is_alive(): + # The thread may still access the registered buffers or ZMQ context. + # It may also still race another process opening a shared name, so + # this deferred-cleanup branch must not unlink host storage. + logger.error("PLE request thread did not stop; deferring resource cleanup") + self._progress.close() + return + self._request_thread = None + + # Inputs, outputs, and flags remain registered and strongly referenced. + # A captured graph or delayed request can still contain their addresses; + # the CUDA driver releases registrations when the process exits. + registered = len(self._pinned_input_buffers) + len(self._mapped_host_buffers) + if registered: + logger.debug( + "PleOffload: %d registered host buffers retained until exit", + registered, + ) + self._d2h_event_pool = None + if self._registration_socket is not None: + self._registration_socket.close(linger=0) + self._registration_socket = None + if self._zmq_ctx is not None: + self._zmq_ctx.term() + self._zmq_ctx = None + self._progress.close() + + # __init__'s failure handler can reach here before the child has rebuilt + # a registration. Startup is already failing in that case; if a payload + # was sent, a late rebuild may also observe the removed name. Successful + # startup accepts every registration before an external close begins. + # + # NOTE: Graceful shutdown unlinks the PLE shared-memory names and the + # PLE ipc socket. The remaining 17 /tmp UUID sockets per boot are + # upstream vLLM VLLM_RPC_BASE_PATH behaviour; a tmpfiles/launch-profile + # sweep is recommended. + # Keep this last: unlinking preserves every existing CPU/CUDA mapping, + # but no transport user may still need to open a name at this point. + self.release_host_storage() diff --git a/vllm/v1/ple_offload/fatal.py b/vllm/v1/ple_offload/fatal.py new file mode 100644 index 0000000..a80cad4 --- /dev/null +++ b/vllm/v1/ple_offload/fatal.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fail-closed helpers for the PLE single-slot transport.""" + +import functools +import os +import time +import traceback +from collections.abc import Callable +from typing import Any, NoReturn, TypeVar + +_F = TypeVar("_F", bound=Callable[..., Any]) + +_TRUSTED_ERROR_LABELS = ( + (RuntimeError, "RuntimeError"), + (TimeoutError, "TimeoutError"), + (ValueError, "ValueError"), + (OSError, "OSError"), + (EOFError, "EOFError"), + (BrokenPipeError, "BrokenPipeError"), + (AssertionError, "AssertionError"), + (KeyError, "KeyError"), + (IndexError, "IndexError"), + (TypeError, "TypeError"), + (MemoryError, "MemoryError"), + (KeyboardInterrupt, "KeyboardInterrupt"), +) + + +def ple_fail_closed( + *, + where: str, + rank: int, + layer: str, + sequence: int, + error: BaseException, +) -> NoReturn: + """Emit one best-effort nonblocking diagnostic and terminate immediately.""" + try: + os.set_blocking(2, False) + error_type = type(error) + error_name = "Exception" + trusted_error = False + for trusted_type, trusted_name in _TRUSTED_ERROR_LABELS: + if error_type is trusted_type: + error_name = trusted_name + trusted_error = True + break + + safe_where = where[:128] if type(where) is str else "unknown" + safe_layer = layer[:256] if type(layer) is str else "unknown" + safe_rank = rank if type(rank) is int else -1 + safe_sequence = sequence if type(sequence) is int else -1 + detail_suffix = "" + if trusted_error: + args = BaseException.args.__get__(error, error_type) + if ( + type(args) is tuple + and len(args) == 1 + and type(args[0]) is str + ): + detail = args[0][:512].replace("\n", "\\n") + detail_suffix = f":{detail}" + message = ( + f"PLE fatal: where={safe_where} rank={safe_rank} layer={safe_layer} " + f"sequence={safe_sequence} error={error_name}{detail_suffix} " + f"pid={os.getpid()} " + f"monotonic={time.monotonic():.9f} wall={time.time():.9f}\n" + ) + os.write(2, message.encode("utf-8", "backslashreplace")[:2048]) + except BaseException: + pass + # THE TRACEBACK, unconditionally, in its own try so it cannot cost us the + # line above. + # + # This path os._exit(1)s, which runs no handlers and flushes nothing, so + # whatever is not written here is gone. The line above names the TYPE only + # when the type is one of _TRUSTED_ERROR_LABELS and flattens everything + # else to the literal string "Exception" -- and the flattened set includes + # the two that matter most here: torch.AcceleratorError (a RuntimeError + # SUBCLASS, so the `is` comparison misses it) and Triton CompilationError. + # An earlier boot died exactly that way: both PP0 ranks + # reported "error=Exception" and there was no traceback anywhere in the + # log, in any rank, so the cause could not be named at all and the next + # step was a re-run rather than a fix. + # + # A fail-closed path is where evidence matters most, not least: it fires + # once, it kills the process, and it is the only witness. Printing the type + # of a wrapped exception without the frames that raised it is not a + # diagnostic. + try: + # Three-arg form: it does not depend on being inside an active except + # block, and this function is called from one frame further out than + # the raise. + tb = "".join( + traceback.format_exception(type(error), error, error.__traceback__) + ) + raw = tb.encode("utf-8", "backslashreplace") + # Truncate from the END, not the start. format_exception puts the + # frames first and the exception TYPE AND MESSAGE last, so a head-side + # cut throws away the one line that names what went wrong and keeps the + # import-time frames nobody needs. A long chained traceback is exactly + # the case where that matters. + clipped = len(raw) > 8192 + raw = raw[-8192:] + os.write( + 2, + b"PLE fatal traceback (tail) follows:\n" if clipped + else b"PLE fatal traceback follows:\n", + ) + os.write(2, raw) + if not raw.endswith(b"\n"): + os.write(2, b"\n") + except BaseException: + # Never silent HERE of all places: without this the reader cannot tell + # "the wrapper predates the traceback fix" from "the traceback itself + # could not be produced", and those call for different next steps. + # Still guarded, because os._exit below must run whatever happens: an + # escape from this handler would leave the worker alive in the state + # the fail-closed path exists to end. + try: + os.write(2, b"PLE fatal: traceback unavailable\n") + except BaseException: + pass + os._exit(1) + + +def fail_ple_execute_model(func: _F) -> _F: + """Fail a non-shutdown worker-step escape only on PLE-owning ranks. + + This includes ``sample_tokens``, the second half of the split execute-model + step. A rank without a PLE connector rethrows the original exception. + """ + + @functools.wraps(func) + def wrapped(self: Any, *args: Any, **kwargs: Any) -> Any: + try: + return func(self, *args, **kwargs) + except SystemExit: + # WorkerProc's signal handler uses SystemExit for clean shutdown. + raise + except BaseException as error: + connector = getattr(self, "_ple_offload_connector", None) + if connector is None: + model_runner = getattr(self, "model_runner", None) + connector = getattr( + model_runner, "_ple_offload_connector", None + ) + if connector is None: + raise + connector._fail_closed(func.__name__, error) + + return wrapped # type: ignore[return-value] diff --git a/vllm/v1/ple_offload/progress.py b/vllm/v1/ple_offload/progress.py new file mode 100644 index 0000000..bf88291 --- /dev/null +++ b/vllm/v1/ple_offload/progress.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Independent progress deadline for PLE transactions.""" + +import math +import os +import threading +import time +from collections import deque +from dataclasses import dataclass +from typing import Any + +from vllm.v1.ple_offload.fatal import ple_fail_closed + +_DEFAULT_SERVING_DEADLINE_S = 300.0 +_DEFAULT_STARTUP_DEADLINE_S = 1800.0 +_MAX_PROGRESS_DEADLINE_S = 86400.0 +_POLL_INTERVAL_S = 0.1 + + +def _validate_deadline(value: float, name: str) -> float: + value = float(value) + if ( + not math.isfinite(value) + or value <= 0 + or value > _MAX_PROGRESS_DEADLINE_S + ): + raise ValueError( + f"{name} must be finite, greater than zero, and no greater than " + f"{_MAX_PROGRESS_DEADLINE_S:g} seconds" + ) + return value + + +def _environment_deadline(name: str, default: float) -> float: + value = float(os.environ.get(name, str(default))) + return _validate_deadline(value, name) + + +@dataclass(frozen=True) +class _Submission: + sequence: int + started: float + startup: bool + num_tokens: int + deadline: float + depth: int + + +@dataclass(frozen=True) +class _PendingCompletion: + sequence: int + started: float + startup: bool + num_tokens: int + deadline: float + event: Any + + +class PleProgressMonitor: + """Bound host submission and the resulting ordered GPU completions.""" + + def __init__( + self, + *, + rank: int, + layer_names: tuple[str, ...], + serving_deadline_s: float, + startup_deadline_s: float, + ) -> None: + serving_deadline_s = _validate_deadline( + serving_deadline_s, "serving_deadline_s" + ) + startup_deadline_s = _validate_deadline( + startup_deadline_s, "startup_deadline_s" + ) + self._rank = rank + self._layer = ",".join(layer_names) + self._serving_deadline_s = serving_deadline_s + self._startup_deadline_s = startup_deadline_s + self._condition = threading.Condition() + self._submission: _Submission | None = None + self._pending: deque[_PendingCompletion] = deque() + self._sequence = 0 + self._stopping = False + self._thread = threading.Thread( + target=self._run, + name=f"ple-progress-rank{rank}", + daemon=True, + ) + self._thread.start() + + @classmethod + def from_environment( + cls, *, rank: int, layer_names: tuple[str, ...] + ) -> "PleProgressMonitor": + return cls( + rank=rank, + layer_names=layer_names, + serving_deadline_s=_environment_deadline( + "VLLM_PLE_PROGRESS_DEADLINE_S", + _DEFAULT_SERVING_DEADLINE_S, + ), + startup_deadline_s=_environment_deadline( + "VLLM_PLE_STARTUP_DEADLINE_S", + _DEFAULT_STARTUP_DEADLINE_S, + ), + ) + + @property + def sequence(self) -> int: + with self._condition: + return self._sequence + + @property + def has_pending_completions(self) -> bool: + with self._condition: + return bool(self._pending) + + def describe_submission(self, sequence: int) -> str: + """Return the armed transaction descriptor for the active sequence.""" + with self._condition: + active = self._submission + if active is None or active.sequence != sequence: + raise RuntimeError("PLE fault correlation has no active submission") + return ( + f"sequence:{active.sequence}" + f"@begin:{active.started:.9f}" + ) + + def completion_is_pending(self, sequence: int) -> bool: + """Whether this serving sequence has a retained, incomplete event. + + This method is called only behind the armed fault guard. An event that + has completed on the GPU but has not yet been dequeued is not counted. + """ + with self._condition: + pending = next( + ( + item + for item in self._pending + if item.sequence == sequence and not item.startup + ), + None, + ) + if pending is None: + return False + return not bool(pending.event.query()) + + def begin(self, *, startup: bool, num_tokens: int) -> int: + now = time.monotonic() + budget = ( + self._startup_deadline_s if startup else self._serving_deadline_s + ) + with self._condition: + active = self._submission + if active is not None: + if not startup or not active.startup: + raise RuntimeError( + "Overlapping PLE host submissions violate the serialized " + "single-slot protocol" + ) + sequence = active.sequence + self._submission = _Submission( + sequence=sequence, + started=active.started, + startup=True, + num_tokens=num_tokens, + deadline=now + budget, + depth=active.depth + 1, + ) + else: + self._sequence += 1 + sequence = self._sequence + self._submission = _Submission( + sequence=sequence, + started=now, + startup=startup, + num_tokens=num_tokens, + deadline=now + budget, + depth=1, + ) + self._condition.notify_all() + return sequence + + def track_completion(self, event: Any) -> None: + """Retain a release event until the monitor observes GPU completion.""" + if not callable(getattr(event, "query", None)): + raise TypeError("PLE completion event must provide query()") + with self._condition: + active = self._submission + if active is None: + raise RuntimeError("PLE completion has no matching submission") + self._pending.append( + _PendingCompletion( + sequence=active.sequence, + started=active.started, + startup=active.startup, + num_tokens=active.num_tokens, + deadline=active.deadline, + event=event, + ) + ) + self._condition.notify_all() + + def complete(self) -> None: + with self._condition: + active = self._submission + if active is None or active.depth <= 0: + raise RuntimeError("PLE release has no matching transaction") + if active.depth == 1: + self._submission = None + else: + self._submission = _Submission( + sequence=active.sequence, + started=active.started, + startup=active.startup, + num_tokens=active.num_tokens, + deadline=active.deadline, + depth=active.depth - 1, + ) + self._condition.notify_all() + + def close(self) -> None: + with self._condition: + self._stopping = True + self._condition.notify_all() + self._thread.join(timeout=1) + + def _run(self) -> None: + try: + self._monitor() + except BaseException as error: + ple_fail_closed( + where="progress-monitor-thread", + rank=self._rank, + layer=self._layer, + sequence=self._sequence, + error=error, + ) + + def _monitor(self) -> None: + while True: + with self._condition: + if self._stopping: + return + submission = self._submission + pending = self._pending[0] if self._pending else None + if submission is None and pending is None: + self._condition.wait(timeout=_POLL_INTERVAL_S) + continue + + # CUDA completion is observed only here. Never synchronize and never + # use the host flag as evidence that the queued reset completed. + if pending is not None and pending.event.query(): + with self._condition: + if self._pending and self._pending[0] is pending: + self._pending.popleft() + self._condition.notify_all() + continue + + now = time.monotonic() + if pending is not None and now >= pending.deadline: + with self._condition: + if not self._pending or self._pending[0] is not pending: + continue + mode = "startup" if pending.startup else "serving" + ple_fail_closed( + where=( + f"progress-deadline:{mode}:gpu-completion:" + f"tokens={pending.num_tokens}" + ), + rank=self._rank, + layer=self._layer, + sequence=pending.sequence, + error=TimeoutError("PLE GPU completion made no bounded progress"), + ) + + if submission is not None and now >= submission.deadline: + with self._condition: + if self._submission is not submission: + continue + mode = "startup" if submission.startup else "serving" + ple_fail_closed( + where=( + f"progress-deadline:{mode}:host-submission:" + f"tokens={submission.num_tokens}" + ), + rank=self._rank, + layer=self._layer, + sequence=submission.sequence, + error=TimeoutError("PLE host submission made no bounded progress"), + ) + + deadlines = [] + if pending is not None: + deadlines.append(pending.deadline) + if submission is not None: + deadlines.append(submission.deadline) + timeout = _POLL_INTERVAL_S + if deadlines: + timeout = min( + timeout, + max(0.0, min(deadlines) - time.monotonic()), + ) + with self._condition: + if self._stopping: + return + if ( + self._submission is not submission + or (self._pending[0] if self._pending else None) is not pending + ): + continue + self._condition.wait(timeout=timeout) diff --git a/vllm/v1/ple_offload/protocol.py b/vllm/v1/ple_offload/protocol.py new file mode 100644 index 0000000..87eba13 --- /dev/null +++ b/vllm/v1/ple_offload/protocol.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""IPC message definitions for PLE CPU offload (host-mapped pull design). + +The offload process owns no CUDA context. Every GPU worker owns, per PLE +layer, a shared host output buffer and a shared host int32 flag, both +page-locked and device-mapped in the worker's context. The offload process +waits for ``flag == 0``, writes the layer result into the host buffer and +stores ``flag = 1`` with plain CPU stores; the GPU's captured forward waits on +the host flag, DMAs host -> GPU on the model stream, and ``release_outputs`` +resets the flag after the forward. +""" + +from dataclasses import dataclass + +import msgspec +import torch + +# --------------------------------------------------------------------------- +# IPC message dataclasses +# --------------------------------------------------------------------------- + + +@dataclass +class PleOffloadRegistration: + """Sent once from each GPU worker during offload setup.""" + + worker_id: int + tp_rank: int + dp_rank: int + # Shared CPU tensors per PLE layer, filled by the offload process. + host_output_buffers: dict[str, torch.Tensor] + # Shared CPU int32 tensors per PLE layer; element 0 is the ready flag. + host_flag_tensors: dict[str, torch.Tensor] + # CPU tensors are allocated in shared memory and registered once. + input_ids_buf: torch.Tensor + query_start_loc_buf: torch.Tensor + ngram_context_buf: torch.Tensor | None + + +@dataclass +class PleOffloadRequest: + """Sent by each DP rank's TP rank zero at every inference step.""" + + dp_rank: int + num_tokens: int + num_reqs: int + + +_PLE_OFFLOAD_REQUEST_DECODER = msgspec.msgpack.Decoder(PleOffloadRequest) diff --git a/vllm/v1/ple_offload/worker.py b/vllm/v1/ple_offload/worker.py new file mode 100644 index 0000000..43caf80 --- /dev/null +++ b/vllm/v1/ple_offload/worker.py @@ -0,0 +1,1298 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Dedicated CPU-offload process for PLE embedding layers. + +This module implements a standalone process that: +1. Loads only the :class:`PleOffloadLayer` weights into CPU memory. +2. Accepts per-step computation requests from GPU worker processes. +3. Runs ``forward_impl()`` on CPU, writes the rows into every TP worker's + shared host output buffer for the requesting DP rank and sets that + worker's shared host flag (host-mapped pull design: this process holds no + CUDA context; the GPU waits on the host flag and DMAs host -> GPU inside + its captured forward). + +The TP workers within one DP rank receive identical inputs, so the CPU result +is computed once per DP rank and fanned out to all of its TP ranks. + +Class structure mirrors the GPU worker pattern in multiproc_executor.py: + + PleOffloadWorkerHandle -- handle held by the spawning GPU worker + PleOffloadWorker -- process lifecycle and READY handshake + PleOffloadRunner -- owns weights and serves inference requests +""" + +import contextlib +import os +import pickle +import select +import signal +import socket +import struct +import subprocess +import sys +import tempfile +import threading +import time +from collections.abc import Iterable +from dataclasses import dataclass +from multiprocessing.connection import Connection +from typing import Any, cast + +import msgspec +import torch +import torch.distributed as dist +import zmq + +import vllm.envs as envs +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.distributed.parallel_state import ( + ensure_model_parallel_initialized, + init_distributed_environment, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.ple_offload_layer import ( + PleOffloadLayer, + mark_as_offload_worker, +) +from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.model_executor.model_loader.utils import ( + initialize_model, + process_weights_after_loading, +) +from vllm.model_executor.model_loader.weight_utils import initialize_dummy_weights +from vllm.utils.system_utils import decorate_logs +from vllm.utils.torch_utils import set_default_torch_dtype +from vllm.v1.ple_offload.protocol import ( + _PLE_OFFLOAD_REQUEST_DECODER, + PleOffloadRegistration, + PleOffloadRequest, +) +from vllm.v1.utils import record_function_or_nullcontext +from vllm_ple_test_fault import ( + _FAULT, + _FaultReconstructionPayload, +) + +logger = init_logger(__name__) + +# Publication fence for the host-mapped pull: a locked read-modify-write (the +# uncontended lock acquire) is a full barrier on x86 and also drains +# write-combining buffers, so every row store of a result is globally visible +# before the flag store that tells the GPU to DMA it (do not rely +# on "memcpy uses no non-temporal stores at this size"). +_PUBLISH_FENCE = threading.Lock() + + +def _unlink_ipc_socket(ipc_addr: str) -> None: + """Best-effort unlink of the exact filesystem IPC endpoint PLE owns.""" + if not ipc_addr.startswith("ipc://"): + return + path = ipc_addr.removeprefix("ipc://") + # PLE receives an absolute UUID path from get_open_zmq_ipc_path(). Avoid + # treating abstract or relative ZMQ addresses as filesystem ownership. + if not os.path.isabs(path): + return + try: + os.unlink(path) + except FileNotFoundError: + pass + except OSError as error: + logger.warning("Failed to unlink PLE IPC socket %s: %s", path, error) + + +def _publish_fence() -> None: + with _PUBLISH_FENCE: + pass + + +def _offload_child_environment() -> dict[str, str]: + """Build the offload environment without mutating the parent.""" + child_env = dict(os.environ) + child_env["VLLM_DISABLE_SHARED_EXPERTS_STREAM"] = "1" + child_env.setdefault("CUTLASS_CUDA_DEVICE_ID", "0") + child_env.pop("VLLM_PP_LAYER_PARTITION", None) + + home_device = child_env.get("VLLM_PLE_OFFLOAD_HOME_DEVICE", "").strip() + if home_device: + parent_visible = child_env.get("CUDA_VISIBLE_DEVICES", "").strip() + try: + physical = ( + parent_visible.split(",")[int(home_device)].strip() + if parent_visible + else home_device + ) + except (IndexError, ValueError) as error: + raise ValueError("Invalid VLLM_PLE_OFFLOAD_HOME_DEVICE") from error + child_env["CUDA_VISIBLE_DEVICES"] = physical + return child_env + + +def _prepare_cuda_free_process() -> None: + """Keep the PLE offload process free of any CUDA context. + + The host-mapped pull design issues no CUDA work here, yet model + construction reached CUDA incidentally (found by the offload-startup + replays): + 1. ``VLLM_DISABLE_SHARED_EXPERTS_STREAM=1``: each MoE layer's shared-experts + runner calls ``aux_stream()`` -> ``torch.cuda.Stream()``, whose C++ + constructor creates the primary context without calling + ``torch.cuda._lazy_init`` (``torch.cuda.is_initialized()`` stays False). + This is the context the replay caught (stack dumped inside + ``Stream.__new__``). This process never runs an MoE layer, so it never + needs the stream. vLLM reads the variable when the layer is built; this + process does not enable vLLM's env cache. + 2. ``CUTLASS_CUDA_DEVICE_ID``: CUTLASS python (``cutlass_cppgen``) runs + ``initialize_cuda_context()`` -> ``cudart.cudaFree(0)`` at import unless + it is set. Not observed on this import path; kept as a cheap defence. + 3. The torch guard (``_install_cuda_guard``). + """ + os.environ["VLLM_DISABLE_SHARED_EXPERTS_STREAM"] = "1" + os.environ.setdefault("CUTLASS_CUDA_DEVICE_ID", "0") + _install_cuda_guard() + + +def _install_cuda_guard() -> None: + """Make torch-level CUDA initialization raise in the PLE offload process. + + An offload-startup check found the host-mapped pull still got a 322 MiB context + on its home card: importing the Qwen4Exp model imports the vendored FLA + ops, whose utils call triton's ``get_current_target()`` at import time, + i.e. ``torch.cuda.current_device()`` -> lazy init. FLA catches + ``RuntimeError`` and falls back to "cpu" (and then skips its Hopper + probe), so raising here keeps this process context-free. Any other call + site fails loudly at offload startup with its stack logged, instead of + silently taking VRAM on a stage-0 card. The GPU workers are untouched. + + torch also has C++ entry points that create the context without calling + ``torch.cuda._lazy_init``: stream construction (``torch.cuda.Stream``) and + the pinned host allocator (``Tensor.pin_memory``). Both raise here too; + this process needs neither (all of its buffers are plain CPU memory). + ``torch.cuda.is_initialized()`` only reflects the Python-level path, so + the per-PID GPU sampler, not that flag, is the authority on contexts. + """ + import traceback + + if torch.cuda.is_initialized(): + # Fail closed: a context that already exists would stay on a + # stage-0 card for the whole run. + raise RuntimeError( + "CUDA is already initialized in the PLE offload process before its " + "guard; the host-mapped pull offload must stay CUDA-free" + ) + + def _refuse(what: str) -> None: + stack = "".join(traceback.format_stack(limit=12)[:-2]) + logger.info( + "PLE offload CUDA guard: blocked a CUDA initialization attempt " + "(%s; the caller must fall back to CPU or fail); call site:\n%s", + what, + stack, + ) + raise RuntimeError( + f"CUDA is disabled in the PLE offload process (host-mapped pull): {what}" + ) + + def _blocked_lazy_init() -> None: + _refuse("torch.cuda lazy init") + + def _blocked_stream(cls, *args, **kwargs): + _refuse("torch.cuda.Stream construction") + + def _blocked_pin_memory(self, *args, **kwargs): + _refuse("Tensor.pin_memory") + + torch.cuda._lazy_init = _blocked_lazy_init # type: ignore[assignment] + torch.cuda.Stream.__new__ = staticmethod(_blocked_stream) # type: ignore[assignment] + torch.Tensor.pin_memory = _blocked_pin_memory # type: ignore[assignment] + logger.info( + "PLE offload CUDA guard installed (torch CUDA lazy init, stream " + "construction and pin_memory raise)" + ) + + +@contextlib.contextmanager +def _platform_device_is_meta(): + """During the offload's meta-model construction, report the meta device. + + Some layer constructors pass an explicit ``current_platform.current_device()`` + (e.g. the GDN output norm in qwen_gdn_linear_attn.py), which overrides the + ``torch.device("meta")`` context and would allocate on CUDA. In the offload + process nothing but the CPU PLE layers is ever materialized, so those + tensors belong on meta too. The patch is an instance attribute removed on + exit, restoring the platform's own method. + """ + from vllm.platforms import current_platform + + current_platform.current_device = lambda *args, **kwargs: torch.device("meta") + try: + yield + finally: + with contextlib.suppress(AttributeError): + del current_platform.current_device + + +# Bounded wait for the GPU to release a result (seconds). A timed-out ownership +# transfer is fatal: the offload process raises and exits, never overwrites. +_PLE_WAIT_TIMEOUT_S = float(os.environ.get("VLLM_PLE_OFFLOAD_WAIT_TIMEOUT", "120")) + + +@dataclass +class PleOffloadOutputTarget: + """Host output destination and ready flag for one TP worker. + + Both are shared CPU tensors owned by the GPU worker, page-locked and + device-mapped in that worker's context. This process only reads/writes + them with plain CPU loads/stores (no CUDA): wait for ``flag == 0`` (the GPU + released the previous result), write the rows, then store ``flag = 1``. + """ + + tp_rank: int + host_output_buffer: torch.Tensor + host_flag: torch.Tensor # int32, element 0 is the flag + + +@dataclass +class PleOffloadInputBuffers: + """Shared-memory input buffers registered for one DP rank.""" + + input_ids_buf: torch.Tensor # int32 (max_num_tokens,) + query_start_loc_buf: torch.Tensor # int32 (max_num_reqs + 1,) + ngram_context_buf: torch.Tensor | None # int32 (max_num_reqs, ngram_context_len) + + +@dataclass(frozen=True) +class _ReceivedPleOffloadRequest: + request: PleOffloadRequest + serving: bool + transaction_id: str | None + correlation_id: str | None + + +def _decode_offload_request( + frames: list[bytes], +) -> _ReceivedPleOffloadRequest: + if len(frames) != 2: + raise RuntimeError("PLE offload request must contain two frames") + request = _PLE_OFFLOAD_REQUEST_DECODER.decode(frames[0]) + metadata = msgspec.msgpack.decode(frames[1]) + if type(metadata) is not dict or set(metadata) != { + "serving", + "transaction", + "correlation", + }: + raise RuntimeError("PLE offload request metadata is invalid") + serving = metadata["serving"] + transaction = metadata["transaction"] + correlation = metadata["correlation"] + if type(serving) is not bool: + raise RuntimeError("PLE offload serving metadata is invalid") + if transaction is not None and type(transaction) is not str: + raise RuntimeError("PLE offload transaction metadata is invalid") + if correlation is not None and type(correlation) is not str: + raise RuntimeError("PLE offload correlation metadata is invalid") + return _ReceivedPleOffloadRequest( + request=request, + serving=serving, + transaction_id=transaction, + correlation_id=correlation, + ) + + +def _receive_offload_request( + socket: zmq.Socket, flags: int = 0 +) -> _ReceivedPleOffloadRequest | PleOffloadRequest: + """Disarmed: the unchanged single-frame request. Armed: request + metadata.""" + if _FAULT is None: + return _PLE_OFFLOAD_REQUEST_DECODER.decode(socket.recv(flags)) + return _decode_offload_request(socket.recv_multipart(flags)) + + +def _send_framed_payload(fd: int, payload: bytes, deadline: float) -> None: + """Send one launch payload without blocking beyond an absolute deadline.""" + frame = memoryview(struct.pack("!Q", len(payload)) + payload) + offset = 0 + os.set_blocking(fd, False) + try: + while offset < len(frame): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("PLE launch payload transmission timed out") + _, writable, _ = select.select([], [fd], [], remaining) + if not writable: + raise TimeoutError("PLE launch payload transmission timed out") + try: + sent = os.write(fd, frame[offset:]) + except BlockingIOError: + continue + if sent <= 0: + raise BrokenPipeError("PLE launch control socket closed") + offset += sent + if ( + _FAULT is not None + and _FAULT.cut == "launch_park_before_payload_read" + ): + _FAULT.acknowledge_launch_send( + bytes_sent=offset, + frame_size=len(frame), + ) + finally: + with contextlib.suppress(OSError): + os.set_blocking(fd, True) + + +_PLE_READY_MAX_PAYLOAD_BYTES = 1 << 20 + +class _BoundedStartupResponseWriter: + """Preflight startup responses against the parent receiver's frame cap.""" + + def __init__(self, writer: Connection) -> None: + self._writer = writer + + def send(self, message: Any) -> None: + payload = pickle.dumps(message, protocol=pickle.HIGHEST_PROTOCOL) + if len(payload) > _PLE_READY_MAX_PAYLOAD_BYTES: + payload = pickle.dumps( + { + "status": "FAILURE", + "error": "PLE startup response exceeded size limit", + }, + protocol=pickle.HIGHEST_PROTOCOL, + ) + if len(payload) > _PLE_READY_MAX_PAYLOAD_BYTES: + raise RuntimeError( + "Bounded PLE startup failure response exceeds its frame cap" + ) + self._writer.send_bytes(payload) + + def close(self) -> None: + self._writer.close() + + +def _read_exact_before_deadline( + fd: int, + size: int, + deadline: float, +) -> bytes: + """Read exactly ``size`` bytes without crossing an absolute deadline.""" + chunks = bytearray() + while len(chunks) < size: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + "PLE offload worker did not become ready before its " + "absolute launch deadline" + ) + readable, _, _ = select.select([fd], [], [], remaining) + if not readable: + raise TimeoutError( + "PLE offload worker did not become ready before its " + "absolute launch deadline" + ) + try: + chunk = os.read(fd, size - len(chunks)) + except BlockingIOError: + continue + if not chunk: + raise EOFError("PLE startup response was truncated") + chunks.extend(chunk) + return bytes(chunks) + + +def _recv_bounded_connection_message( + reader: Connection, + deadline: float, +) -> Any: + """Receive one bounded multiprocessing Connection frame by ``deadline``.""" + fd = reader.fileno() + os.set_blocking(fd, False) + try: + header = _read_exact_before_deadline(fd, 4, deadline) + size = struct.unpack("!i", header)[0] + if size == -1: + size = struct.unpack( + "!Q", _read_exact_before_deadline(fd, 8, deadline) + )[0] + elif size < 0: + raise RuntimeError("PLE startup response has an invalid frame size") + if size > _PLE_READY_MAX_PAYLOAD_BYTES: + raise RuntimeError( + f"PLE startup response exceeds " + f"{_PLE_READY_MAX_PAYLOAD_BYTES} bytes" + ) + frame = _read_exact_before_deadline(fd, size, deadline) + finally: + with contextlib.suppress(OSError): + os.set_blocking(fd, True) + try: + return pickle.loads(frame) + except BaseException as error: + raise RuntimeError( + "PLE startup response is not a valid serialized message" + ) from error + + +@dataclass +class PleOffloadWorkerHandle: + """Resources owned by the GPU worker that spawned the offload process.""" + + proc: subprocess.Popen + death_writer: Connection | None + ready_pipe_reader: Connection | None + launch_deadline: float + ipc_addr: str | None = None + + def close(self, *, graceful: bool = True) -> None: + """Release all process resources. Safe to call more than once.""" + if self.ready_pipe_reader is not None: + self.ready_pipe_reader.close() + self.ready_pipe_reader = None + if graceful and self.proc.poll() is None: + # Keep parent-death supervision connected while SIGTERM lets + # proc_main run its shutdown event and finally cleanup. + self.proc.terminate() + try: + self.proc.wait(timeout=5) + except subprocess.TimeoutExpired: + pass + if self.death_writer is not None: + self.death_writer.close() + self.death_writer = None + if not graceful and self.proc.poll() is None: + # Expired launch: retain the immediate EOF -> TERM path. + self.proc.terminate() + try: + self.proc.wait(timeout=5) + except subprocess.TimeoutExpired: + pass + if self.proc.poll() is None: + self.proc.kill() + try: + self.proc.wait(timeout=5) + except subprocess.TimeoutExpired: + pass + # The child normally removes the endpoint in proc_main's finally. + # Once it has exited, this exact PLE-owned address is safe to remove as + # a fallback for SIGKILL or an interrupted child cleanup. + if self.proc.poll() is not None and self.ipc_addr is not None: + _unlink_ipc_socket(self.ipc_addr) + + +def _init_offload_distributed() -> None: + """Initialize the single-rank Gloo world required by TP-aware layers.""" + if dist.is_initialized(): + return + + # VocabParallelEmbedding reads the TP process group during construction. + # The offload process owns the full embedding table, so it uses an isolated + # TP1/PP1 Gloo world and never joins the GPU workers' NCCL groups. + store_dir = tempfile.mkdtemp(prefix="vllm_ple_offload_") + init_distributed_environment( + world_size=1, + rank=0, + distributed_init_method=f"file://{store_dir}/store", + local_rank=0, + backend="gloo", + ) + # initialize_model_parallel reads the active VllmConfig in the current + # vLLM version. Explicitly configure DP1/TP1/PP1 to match the isolated + # world, regardless of any DP environment variables inherited from the GPU + # worker. The real DP/TP configuration is used later for model construction, + # registration, and request routing. + offload_config = VllmConfig() + offload_parallel_config = offload_config.parallel_config + offload_parallel_config.data_parallel_size = 1 + offload_parallel_config.data_parallel_size_local = 1 + offload_parallel_config.data_parallel_rank = 0 + offload_parallel_config.data_parallel_rank_local = 0 + offload_parallel_config.data_parallel_index = 0 + offload_parallel_config.tensor_parallel_size = 1 + offload_parallel_config.pipeline_parallel_size = 1 + offload_parallel_config.prefill_context_parallel_size = 1 + offload_parallel_config.decode_context_parallel_size = 1 + offload_parallel_config.world_size = 1 + offload_parallel_config.nnodes = 1 + offload_parallel_config.node_rank = 0 + with set_current_vllm_config(offload_config): + ensure_model_parallel_initialized( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + backend="gloo", + ) + logger.info( + "Distributed environment initialized (backend=gloo, rank=0, world_size=1)." + ) + + +class PleOffloadWorker: + """Manage process creation, READY handshake, and the child entry point.""" + + READY_STR = "READY" + + @staticmethod + def make_process( + vllm_config: VllmConfig, + num_workers: int, + ipc_addr: str, + ) -> PleOffloadWorkerHandle: + """Start one CPU offload process in a fresh interpreter.""" + launch_deadline = ( + time.monotonic() + float(envs.VLLM_PLE_OFFLOAD_READY_TIMEOUT) + ) + config_payload = pickle.dumps( + vllm_config, protocol=pickle.HIGHEST_PROTOCOL + ) + child_env = _offload_child_environment() + + # READY and parent-death supervision must never compete for reads. + parent_control, child_control = socket.socketpair() + parent_death, child_death = socket.socketpair() + ready_reader = Connection(parent_control.detach()) + death_writer = Connection(parent_death.detach()) + try: + proc = subprocess.Popen( + [ + sys.executable, + "-m", + "vllm_ple_offload_bootstrap", + str(child_control.fileno()), + str(child_death.fileno()), + ], + env=child_env, + close_fds=True, + pass_fds=(child_control.fileno(), child_death.fileno()), + ) + except BaseException: + ready_reader.close() + death_writer.close() + raise + finally: + child_control.close() + child_death.close() + + handle = PleOffloadWorkerHandle( + proc=proc, + death_writer=death_writer, + ready_pipe_reader=ready_reader, + launch_deadline=launch_deadline, + ipc_addr=ipc_addr, + ) + try: + config_for_launch: bytes | _FaultReconstructionPayload = config_payload + if ( + _FAULT is not None + and _FAULT.cut == "launch_block_reconstruction" + ): + config_for_launch = _FaultReconstructionPayload(config_payload) + launch_state = { + "vllm_config_payload": config_for_launch, + "num_workers": num_workers, + "ipc_addr": ipc_addr, + } + if ( + _FAULT is not None + and _FAULT.cut == "launch_park_before_payload_read" + ): + # Larger than ordinary AF_UNIX socket buffers: an acknowledged + # non-reading child leaves the parent inside framed transmission. + launch_state["_test_padding"] = b"\0" * (8 << 20) + launch_payload = pickle.dumps( + launch_state, + protocol=pickle.HIGHEST_PROTOCOL, + ) + _send_framed_payload( + ready_reader.fileno(), launch_payload, launch_deadline + ) + except BaseException: + handle.close(graceful=False) + raise + + return handle + + @staticmethod + def wait_for_ready(handle: PleOffloadWorkerHandle) -> None: + """Wait until weights and all GPU registrations are ready to serve.""" + reader = handle.ready_pipe_reader + if reader is None: + return + try: + message = _recv_bounded_connection_message( + reader, handle.launch_deadline + ) + if type(message) is not dict: + raise RuntimeError("PLE offload worker sent an invalid response") + status = dict.get(message, "status") + if type(status) is not str: + raise RuntimeError( + "PLE offload worker sent an invalid response status" + ) + if status == PleOffloadWorker.READY_STR: + pass + elif status == "FAILURE": + error_detail = dict.get(message, "error") + if type(error_detail) is not str: + error_detail = "unknown error" + raise RuntimeError( + "PLE offload worker failed during startup: " + f"{error_detail}" + ) + else: + raise RuntimeError("PLE offload worker sent unsupported status") + layer_names = dict.get(message, "layer_names") + if ( + type(layer_names) is not list + or any( + type(name) is not str for name in layer_names + ) + ): + raise RuntimeError( + "PLE offload worker sent invalid READY layer_names" + ) + logger.info( + "Worker ready - %d PleOffloadLayer(s): %s", + len(layer_names), + layer_names, + ) + except EOFError as error: + handle.close(graceful=False) + raise RuntimeError("PLE offload worker exited during startup") from error + except BaseException: + handle.close(graceful=False) + raise + finally: + if handle.ready_pipe_reader is not None: + handle.ready_pipe_reader.close() + handle.ready_pipe_reader = None + + @staticmethod + def proc_main( + vllm_config: VllmConfig, + num_workers: int, + ipc_addr: str, + ready_pipe: Connection, + ) -> None: + """Load PLE weights, accept registrations, and run the request loop.""" + decorate_logs("PleOffloadWorker") + ready_writer = _BoundedStartupResponseWriter(ready_pipe) + shutdown_event = threading.Event() + + def handle_signal(_signum: int, _frame: object) -> None: + shutdown_event.set() + + signal.signal(signal.SIGTERM, handle_signal) + signal.signal(signal.SIGINT, handle_signal) + + zmq_context: zmq.Context | None = None + pull_socket: zmq.Socket | None = None + owns_ipc_socket = False + try: + # The flag lets PleOffloadLayer subclasses execute their complete + # constructors instead of becoming empty GPU-worker placeholders. + mark_as_offload_worker() + + # The bootstrap established child-only visibility before importing + # torch or vLLM. Retain the Python-level guards for known CUDA + # initialization paths and fail if initialization already occurred. + _prepare_cuda_free_process() + home_device = os.environ.get("VLLM_PLE_OFFLOAD_HOME_DEVICE", "").strip() + if home_device: + logger.info( + "PLE offload bootstrap visibility is CUDA_VISIBLE_DEVICES=%s " + "(host-mapped pull: no CUDA context expected)", + os.environ.get("CUDA_VISIBLE_DEVICES", ""), + ) + + # Initialize Gloo before installing the real VllmConfig. This keeps + # the CPU process in an isolated rank-zero, world-size-one group. + # The offload worker is a PP=1 world; a VLLM_PP_LAYER_PARTITION set + # for the GPU workers (e.g. "25,23") would make make_layers raise + # "len(partitions) does not match pp_size=1" while the meta model + # is built here. Drop it in this process only. + os.environ.pop("VLLM_PP_LAYER_PARTITION", None) + _init_offload_distributed() + + # Model components read the active VllmConfig while the meta model + # is constructed, so keep the context around runner initialization. + with set_current_vllm_config(vllm_config): + runner = PleOffloadRunner(vllm_config) + + zmq_context = zmq.Context() + pull_socket = zmq_context.socket(zmq.PULL) + pull_socket.bind(ipc_addr) + owns_ipc_socket = True + logger.info( + "Bound IPC address %s; waiting for %d GPU worker registration(s).", + ipc_addr, + num_workers, + ) + + # READY means that the process can immediately serve requests. Wait + # for every DP/TP worker to register before notifying the parent. + runner.accept_registrations(pull_socket, num_workers) + # Checked: the host-mapped pull design must not create a context. + logger.info( + "PLE offload host-mapped pull: CUDA initialized in this process = %s", + torch.cuda.is_initialized(), + ) + ready_writer.send( + { + "status": PleOffloadWorker.READY_STR, + "layer_names": sorted(runner.layer_names), + } + ) + ready_writer.close() + ready_writer = None # type: ignore[assignment] + + runner.busy_loop(pull_socket, shutdown_event) + except Exception as error: + logger.exception("Unexpected failure in PLE offload worker.") + if ready_writer is not None: + with contextlib.suppress(Exception): + ready_writer.send( + {"status": "FAILURE", "error": repr(error)[:4096]} + ) + raise + finally: + if pull_socket is not None: + pull_socket.close(linger=0) + if zmq_context is not None: + zmq_context.term() + if owns_ipc_socket: + _unlink_ipc_socket(ipc_addr) + if ready_writer is not None: + ready_writer.close() + + +class PleOffloadRunner: + """Own all discovered PLE tables and serve every local DP rank.""" + + def __init__(self, vllm_config: VllmConfig) -> None: + self.vllm_config = vllm_config + self._clamp_input_ids = ( + getattr(vllm_config, "speculative_config", None) is not None + ) + # name -> PleOffloadLayer (CPU) + self._layers: dict[str, PleOffloadLayer] = {} + # dp_rank -> layer_name -> one destination per TP rank + self._worker_targets: dict[int, dict[str, list[PleOffloadOutputTarget]]] = {} + # Each (dp_rank, layer_name) pair owns a plain CPU scratch buffer that + # forward_impl writes; the rows are then copied into every TP worker's + # shared host buffer. Nothing in this process is pinned or on a GPU. + self._scratch_bufs: dict[int, dict[str, torch.Tensor]] = {} + self._shutdown_event: threading.Event | None = None + self._step = 0 # number of handled request batches (diagnostics) + self._request_index = 0 + # Shared-memory inputs are registered once per DP rank by TP rank zero. + self._input_bufs: dict[int, PleOffloadInputBuffers] = {} + self._load_weights() + + @property + def layer_names(self) -> list[str]: + """Return PleOffloadLayer names in model traversal order.""" + return list(self._layers) + + def _load_weights(self) -> None: + """Load only :class:`PleOffloadLayer` subtrees into CPU memory. + + Strategy: + 1. Build the entire model on ``meta`` so non-offloaded parameters use no + physical memory. PleOffloadLayer constructors explicitly target CPU. + 2. Discover all PleOffloadLayer modules from the complete model. + 3. Stream the checkpoint through a prefix filter so only matching PLE + tensors are materialized and passed to ``model.load_weights``. + 4. Run post-load processing only on the CPU-owned PLE subtrees. + """ + model_config = self.vllm_config.model_config + load_config = self.vllm_config.load_config + + # Step 1: build complete structure, while only PLE subtrees allocate CPU + # memory. All transformer, MoE, and vision parameters remain on meta. + logger.info("Initializing model structure for PLE weight discovery ...") + model_dtype = cast(torch.dtype, model_config.dtype) + with ( + set_default_torch_dtype(model_dtype), + torch.device("meta"), + _platform_device_is_meta(), + ): + model = initialize_model( + vllm_config=self.vllm_config, + model_config=model_config, + ) + + # Step 2: preserve named_modules DFS order so CPU execution follows the + # same layer order as the GPU model forward. + offload_layers = { + name: module + for name, module in model.named_modules() + if isinstance(module, PleOffloadLayer) + } + if not offload_layers: + raise RuntimeError( + "VLLM_PLE_CPU_OFFLOAD is enabled, but no PleOffloadLayer " + "was found in the initialized model" + ) + logger.info( + "Found %d PleOffloadLayer(s): %s", + len(offload_layers), + sorted(offload_layers), + ) + offload_prefixes = tuple(f"{name}." for name in offload_layers) + + # Step 3: filter checkpoint tensors before model.load_weights(). The + # conditional-generation checkpoint uses HF names such as + # ``model.language_model.*`` while named_modules exposes mapped vLLM + # names such as ``language_model.model.*``. Apply the model mapper only + # for matching, then yield the original pair so load_weights performs + # its normal single mapping pass. + mapper = getattr(model, "hf_to_vllm_mapper", None) + matched_checkpoint_tensors = 0 + + def offload_only_iter( + weights: Iterable[tuple[str, torch.Tensor]], + ) -> Iterable[tuple[str, torch.Tensor]]: + nonlocal matched_checkpoint_tensors + for weight_name, tensor in weights: + mapped_name: str | None = weight_name + if mapper is not None: + mapped_names = mapper.apply_list([weight_name]) + mapped_name = mapped_names[0] if mapped_names else None + if mapped_name is not None and mapped_name.startswith(offload_prefixes): + matched_checkpoint_tensors += 1 + yield weight_name, tensor + + loader = get_model_loader(load_config) + if isinstance(loader, DummyModelLoader): + logger.info( + "Initializing dummy weights for %d PleOffloadLayer(s) ...", + len(offload_layers), + ) + for layer in offload_layers.values(): + initialize_dummy_weights(layer, model_config) + elif isinstance(loader, DefaultModelLoader): + all_weights = loader.get_all_weights(model_config, model) + loaded_params = model.load_weights(offload_only_iter(all_weights)) + if matched_checkpoint_tensors == 0: + raise RuntimeError( + "PLE offload checkpoint filter matched no weights for " + f"layers: {sorted(offload_layers)}" + ) + + expected_offload_params = { + f"{layer_name}.{param_name}" + for layer_name, layer in offload_layers.items() + for param_name, _ in layer.named_parameters() + } + loaded_offload_entries = { + name for name in loaded_params if name.startswith(offload_prefixes) + } + loaded_expected_params = expected_offload_params.intersection(loaded_params) + missing_offload_params = sorted( + expected_offload_params.difference(loaded_expected_params) + ) + if missing_offload_params: + raise RuntimeError( + "PLE offload checkpoint did not load all materialized " + f"parameters: {missing_offload_params}" + ) + logger.info( + "PLE offload matched %d checkpoint tensor(s), loaded %d " + "offload entries, and verified %d/%d materialized " + "parameter(s) for layers: %s", + matched_checkpoint_tensors, + len(loaded_offload_entries), + len(loaded_expected_params), + len(expected_offload_params), + sorted(offload_layers), + ) + else: + raise NotImplementedError( + "PLE offload requires the default or dummy model loader, got " + f"{type(loader).__name__}" + ) + + # Step 4: post-load processing is restricted to CPU-owned PLE modules; + # the remainder of the model is still on meta and must not be visited. + for layer in offload_layers.values(): + process_weights_after_loading(layer, model_config, torch.device("cpu")) + + self._layers.update(offload_layers) + del model + logger.info("PLE weight loading complete.") + + def accept_registrations( + self, + pull_socket: zmq.Socket, + num_workers: int, + ) -> None: + """Receive every local DP/TP worker's IPC and shared-memory buffers.""" + logger.info("Waiting for %d GPU worker registration(s) ...", num_workers) + registrations: list[PleOffloadRegistration] = [] + for index in range(num_workers): + item = pickle.loads(pull_socket.recv()) + if not isinstance(item, PleOffloadRegistration): + raise RuntimeError( + "Expected PleOffloadRegistration during setup, got " + f"{type(item).__name__} ({index + 1}/{num_workers})" + ) + registrations.append(item) + logger.info( + "GPU worker %d registered (dp_rank=%d, tp_rank=%d, layers=%s).", + item.worker_id, + item.dp_rank, + item.tp_rank, + sorted(item.host_output_buffers), + ) + + dp_size = self.vllm_config.parallel_config.data_parallel_size + tp_size = self.vllm_config.parallel_config.tensor_parallel_size + if num_workers != dp_size * tp_size: + raise RuntimeError( + f"Expected {dp_size * tp_size} registrations for DP={dp_size}, " + f"TP={tp_size}, got {num_workers}" + ) + + registrations_by_dp: dict[int, list[PleOffloadRegistration]] = {} + for registration in registrations: + registrations_by_dp.setdefault(registration.dp_rank, []).append( + registration + ) + if set(registrations_by_dp) != set(range(dp_size)): + raise RuntimeError( + f"Expected DP ranks {set(range(dp_size))}, " + f"got {set(registrations_by_dp)}" + ) + for dp_rank, dp_registrations in registrations_by_dp.items(): + tp_ranks = {registration.tp_rank for registration in dp_registrations} + if tp_ranks != set(range(tp_size)): + raise RuntimeError( + f"DP rank {dp_rank} expected TP ranks {set(range(tp_size))}, " + f"got {tp_ranks}" + ) + + for registration in registrations: + if set(registration.host_output_buffers) != set(self.layer_names) or set( + registration.host_flag_tensors + ) != set(self.layer_names): + raise RuntimeError( + "Registered PLE layers do not match CPU layers: " + f"registered={sorted(registration.host_output_buffers)} / " + f"{sorted(registration.host_flag_tensors)}, " + f"cpu={sorted(self.layer_names)}" + ) + targets_for_dp = self._worker_targets.setdefault(registration.dp_rank, {}) + for layer_name, host_buffer in registration.host_output_buffers.items(): + host_flag = registration.host_flag_tensors[layer_name] + for tensor in (host_buffer, host_flag): + if tensor.device.type != "cpu" or not tensor.is_shared(): + raise RuntimeError( + f"PLE host buffer for {layer_name} (dp=" + f"{registration.dp_rank}, tp={registration.tp_rank}) " + "is not a shared CPU tensor" + ) + if host_flag.dtype != torch.int32 or host_flag.numel() < 1: + raise RuntimeError(f"PLE host flag for {layer_name} must be int32") + target = PleOffloadOutputTarget( + tp_rank=registration.tp_rank, + host_output_buffer=host_buffer, + host_flag=host_flag, + ) + targets_for_dp.setdefault(layer_name, []).append(target) + # All TP ranks in one DP group receive the same input, so buffers + # registered by TP rank zero are sufficient for that DP rank. + if registration.tp_rank == 0: + self._input_bufs[registration.dp_rank] = PleOffloadInputBuffers( + input_ids_buf=registration.input_ids_buf, + query_start_loc_buf=registration.query_start_loc_buf, + ngram_context_buf=registration.ngram_context_buf, + ) + + if set(self._input_bufs) != set(range(dp_size)): + raise RuntimeError( + "TP rank zero did not register PLE input buffers for every DP " + f"rank: expected={set(range(dp_size))}, got={set(self._input_bufs)}" + ) + + config = self.vllm_config.model_config.hf_text_config + max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens + for dp_rank, layer_targets in self._worker_targets.items(): + self._scratch_bufs[dp_rank] = {} + for layer_name, targets in layer_targets.items(): + if len(targets) != tp_size: + raise RuntimeError( + f"PLE layer {layer_name} for DP rank {dp_rank} received " + f"{len(targets)} targets, expected {tp_size}" + ) + targets.sort(key=lambda target: target.tp_rank) + out_dim = self._layers[layer_name].get_offload_output_dim( + int(config.ple_embed_dim) + ) + out_dtype = self._layers[layer_name].get_offload_output_dtype( + self.vllm_config.model_config.dtype + ) + for target in targets: + shape = tuple(target.host_output_buffer.shape) + if shape != (max_tokens, out_dim) or ( + target.host_output_buffer.dtype != out_dtype + ): + raise RuntimeError( + f"PLE host buffer for {layer_name} (dp={dp_rank}, " + f"tp={target.tp_rank}) is {shape} " + f"{target.host_output_buffer.dtype}, expected " + f"{(max_tokens, out_dim)} {out_dtype}" + ) + self._scratch_bufs[dp_rank][layer_name] = torch.empty( + max_tokens, out_dim, dtype=out_dtype, device="cpu" + ) + logger.info( + "Registrations complete (dp_size=%d, tp_size=%d, layers=%s).", + dp_size, + tp_size, + sorted(self.layer_names), + ) + + @torch.inference_mode() + def busy_loop( + self, + pull_socket: zmq.Socket, + shutdown_event: threading.Event, + ) -> None: + """Decode and batch available requests by DP rank until shutdown.""" + logger.info("Busy-loop started.") + self._shutdown_event = shutdown_event + poller = zmq.Poller() + poller.register(pull_socket, zmq.POLLIN) + while not shutdown_event.is_set(): + if pull_socket not in dict(poller.poll(timeout=100)): + continue + + requests = [] + try: + requests.append(_receive_offload_request(pull_socket)) + while True: + requests.append( + _receive_offload_request(pull_socket, zmq.NOBLOCK) + ) + except zmq.Again: + pass + except (msgspec.DecodeError, RuntimeError) as error: + raise RuntimeError("Unexpected PLE offload request") from error + + self._handle_requests(requests) + + def _handle_requests( + self, requests: list[_ReceivedPleOffloadRequest] + ) -> None: + """Run requests layer-first so each DP rank can resume promptly.""" + with record_function_or_nullcontext("ple_offload: handle_requests"): + self._handle_requests_impl(requests) + + def _validate_requests_before_publication( + self, requests: list[PleOffloadRequest] + ) -> None: + """Validate a drained request batch before computation or publication.""" + for request in requests: + dp_rank = request.dp_rank + if dp_rank not in self._worker_targets or dp_rank not in self._input_bufs: + raise RuntimeError( + f"No PLE buffers or output targets for dp_rank={dp_rank}; " + "refusing request" + ) + + input_bufs = self._input_bufs[dp_rank] + if not 0 < request.num_tokens <= input_bufs.input_ids_buf.numel(): + raise RuntimeError( + f"PLE num_tokens={request.num_tokens} is outside capacity " + f"{input_bufs.input_ids_buf.numel()} for dp_rank={dp_rank}" + ) + max_reqs = input_bufs.query_start_loc_buf.numel() - 1 + if not 0 < request.num_reqs <= max_reqs: + raise RuntimeError( + f"PLE num_reqs={request.num_reqs} is outside capacity " + f"{max_reqs} for dp_rank={dp_rank}" + ) + if ( + input_bufs.query_start_loc_buf.dtype != torch.int32 + or input_bufs.query_start_loc_buf.dim() != 1 + ): + raise RuntimeError( + f"PLE query_start_loc for dp_rank={dp_rank} must be 1-D int32" + ) + if ( + input_bufs.ngram_context_buf is not None + and request.num_reqs > input_bufs.ngram_context_buf.shape[0] + ): + raise RuntimeError( + f"PLE num_reqs={request.num_reqs} exceeds ngram context " + f"capacity for dp_rank={dp_rank}" + ) + + query_start_loc = input_bufs.query_start_loc_buf[ + : request.num_reqs + 1 + ].tolist() + if ( + query_start_loc[0] != 0 + or query_start_loc[-1] > request.num_tokens + or any(a > b for a, b in zip(query_start_loc, query_start_loc[1:])) + ): + raise RuntimeError( + f"Invalid PLE query_start_loc for dp_rank={dp_rank}: " + f"{query_start_loc}, num_tokens={request.num_tokens}" + ) + + for layer_name, targets in self._worker_targets[dp_rank].items(): + scratch = self._scratch_bufs[dp_rank][layer_name] + if request.num_tokens > scratch.shape[0] or any( + request.num_tokens > target.host_output_buffer.shape[0] + for target in targets + ): + raise RuntimeError( + f"PLE num_tokens={request.num_tokens} exceeds output " + f"capacity for {layer_name}, dp_rank={dp_rank}" + ) + + def _handle_requests_impl( + self, + requests: list[_ReceivedPleOffloadRequest | PleOffloadRequest], + ) -> None: + """Process one batch of PLE requests.""" + received = [ + item + if isinstance(item, _ReceivedPleOffloadRequest) + else _ReceivedPleOffloadRequest(item, False, None, None) + for item in requests + ] + plain_requests = [item.request for item in received] + self._validate_requests_before_publication(plain_requests) + received_by_dp: dict[int, _ReceivedPleOffloadRequest] = {} + for item in received: + request = item.request + if request.dp_rank not in self._worker_targets: + raise RuntimeError( + f"No PLE output targets for dp_rank={request.dp_rank}; " + "refusing request" + ) + if request.dp_rank in received_by_dp: + # One input slot and one output slot per DP rank: a second + # request in the same batch would alias them. Fatal, never + # silently dropped. + raise RuntimeError( + f"Two PLE requests for dp_rank={request.dp_rank} in one batch " + f"(step {self._step}); the single-slot protocol is violated" + ) + received_by_dp[request.dp_rank] = item + self._step += 1 + if _FAULT is not None: + for item in received_by_dp.values(): + if not item.serving: + continue + self._request_index += 1 + transaction = f"request:{self._request_index}" + if item.transaction_id is not None: + transaction = f"{transaction}@{item.transaction_id}" + for cut in ("park_before_publish", "park_before_publish_ack"): + if _FAULT is not None and _FAULT.fires( + cut, + role="offload", + unit="offload-request", + eligible=True, + dp_rank=item.request.dp_rank, + transaction_id=transaction, + correlation_id=item.correlation_id, + ): + # No multiprocessing or fault-state lock is held here. + threading.Event().wait() + + requests_by_dp = { + dp_rank: item.request + for dp_rank, item in received_by_dp.items() + } + + # Speculative placeholders are not vocabulary IDs. Normalize each DP + # input once before all PLE layers consume the shared buffer. + if self._clamp_input_ids: + for dp_rank, request in requests_by_dp.items(): + self._input_bufs[dp_rank].input_ids_buf[ + : request.num_tokens + ].clamp_min_(0) + + for layer_name, layer in self._layers.items(): + for dp_rank, request in requests_by_dp.items(): + targets = self._worker_targets[dp_rank][layer_name] + + input_bufs = self._input_bufs[dp_rank] + ngram_context = ( + input_bufs.ngram_context_buf[: request.num_reqs] + if input_bufs.ngram_context_buf is not None + else None + ) + result = layer.forward_impl( + input_bufs.input_ids_buf[: request.num_tokens], + input_bufs.input_ids_buf[: request.num_tokens], + input_bufs.query_start_loc_buf[: request.num_reqs + 1], + ngram_context, + output_buffer=self._scratch_bufs[dp_rank][layer_name], + ) + if _rowdump.enabled(): + _seq = getattr(self, "_rowdump_seq", 0) + self._rowdump_seq = _seq + 1 + _rowdump.write(f"offload-dp{dp_rank}", 1, dp_rank, _seq, + input_bufs.input_ids_buf[: request.num_tokens], result) + + # The result is identical on every TP rank in this DP group. + # A worker's host buffer may be overwritten only after its GPU + # released the previous result (``release_outputs`` stores + # flag=0 on the model stream after the forward, which is also + # after the in-graph DMA read it). Then write rows, then flag=1. + scratch = self._scratch_bufs[dp_rank][layer_name] + if result.dim() != 2 or result.shape[0] != request.num_tokens or ( + result.shape[1] != scratch.shape[1] + ) or result.dtype != scratch.dtype or result.device.type != "cpu": + raise RuntimeError( + f"PLE result {tuple(result.shape)} {result.dtype} " + f"{result.device} for {layer_name} does not " + f"match the request ({request.num_tokens} rows x " + f"{scratch.shape[1]} {scratch.dtype} CPU); refusing to " + f"publish (step {self._step})" + ) + slices = tuple(slice(0, size) for size in result.shape) + for target in targets: + self._wait_flag_released(target, layer_name, dp_rank) + target.host_output_buffer[slices].copy_(result[slices]) + _publish_fence() + target.host_flag[0] = 1 + + def _wait_flag_released( + self, target: PleOffloadOutputTarget, layer_name: str, dp_rank: int + ) -> None: + """Spin until the GPU reset this target's flag (previous result used).""" + flag = target.host_flag + if int(flag[0]) == 0: + return + spins = 0 + warned = False + start = time.monotonic() + while int(flag[0]) != 0: + spins += 1 + if self._shutdown_event is not None and self._shutdown_event.is_set(): + raise RuntimeError("PLE offload shutting down while waiting for GPU") + if spins > 200: + time.sleep(0.0001) + waited = time.monotonic() - start + if not warned and waited > 5.0: + warned = True + logger.warning( + "PLE offload waited >5 s for the GPU to release %s " + "(dp=%d, tp=%d, step %d); flag=%d", + layer_name, + dp_rank, + target.tp_rank, + self._step, + int(flag[0]), + ) + if waited > _PLE_WAIT_TIMEOUT_S: + raise RuntimeError( + f"PLE offload: GPU did not release {layer_name} (dp={dp_rank}, " + f"tp={target.tp_rank}, step {self._step}) within " + f"{_PLE_WAIT_TIMEOUT_S:.0f} s; refusing to overwrite" + ) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 9657c28..0e132b6 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib import gc import itertools from collections import defaultdict -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Iterator from contextlib import ExitStack from dataclasses import dataclass, replace from itertools import groupby, product @@ -30,6 +31,9 @@ from vllm.distributed.parallel_state import ( ) from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger +from vllm.model_executor.layers.ple_offload_layer import ( + observe_ple_copy_geometry, +) from vllm.model_executor.offloader.base import get_offloader from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors @@ -170,6 +174,14 @@ class CudaGraphManager: self._lora_dispatch_map, self._max_lora_case = self._build_lora_dispatch_map() self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {} + # FULL graphs bake every PLE wait/copy into their native graph. Metadata + # comes from the actual custom-op implementation during native capture. + self._ple_full_copy_geometry: dict[ + BatchExecutionDescriptor, tuple[tuple[str, int, int], ...] + ] = {} + self._expected_ple_copy_layers: tuple[str, ...] = () + # Only the target-model manager is allowed to capture PLE operations. + self._allow_ple_full_graph = False self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None self._graphs_captured = False @@ -469,15 +481,16 @@ class CudaGraphManager: if self._capture_mem_samples is not None: torch.accelerator.synchronize() free_before = torch.accelerator.get_memory_info()[0] - with torch.cuda.graph( - graph, self.pool, stream=self._capture_stream(desc) - ): - forward_fn(CUDAGraphMode.NONE) - # Join offloader's copy stream after forward to avoid - # unjoined stream error. The last layer's start_prefetch - # forks copy_stream, but wait_prefetch only happens in - # the next forward pass. - get_offloader().join_after_forward() + with self.capture_ple_copy_geometry(desc): + with torch.cuda.graph( + graph, self.pool, stream=self._capture_stream(desc) + ): + forward_fn(CUDAGraphMode.NONE) + # Join offloader's copy stream after forward to avoid + # unjoined stream error. The last layer's start_prefetch + # forks copy_stream, but wait_prefetch only happens in + # the next forward pass. + get_offloader().join_after_forward() if self._capture_mem_samples is not None: torch.accelerator.synchronize() free_after = torch.accelerator.get_memory_info()[0] @@ -493,6 +506,78 @@ class CudaGraphManager: {desc.num_tokens for desc in self.graphs if desc.num_active_loras == 0} ) + @contextlib.contextmanager + def capture_ple_copy_geometry( + self, desc: BatchExecutionDescriptor + ) -> Iterator[None]: + """Production capture callback for actual PLE custom-op geometry.""" + with observe_ple_copy_geometry() as geometry: + yield + + if geometry and not self._allow_ple_full_graph: + raise RuntimeError( + "PLE wait-copy reached an unsupported FULL graph path; " + "speculator/draft FULL graphs may not contain PLE operations" + ) + if geometry: + self.record_ple_copy_geometry(desc, geometry) + elif self._expected_ple_copy_layers: + raise RuntimeError( + f"PLE copy geometry was not observed while capturing FULL graph {desc}" + ) + + def record_ple_copy_geometry( + self, + desc: BatchExecutionDescriptor, + geometry: Iterable[tuple[str, int, int]], + ) -> None: + if desc.cg_mode != CUDAGraphMode.FULL: + raise ValueError("PLE capture geometry is recorded only for FULL graphs") + observed = tuple( + (str(layer), int(rows), int(nbytes)) + for layer, rows, nbytes in geometry + ) + if not observed: + raise RuntimeError(f"PLE FULL graph {desc} observed no wait-copy operations") + observed_layers = tuple(layer for layer, _, _ in observed) + expected_layers = self._expected_ple_copy_layers + if expected_layers and sorted(observed_layers) != sorted(expected_layers): + raise RuntimeError( + f"PLE FULL graph {desc} observed layers {observed_layers}, " + f"expected {expected_layers}" + ) + inconsistent = [ + (layer, rows, nbytes) + for layer, rows, nbytes in observed + if rows != desc.num_tokens or nbytes < 0 + ] + if inconsistent: + raise RuntimeError( + f"PLE FULL graph {desc} actual PLE copy geometry " + f"{inconsistent} does not match descriptor rows={desc.num_tokens}" + ) + previous = self._ple_full_copy_geometry.get(desc) + if previous is not None and previous != observed: + raise RuntimeError( + f"PLE FULL graph {desc} changed copy geometry from " + f"{previous} to {observed}" + ) + self._ple_full_copy_geometry[desc] = observed + + def ple_copy_rows(self, desc: BatchExecutionDescriptor) -> int: + try: + geometry = self._ple_full_copy_geometry[desc] + except KeyError as error: + raise RuntimeError( + f"PLE copy geometry was not recorded for FULL graph {desc}" + ) from error + rows = {num_rows for _, num_rows, _ in geometry} + if len(rows) != 1: + raise RuntimeError( + f"PLE FULL graph {desc} has inconsistent per-layer rows: {geometry}" + ) + return rows.pop() + def dispatch( self, num_reqs: int, @@ -532,6 +617,11 @@ class CudaGraphManager: f"Expected FULL mode, got {desc.cg_mode}" ) assert desc in self.graphs, f"No cudagraph for {desc}" + if ( + desc in self._ple_full_copy_geometry + and not self._allow_ple_full_graph + ): + raise RuntimeError("refusing unsupported PLE speculator FULL replay") # Sync offloader before replay - needed when transitioning from # eager/piecewise to full cudagraph (e.g., prefill → decode). # The previous eager iteration's start_prefetch may have queued @@ -577,11 +667,32 @@ class ModelCudaGraphManager(CudaGraphManager): varlen_decode=varlen_decode, ubatch_runner=ubatch_runner, ) + self._allow_ple_full_graph = True self.hidden_states: torch.Tensor | None = None self.aux_hidden_states: list[torch.Tensor] = [] self.use_aux_hidden_state_outputs = False self.intermediate_tensors: IntermediateTensors | None = None + def _prepare_ple_capture( + self, model: nn.Module, *, initialize_breakable: bool + ) -> None: + expected_ple_copy_layers = tuple( + str(getattr(module, "_ple_offload_name", name)) + for name, module in model.named_modules() + if getattr(module, "_is_cpu_offloaded", False) + and getattr(module, "_host_output_ptr", None) is not None + ) + if ( + "" in expected_ple_copy_layers + or len(set(expected_ple_copy_layers)) != len(expected_ple_copy_layers) + ): + raise RuntimeError( + "PLE FULL capture requires distinct, bound connector layer identities" + ) + self._expected_ple_copy_layers = expected_ple_copy_layers + if initialize_breakable: + self.init_breakable_cg_runner(model) + def capture( self, model: nn.Module, @@ -599,8 +710,9 @@ class ModelCudaGraphManager(CudaGraphManager): ) -> None: """Capture CUDA graphs for model forward pass.""" self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs - if self.use_breakable_cg: - self.init_breakable_cg_runner(model) + self._prepare_ple_capture( + model, initialize_breakable=self.use_breakable_cg + ) if self.cudagraph_mode.has_piecewise_cudagraphs() and not ( self.use_breakable_cg or has_compiled_submodule(model) @@ -901,7 +1013,9 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int: if manager.use_breakable_cg: # Create the breakable runner before the wrapper pool swap so # its pool is covered as well. - manager.init_breakable_cg_runner(runner.model) + manager._prepare_ple_capture( + runner.model, initialize_breakable=True + ) all_wrappers = list(CUDAGraphWrapper._all_instances) + list( BreakableCUDAGraphWrapper._all_instances ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 744fe5b..e2ddc52 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -77,6 +77,9 @@ from vllm.v1.outputs import ( ModelRunnerOutput, RoutedExpertsTensors, ) +from vllm.v1.ple_offload.connector import PleOffloadConnector +from vllm.v1.ple_offload.fatal import fail_ple_execute_model +from vllm.v1.ple_offload.test_fault import _FAULT from vllm.v1.watermarking import create_watermarker from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler from vllm.v1.watermarking.spec_decode import ( @@ -1689,6 +1692,45 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.ec_connector.no_forward(scheduler_output).ec_connector_output, ) + def _prepare_ple_forward( + self, + input_batch: InputBatch, + batch_desc: BatchExecutionDescriptor, + dummy_run: bool, + *, + serving: bool = False, + transaction_id: str | None = None, + correlation_id: str | None = None, + ) -> None: + connector = self._ple_offload_connector + assert connector is not None + request_num_tokens = int(input_batch.num_tokens_after_padding) + if batch_desc.cg_mode == CUDAGraphMode.FULL: + manager = self.cudagraph_manager + if manager is None: + raise RuntimeError("PLE FULL replay has no CUDA graph manager") + copy_num_tokens = manager.ple_copy_rows(batch_desc) + else: + # PIECEWISE/eager execute the Python op with this actual padded view. + copy_num_tokens = int(input_batch.input_ids.shape[0]) + + if request_num_tokens != copy_num_tokens: + raise RuntimeError( + "PLE request/selected-forward copy-count mismatch: " + f"request={request_num_tokens}, copy={copy_num_tokens}, " + f"mode={batch_desc.cg_mode}" + ) + connector.prepare_forward( + input_batch.num_reqs, + request_num_tokens, + dummy_run, + consumer_num_tokens=copy_num_tokens, + serving=serving, + transaction_id=transaction_id, + correlation_id=correlation_id, + ) + + @fail_ple_execute_model @torch.inference_mode() def execute_model( self, @@ -1956,6 +1998,54 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) self.step_timing.forward_start() + # prepare_inputs has finalized MRV2's GPU buffers. Stage D2H on this + # stream before the PLE placeholder can wait for CPU output. + connector = self._ple_offload_connector + ple_serving = False + ple_sequence = None + ple_transaction = None + ple_correlation = None + if connector is not None: + ple_serving = not dummy_run and not is_profile + ple_sequence = connector.begin_forward( + dummy_run=dummy_run, + num_tokens=int(input_batch.num_tokens_after_padding), + ) + if _FAULT is not None and ple_serving: + ple_correlation = f"sequence:{ple_sequence}" + ple_transaction = connector.fault_transaction(ple_sequence) + if _FAULT is not None and _FAULT.fires( + "exc_before_prepare", + role="gpu", + unit="runner-forward", + eligible=ple_serving, + tp_rank=connector.tp_rank, + pp_rank=connector.pp_rank, + dp_rank=connector.dp_rank, + transaction_id=ple_transaction, + correlation_id=ple_correlation, + ): + raise RuntimeError("PLE test fault: exc_before_prepare") + self._prepare_ple_forward( + input_batch, + batch_desc, + dummy_run, + serving=ple_serving, + transaction_id=ple_transaction, + correlation_id=ple_correlation, + ) + if _FAULT is not None and _FAULT.fires( + "exc_after_prepare", + role="gpu", + unit="runner-forward", + eligible=ple_serving, + tp_rank=connector.tp_rank, + pp_rank=connector.pp_rank, + dp_rank=connector.dp_rank, + transaction_id=ple_transaction, + correlation_id=ple_correlation, + ): + raise RuntimeError("PLE test fault: exc_after_prepare") connector_kwargs = dict( scheduler_output=scheduler_output, request_state_indices=input_batch.idx_mapping, @@ -2013,6 +2103,26 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Eager (NONE): call the raw model directly. model_output = self.model(**model_inputs) + # The model has consumed every PLE output, so the CPU worker may reuse + # the registered buffers for the next request. + if connector is not None: + connector.release_outputs( + serving=ple_serving, + transaction_id=ple_transaction, + correlation_id=ple_correlation, + ) + if _FAULT is not None and _FAULT.fires( + "exc_after_release", + role="gpu", + unit="runner-forward", + eligible=ple_serving, + tp_rank=connector.tp_rank, + pp_rank=connector.pp_rank, + dp_rank=connector.dp_rank, + transaction_id=ple_transaction, + correlation_id=ple_correlation, + ): + raise RuntimeError("PLE test fault: exc_after_release") self.kv_connector.finish_forward() if self.is_last_pp_rank: @@ -2047,6 +2157,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): ec_connector_output=ec_connector_output, routed_experts=routed_experts, cudagraph_stats=cudagraph_stats, + ple_serving=ple_serving, + ple_sequence=ple_sequence, + ple_transaction=ple_transaction, ) if not self.is_last_pp_rank: @@ -2058,6 +2171,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) return None + @fail_ple_execute_model @torch.inference_mode() @step_eplb_after() def sample_tokens( @@ -2077,6 +2191,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): ec_connector_output = self.execute_model_state.ec_connector_output routed_experts = self.execute_model_state.routed_experts cudagraph_stats = self.execute_model_state.cudagraph_stats + ple_serving = self.execute_model_state.ple_serving + ple_sequence = self.execute_model_state.ple_sequence + ple_transaction = self.execute_model_state.ple_transaction self.execute_model_state = None if not self.is_last_pp_rank: @@ -2084,6 +2201,26 @@ class GPUModelRunner(LoRAModelRunnerMixin): # IntermediateTensors instead of final hidden states. Receive the # sampled tokens broadcast from the last rank and update local state. assert self.pp_handler is not None + connector = self._ple_offload_connector + if ( + connector is not None + and _FAULT is not None + and ple_serving + and ple_sequence is not None + and connector.completion_is_pending(ple_sequence) + and _FAULT.fires( + "exc_async_output", + role="gpu", + unit="runner-output-relay", + eligible=True, + tp_rank=connector.tp_rank, + pp_rank=connector.pp_rank, + dp_rank=connector.dp_rank, + transaction_id=ple_transaction, + correlation_id=f"sequence:{ple_sequence}", + ) + ): + raise RuntimeError("PLE test fault: exc_async_output") all_decode_next = self.pp_handler.receive(input_batch) # Optimistically update num_computed_tokens for entire batch here. # Will be adjusted for rejections if necessary in update_requests. @@ -2379,6 +2516,9 @@ class ExecuteModelState(NamedTuple): ec_connector_output: ECConnectorOutput | None routed_experts: RoutedExpertsTensors | None cudagraph_stats: CUDAGraphStat | None + ple_serving: bool + ple_sequence: int | None + ple_transaction: str | None class BatchReqState(NamedTuple): diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 09f91f9..43bc705 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -87,6 +87,7 @@ from vllm.v1.outputs import ( DraftTokenIds, ModelRunnerOutput, ) +from vllm.v1.ple_offload.fatal import fail_ple_execute_model from vllm.v1.utils import compute_iteration_details, report_usage_stats from vllm.v1.worker import guard_telemetry from vllm.v1.worker.sentinel.gpu_worker_sentinel import WorkerSentinel @@ -231,6 +232,148 @@ class Worker(WorkerBase): # Resolved lazily on first sleep/wake; persists worker-process state. self._sleep_mode_backend: SleepModeBackend | None = None + def _has_ple_layers(self) -> bool: + """Return whether this model configuration constructs PLE layers.""" + if not envs.VLLM_PLE_CPU_OFFLOAD: + return False + text_config = self.model_config.hf_text_config + ple_layer_ids = getattr(text_config, "ple_layer_ids", None) + if not ple_layer_ids: + return False + # PP (vllm#54709, carried from the earlier overlay): only the pipeline rank + # that owns a PLE layer takes part in the offload (rank 0 with the + # default partition). Other ranks build PPMissingLayer there and must + # skip the connector entirely. + pp_size = self.parallel_config.pipeline_parallel_size + if pp_size > 1: + from vllm.distributed.utils import get_pp_indices + + tp_size = self.parallel_config.tensor_parallel_size + pp_rank = (self.rank // tp_size) % pp_size + start, end = get_pp_indices( + int(text_config.num_hidden_layers), pp_rank, pp_size + ) + return any(start <= int(i) - 1 < end for i in ple_layer_ids) + return True + + def _validate_ple_offload_config(self) -> None: + """Reject unsupported PLE offload execution modes.""" + parallel_config = self.parallel_config + unsupported = [] + if not self.use_v2_model_runner: + unsupported.append("model runner V1") + if not current_platform.is_cuda(): + unsupported.append(f"device={current_platform.device_type}") + if parallel_config.nnodes != 1: + unsupported.append(f"nnodes={parallel_config.nnodes}") + if parallel_config.data_parallel_backend != "mp": + unsupported.append(f"DP backend={parallel_config.data_parallel_backend}") + if ( + parallel_config.data_parallel_size_local + != parallel_config.data_parallel_size + ): + unsupported.append( + "non-local DP " + f"({parallel_config.data_parallel_size_local}/" + f"{parallel_config.data_parallel_size} local ranks)" + ) + # PP allowed (vllm#54709): PLE layers must live on pipeline rank 0 + # (enforced placement-based in models/config.py); num_workers stays + # dp_size * tp_size = the ranks of the owning stage. + if parallel_config.prefill_context_parallel_size != 1: + unsupported.append(f"PCP={parallel_config.prefill_context_parallel_size}") + if parallel_config.decode_context_parallel_size != 1: + unsupported.append(f"DCP={parallel_config.decode_context_parallel_size}") + if parallel_config.use_ubatching: + unsupported.append("ubatching/DBO") + if self.model_config.architecture not in { + "Qwen4ExpForCausalLM", + "Qwen4ExpForConditionalGeneration", + }: + unsupported.append("architecture") + if self.vllm_config.weight_transfer_config is not None: + unsupported.append("weight transfer") + + if unsupported: + raise ValueError( + "VLLM_PLE_CPU_OFFLOAD does not support the requested " + "configuration. Unsupported settings: " + f"{', '.join(unsupported)}" + ) + + def spawn_ple_offload(self) -> None: + """Spawn one node-local PLE CPU worker from DP0/TP0.""" + if ( + not self._ple_offload_enabled + or self.rank != 0 + or self.parallel_config.data_parallel_rank != 0 + ): + return + + from vllm.v1.ple_offload.worker import PleOffloadWorker + + ipc_addr = self.parallel_config._ple_offload_ipc_path + if not ipc_addr: + raise RuntimeError("PLE offload IPC address was not initialized") + dp_size = self.parallel_config.data_parallel_size + tp_size = self.parallel_config.tensor_parallel_size + num_workers = dp_size * tp_size + logger.info( + "PleOffload: spawning worker " + "(rank=%d, local_rank=%d, dp_size=%d, tp_size=%d, " + "num_workers=%d, ipc_addr=%s).", + self.rank, + self.local_rank, + dp_size, + tp_size, + num_workers, + ipc_addr, + ) + self._ple_offload_worker_handle = PleOffloadWorker.make_process( + self.vllm_config, + num_workers, + ipc_addr, + ) + self._start_ple_offload_watchdog() + + def _start_ple_offload_watchdog(self) -> None: + """Fail this worker if the PLE offload process exits unexpectedly. + + Host-mapped pull: a dead producer would leave the captured forward + waiting on its host flag forever; turn that into a fast, visible + engine failure instead of a hang. ``shutdown`` sets the stop marker + before closing the handle, so a normal close is not a failure. + """ + import threading + + handle = self._ple_offload_worker_handle + self._ple_offload_stopping = False + + def watch() -> None: + returncode = handle.proc.wait() + if getattr(self, "_ple_offload_stopping", False): + return + logger.error( + "PLE offload process %s exited unexpectedly (exitcode %s); " + "failing this worker so the engine does not hang on a PLE wait", + handle.proc.pid, + returncode, + ) + os._exit(1) + + threading.Thread( + target=watch, name="ple-offload-watchdog", daemon=True + ).start() + + def wait_ple_offload_ready(self) -> None: + """Wait until weights and every TP registration are ready.""" + if self._ple_offload_worker_handle is None: + return + + from vllm.v1.ple_offload.worker import PleOffloadWorker + + PleOffloadWorker.wait_for_ready(self._ple_offload_worker_handle) + @property def sleep_mode_backend(self) -> "SleepModeBackend": if self._sleep_mode_backend is None: @@ -1152,6 +1295,7 @@ class Worker(WorkerBase): ) -> ModelRunnerOutput | AsyncModelRunnerOutput: return self.model_runner.sample_tokens(grammar_output) + @fail_ple_execute_model @torch.inference_mode() @with_gpu_sync_check def execute_model( @@ -1487,6 +1631,7 @@ class Worker(WorkerBase): self.model_runner.reset_lora_state() def shutdown(self) -> None: + self._ple_offload_stopping = True gc.unfreeze() # Bounds-guard instrumentation: synchronous final read of the guard counters, so diff --git a/vllm_ple_offload_bootstrap/__init__.py b/vllm_ple_offload_bootstrap/__init__.py new file mode 100644 index 0000000..60ddda4 --- /dev/null +++ b/vllm_ple_offload_bootstrap/__init__.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Minimal PLE child bootstrap; intentionally imports no torch or vLLM.""" + +import importlib.abc +import importlib.machinery +import os +import pickle +import struct +import sys +import threading +from multiprocessing.connection import Connection +from types import ModuleType +from typing import Any + +_FAULT = None +if os.environ.get("VLLM_PLE_TEST_FAULT") is not None: + from vllm_ple_test_fault import ( + _FAULT as _FAULT, + _block_on_harness_gate, + ) + + +def _send_partial_ready_frame(fd: int) -> tuple[int, int]: + """Write the complete header and an incomplete body in one syscall.""" + body = pickle.dumps( + {"status": "READY", "layer_names": []}, + protocol=pickle.HIGHEST_PROTOCOL, + ) + frame = struct.pack("!i", len(body)) + body + fragment = frame[: 4 + max(1, len(body) // 2)] + try: + return os.write(fd, fragment), len(frame) + except OSError: + return 0, len(frame) + + +def _send_partial_ready_header(fd: int) -> tuple[int, int]: + """Write fewer than the four required header bytes in one syscall.""" + header = struct.pack("!i", 1) + try: + return os.write(fd, header[:2]), len(header) + except OSError: + return 0, len(header) + + +def _read_exact(fd: int, size: int) -> bytes: + chunks = bytearray() + while len(chunks) < size: + chunk = os.read(fd, size - len(chunks)) + if not chunk: + raise EOFError("PLE launch payload was truncated") + chunks.extend(chunk) + return bytes(chunks) + + +def _recv_launch_payload(fd: int) -> dict: + size = struct.unpack("!Q", _read_exact(fd, 8))[0] + payload = pickle.loads(_read_exact(fd, size)) + if not isinstance(payload, dict): + raise RuntimeError("invalid PLE bootstrap payload") + return payload + + +def _fires( + cut: str, + *, + unit: str = "bootstrap-launch", + transaction: str = "launch:1", + bytes_written: int | None = None, + frame_size: int | None = None, +) -> bool: + return _FAULT is not None and _FAULT.fires( + cut, + role="bootstrap", + unit=unit, + eligible=True, + transaction_id=transaction, + correlation_id="launch:1", + bytes_written=bytes_written, + frame_size=frame_size, + ) + + +def _write_and_ack_partial( + fd: int, + cut: str, + sender: Any, +) -> bool: + """A short or failed prefix write cannot acknowledge or enter the park.""" + if _FAULT is None or _FAULT.cut != cut: + # Disarmed or armed for another cut: the control stream is untouched. + return False + written, frame_size = sender(fd) + if written <= 0 or written >= frame_size: + return False + expected = 2 if cut == "launch_partial_ready_header" else None + if expected is not None and written != expected: + return False + if expected is None and written <= 4: + return False + return _fires( + cut, + bytes_written=written, + frame_size=frame_size, + ) + + +class _FaultImportLoader(importlib.abc.Loader): + def __init__(self, delegate: importlib.abc.Loader) -> None: + self._delegate = delegate + + def create_module(self, spec): + create = getattr(self._delegate, "create_module", None) + return None if create is None else create(spec) + + def exec_module(self, module: ModuleType) -> None: + if _fires( + "launch_block_worker_import", + unit="bootstrap-import", + transaction="module:vllm.v1.ple_offload.worker", + ): + _block_on_harness_gate() + self._delegate.exec_module(module) + + +class _FaultImportFinder(importlib.abc.MetaPathFinder): + _TARGET = "vllm.v1.ple_offload.worker" + + def find_spec(self, fullname: str, path, target=None): + if fullname != self._TARGET: + return None + spec = importlib.machinery.PathFinder.find_spec(fullname, path, target) + if spec is None or spec.loader is None: + raise ImportError(f"cannot find {fullname}") + if not isinstance(spec.loader, importlib.abc.Loader): + raise ImportError(f"{fullname} has no executable loader") + spec.loader = _FaultImportLoader(spec.loader) + return spec + + +def _install_import_gate() -> None: + target = _FaultImportFinder._TARGET + if target in sys.modules: + raise RuntimeError("PLE worker import gate was a cache hit") + sys.meta_path.insert(0, _FaultImportFinder()) + + +def main() -> None: + """Receive launch state before importing torch or vLLM.""" + if len(sys.argv) != 3: + raise SystemExit("expected inherited PLE control and parent-death fds") + control_fd = int(sys.argv[1]) + death_fd = int(sys.argv[2]) + if control_fd == death_fd: + raise RuntimeError("PLE control and parent-death fds must be distinct") + + os.set_inheritable(control_fd, False) + os.set_inheritable(death_fd, False) + death_pipe = Connection(death_fd) + + def monitor_parent() -> None: + try: + death_pipe.recv() + except (EOFError, OSError): + pass + os._exit(1) + + threading.Thread( + target=monitor_parent, + daemon=True, + name="PleOffloadBootstrapDeathMonitor", + ).start() + + if _fires("launch_park_before_payload_read"): + threading.Event().wait() + + # An armed reconstruction payload invokes its reducer inside this loads(). + payload = _recv_launch_payload(control_fd) + + if _write_and_ack_partial( + control_fd, + "launch_partial_ready", + _send_partial_ready_frame, + ): + threading.Event().wait() + + if _write_and_ack_partial( + control_fd, + "launch_partial_ready_header", + _send_partial_ready_header, + ): + threading.Event().wait() + + if _fires("launch_silent_ready"): + threading.Event().wait() + + ready_pipe = Connection(control_fd) + vllm_config = pickle.loads(payload["vllm_config_payload"]) + + if _FAULT is not None and _FAULT.cut == "launch_block_worker_import": + _install_import_gate() + from vllm.v1.ple_offload.worker import PleOffloadWorker + + PleOffloadWorker.proc_main( + vllm_config=vllm_config, + num_workers=int(payload["num_workers"]), + ipc_addr=str(payload["ipc_addr"]), + ready_pipe=ready_pipe, + ) diff --git a/vllm_ple_offload_bootstrap/__main__.py b/vllm_ple_offload_bootstrap/__main__.py new file mode 100644 index 0000000..5b3ca57 --- /dev/null +++ b/vllm_ple_offload_bootstrap/__main__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fresh-interpreter entry point for the PLE CPU-offload worker.""" + +from vllm_ple_offload_bootstrap import main + +main() -- 2.43.0