diff --git a/tests/entrypoints/launchers/test_launch_cli.py b/tests/entrypoints/launchers/test_launch_cli.py deleted file mode 100644 index 55093c5685..0000000000 --- a/tests/entrypoints/launchers/test_launch_cli.py +++ /dev/null @@ -1,1436 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for the `vllm launch` CLI subcommand.""" - -import argparse -import hashlib -import json -import os -import signal -import subprocess -import sys -from contextlib import asynccontextmanager -from pathlib import Path -from typing import Any -from unittest.mock import create_autospec, patch - -import pytest - -import vllm.snapshot.runtime as snapshot_runtime -import vllm.snapshot.server as snapshot_server -from vllm.entrypoints.cli import snapshot as snapshot_cli -from vllm.entrypoints.cli.launch import ( - LaunchSubcommand, - RenderSubcommand, - cmd_init, -) -from vllm.snapshot.controller import ( - LocalSnapshotTools, - SnapshotCreateError, - SnapshotRestoreError, - _child_engine_argv, - create_snapshot, - restore_snapshot, -) -from vllm.snapshot.manifest import ( - SnapshotCompatibilityError, - SnapshotManifest, - SnapshotRuntimeIdentity, - SnapshotSecurityError, - read_manifest, - validate_artifact_root, - validate_identity, - write_manifest_atomic, -) -from vllm.snapshot.runtime import ProcessInventory, _decode_endpoint, _TcpSocketRecord -from vllm.snapshot.server import SnapshotCanaryError -from vllm.snapshot.types import Oracle, oracles_match -from vllm.utils.argparse_utils import FlexibleArgumentParser - - -@pytest.fixture -def launch_parser(): - parser = FlexibleArgumentParser(description="test") - subparsers = parser.add_subparsers(required=False, dest="subparser") - LaunchSubcommand().subparser_init(subparsers) - return parser - - -def test_subcommand_name(): - assert LaunchSubcommand().name == "launch" - - -def test_cmd_init_returns_subcommand(): - result = cmd_init() - assert len(result) == 1 - assert isinstance(result[0], LaunchSubcommand) - - -# -- Parsing: `vllm launch render` -- - - -def test_parse_launch_render(launch_parser): - args = launch_parser.parse_args(["launch", "render", "--model", "test-model"]) - assert args.launch_component == "render" - - -def test_parse_launch_requires_component(launch_parser): - with pytest.raises(SystemExit): - launch_parser.parse_args(["launch", "--model", "test-model"]) - - -def test_parse_launch_invalid_component(launch_parser): - with pytest.raises(SystemExit): - launch_parser.parse_args(["launch", "unknown", "--model", "test-model"]) - - -# -- Dispatch -- - - -def test_cmd_launch_render_calls_run(): - args = argparse.Namespace(model_tag=None, model="test-model") - with patch("vllm.entrypoints.cli.launch.uvloop.run") as mock_uvloop_run: - RenderSubcommand.cmd(args) - mock_uvloop_run.assert_called_once() - - -def test_cmd_launch_model_tag_overrides(): - args = argparse.Namespace( - model_tag="tag-model", - model="original-model", - launch_command=lambda a: None, - ) - LaunchSubcommand.cmd(args) - assert args.model == "tag-model" - - -def test_cmd_launch_model_tag_none(): - args = argparse.Namespace( - model_tag=None, - model="original-model", - launch_command=lambda a: None, - ) - LaunchSubcommand.cmd(args) - assert args.model == "original-model" - - -def test_cmd_dispatches(): - called = {} - - def fake_dispatch(args): - called["args"] = args - - args = argparse.Namespace(launch_command=fake_dispatch) - LaunchSubcommand.cmd(args) - assert "args" in called - - -# -- Module registration -- - - -def test_subparser_init_returns_parser(): - parser = FlexibleArgumentParser(description="test") - subparsers = parser.add_subparsers(required=False, dest="subparser") - result = LaunchSubcommand().subparser_init(subparsers) - assert isinstance(result, FlexibleArgumentParser) - - -def test_launch_registered_in_main(): - """Verify that launch module is importable as a CLI module.""" - import vllm.entrypoints.cli.launch as launch_module - - assert hasattr(launch_module, "cmd_init") - subcmds = launch_module.cmd_init() - assert any(s.name == "launch" for s in subcmds) - - -_MODEL_REVISION = "c1899de289a04d12100db370d81485cdf75e47ca" - - -def _run_cli(argv: list[str]) -> None: - from vllm.entrypoints.cli import main as cli_main - - with patch.object(sys, "argv", ["vllm", *argv]): - cli_main.main() - - -def test_snapshot_environment_contract( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -): - from vllm.entrypoints.serve.utils import api_utils - - monkeypatch.setattr(api_utils, "cli_env_setup", pytest.fail) - secret = "snapshot-secret" - monkeypatch.setenv("VLLM_API_KEY", secret) - monkeypatch.setenv("VLLM_USER_SETTING", secret) - monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", secret) - environment = dict(LocalSnapshotTools()._environment_identity()) - assert ( - environment["VLLM_USER_SETTING"] - == hashlib.sha256(b"VLLM_USER_SETTING\0snapshot-secret").hexdigest() - ) - assert ( - environment["VLLM_USER_SETTING"] != environment["VLLM_WORKER_MULTIPROC_METHOD"] - ) - assert secret not in str(environment) - assert "VLLM_API_KEY" not in environment - - monkeypatch.delenv("VLLM_WORKER_MULTIPROC_METHOD") - implicit_worker_default = dict(LocalSnapshotTools()._environment_identity()) - monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") - assert dict(LocalSnapshotTools()._environment_identity()) == implicit_worker_default - - monkeypatch.setenv("VLLM_USER_SETTING", f"{secret}-changed") - changed_environment = dict(LocalSnapshotTools()._environment_identity()) - assert changed_environment["VLLM_USER_SETTING"] != environment["VLLM_USER_SETTING"] - - artifact = tmp_path / "snapshot" - artifact.mkdir(mode=0o700) - write_manifest_atomic(artifact, _manifest(environment=tuple(environment.items()))) - before = os.environ.copy() - - _run_cli(["snapshot", "inspect", str(artifact)]) - - assert os.environ == before - output = capsys.readouterr().out - assert secret not in output - inspected = json.loads(output) - inspected_environment = dict(inspected["environment"]) - assert inspected_environment == environment - - tools = LocalSnapshotTools() - monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity()) - manifest = tools.make_manifest( - argparse.Namespace( - hf_token=secret, - model_tag="Qwen/Qwen3-0.6B", - revision=_MODEL_REVISION, - served_model_name=None, - tokenizer_revision=None, - ), - ( - "Qwen/Qwen3-0.6B", - "--api-key", - secret, - f"{secret}-alternate", - f"--api-key={secret}", - "--hf-token", - secret, - f"--hf-token={secret}", - ), - ProcessInventory(100, (100, 101), (101,), "GPU-abc"), - _oracle(), - tmp_path, - ) - - assert manifest.engine_argv == ( - "Qwen/Qwen3-0.6B", - "--api-key", - "***", - "***", - "--api-key=***", - "--hf-token", - "***", - "--hf-token=***", - ) - - -@pytest.mark.parametrize( - ("configured", "expected"), - [(None, "1"), ("0", "0"), ("1", "1")], -) -def test_snapshot_child_defaults_nccl_ib_without_mutating_parent( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - configured: str | None, - expected: str, -): - if configured is None: - monkeypatch.delenv("NCCL_IB_DISABLE", raising=False) - else: - monkeypatch.setenv("NCCL_IB_DISABLE", configured) - parent_environment = os.environ.copy() - process = create_autospec(subprocess.Popen, instance=True) - process.pid = 100 - popen = create_autospec(subprocess.Popen, return_value=process) - monkeypatch.setattr(snapshot_runtime.subprocess, "Popen", popen) - - LocalSnapshotTools().launch_child(tmp_path, ("model",)) - - assert popen.call_args.kwargs["env"]["NCCL_IB_DISABLE"] == expected - assert os.environ == parent_environment - - -def test_snapshot_environment_identity_applies_effective_nccl_ib_default( - monkeypatch: pytest.MonkeyPatch, -): - monkeypatch.delenv("NCCL_IB_DISABLE", raising=False) - implicit = LocalSnapshotTools()._environment_identity() - monkeypatch.setenv("NCCL_IB_DISABLE", "1") - explicit_disabled = LocalSnapshotTools()._environment_identity() - monkeypatch.setenv("NCCL_IB_DISABLE", "0") - explicit_enabled = LocalSnapshotTools()._environment_identity() - - assert implicit == explicit_disabled - assert explicit_enabled != implicit - - -@pytest.mark.parametrize( - "case", - [ - (("--hf_token", "SECRET"), ("--hf_token", "***")), - (("--hf_token=SECRET",), ("--hf_token=***",)), - (("--api_key", "SECRET"), ("--api_key", "***")), - (("--api_key=SECRET",), ("--api_key=***",)), - (("--hf-tok", "SECRET"), ("--hf-tok", "***")), - (("--hf-tok=SECRET",), ("--hf-tok=***",)), - (("--hf_tok", "SECRET"), ("--hf_tok", "***")), - (("--api-k", "SECRET", "SECOND"), ("--api-k", "***", "***")), - (("--api-k=SECRET",), ("--api-k=***",)), - (("--hf-overrides", "{}"), ("--hf-overrides", "{}")), - ], -) -def test_snapshot_manifest_redacts_accepted_secret_options(case): - engine_argv, expected = case - tools = create_autospec(LocalSnapshotTools, instance=True) - tools.current_identity.return_value = _runtime_identity() - tools._artifact_bytes.return_value = 0 - manifest = LocalSnapshotTools.make_manifest( - tools, - snapshot_server.parse_vllm_args(["model", *engine_argv]), - engine_argv, - ProcessInventory(100, (100, 101), (101,), "GPU-abc"), - _oracle(), - Path(), - ) - - assert manifest.engine_argv == expected - if "SECRET" in " ".join(engine_argv): - assert "***" in " ".join(manifest.engine_argv) - assert "SECRET" not in " ".join(manifest.engine_argv) - - -def parse_snapshot(*argv: str): - parser = FlexibleArgumentParser() - subparsers = parser.add_subparsers(dest="subparser", required=True) - snapshot_cli.SnapshotSubcommand( - create_requested=argv[:1] == ("create",) - ).subparser_init(subparsers) - return parser.parse_args(["snapshot", *argv]) - - -@pytest.mark.parametrize("option", ["--snapshot-dir", "--snapshot_dir", "--snapshot-d"]) -@pytest.mark.parametrize("joined", [False, True]) -def test_snapshot_create_forwards_only_engine_options(option, joined, monkeypatch): - directory = [f"{option}=/tmp/snapshot"] if joined else [option, "/tmp/snapshot"] - argv = ["create", "Qwen/Qwen3-0.6B", "--revision", _MODEL_REVISION, *directory] - args = parse_snapshot(*argv) - monkeypatch.setattr(sys, "argv", ["vllm", "snapshot", *argv]) - - child = snapshot_server.parse_vllm_args(list(_child_engine_argv(args, None))) - - assert child.model == "Qwen/Qwen3-0.6B" - assert child.revision == _MODEL_REVISION - assert child.enable_sleep_mode - - -@pytest.mark.parametrize("model_in_cli", [False, True]) -def test_snapshot_create_config_preserves_cli_precedence( - tmp_path, monkeypatch, model_in_cli -): - config = tmp_path / "engine.yaml" - config.write_text("model: config-model\nmax-model-len: 512\ndtype: float16\n") - model = ["Qwen/Qwen3-0.6B"] if model_in_cli else [] - argv = [ - "create", - *model, - "--snapshot-dir", - str(tmp_path / "snapshot"), - "--config", - str(config), - "--revision", - _MODEL_REVISION, - "--max-model-len", - "1024", - ] - args = parse_snapshot(*argv) - monkeypatch.setattr(sys, "argv", ["vllm", "snapshot", *argv]) - child = snapshot_server.parse_vllm_args(list(_child_engine_argv(args, None))) - - assert ( - (args.model_tag or args.model) - == child.model - == ("Qwen/Qwen3-0.6B" if model_in_cli else "config-model") - ) - assert args.max_model_len == child.max_model_len == 1024 - assert args.dtype == child.dtype == "float16" - - -@pytest.mark.parametrize( - ("probe", "status", "output"), - [ - (None, 0, ""), - ("nvidia-smi", 1, ""), - ("find", 1, ""), - ("ss", 1, ""), - ("nvidia-smi", 0, "12345"), - ("find", 0, "link_remap.fixture"), - ("ss", 0, "LISTEN"), - ], -) -def test_snapshot_e2e_cleanup_requires_successful_empty_probes(probe, status, output): - script = ( - Path(__file__).resolve().parents[3] - / ".buildkite/scripts/initialized-snapshot-e2e.sh" - ) - if not script.exists(): - pytest.skip("snapshot E2E script is not packaged in this test image") - source = script.read_text() - # Exercise the real probes without entering Docker/CRIU setup. - functions = source[source.index("link_remaps() {") : source.index("wait_clean() {")] - shell = ( - """set -Eeuo pipefail -GPU_UUID=fixture PORT_ONE=18001 PORT_TWO=18002 LINK_REMAP_BASELINE='' -SNAPSHOT_PIDS=(999999999) -nvidia-smi() { :; } -find() { :; } -ss() { :; } -""" - + functions - ) - if probe: - shell += f'\n{probe}() {{ printf %s "{output}"; return {status}; }}\n' - shell += "\nif state_clean; then exit 0; else exit 1; fi\n" - - result = subprocess.run(["bash", "-c", shell], capture_output=True, timeout=10) - - assert result.returncode == (0 if probe is None else 1), result.stderr - - -@pytest.mark.parametrize("container_attempted", [False, True]) -def test_snapshot_e2e_cleanup_checks_state_only_after_container_attempt( - container_attempted: bool, -): - script = ( - Path(__file__).resolve().parents[3] - / ".buildkite/scripts/initialized-snapshot-e2e.sh" - ) - if not script.exists(): - pytest.skip("snapshot E2E script is not packaged in this test image") - source = script.read_text() - cleanup = source[source.index("cleanup() {") : source.index("on_exit() {")] - shell = ( - "set -Eeuo pipefail\n" - "CONTAINER_NAME='' PORT_TWO=18002\n" - f"CHECK_CLEANUP_STATE={int(container_attempted)}\n" - "wait_clean() { echo checked-state; return 1; }\n" - + cleanup - + "\nif cleanup; then exit 0; else exit 1; fi\n" - ) - - result = subprocess.run(["bash", "-c", shell], capture_output=True, timeout=10) - - assert result.returncode == int(container_attempted), result.stderr - assert (b"checked-state" in result.stdout) == container_attempted - - -def test_snapshot_create_cli_accepts_only_pinned_compact_mode( - monkeypatch: pytest.MonkeyPatch, -): - monkeypatch.setattr(snapshot_cli.platform, "system", lambda: "Linux") - monkeypatch.setattr(snapshot_cli.platform, "machine", lambda: "x86_64") - base = ("create", "Qwen/Qwen3-0.6B", "--snapshot-dir=/tmp/snapshot") - compact = parse_snapshot(*base, "--revision", _MODEL_REVISION) - - snapshot_cli.validate_create_args(compact) - assert compact.model_tag == "Qwen/Qwen3-0.6B" - - dispatched: list[argparse.Namespace] = [] - monkeypatch.setattr(snapshot_cli, "run_create", dispatched.append) - _run_cli(["snapshot", *base, "--revision", _MODEL_REVISION]) - assert dispatched[0].revision == _MODEL_REVISION - - -@pytest.mark.parametrize( - ("attribute", "value", "message"), - [ - ("prefill_context_parallel_size", 2, "prefill context parallel size 1"), - ("model_tag", None, "immutable remote model"), - ("logprobs_mode", "raw_logits", "log-probability mode"), - ("speculative_config", {}, "does not support speculative decoding"), - ("spec_method", "ngram", "does not support speculative decoding"), - ("spec_model", "draft-model", "does not support speculative decoding"), - ("spec_tokens", 3, "does not support speculative decoding"), - ], -) -def test_snapshot_create_rejects_invalid_config( - attribute: str, - value: object, - message: str, - tmp_path: Path, -): - args = parse_snapshot( - "create", - "Qwen/Qwen3-0.6B", - "--snapshot-dir=/tmp/snapshot", - "--revision", - _MODEL_REVISION, - ) - if attribute == "model_tag": - local_model = tmp_path / "model" - local_model.mkdir() - value = str(local_model) - setattr(args, attribute, value) - with pytest.raises(ValueError, match=message): - snapshot_cli.validate_create_args(args) - - -def _oracle( - token_ids: tuple[int, ...] = (12095,), - text: str = " Paris", - logprob: float = -0.125, -) -> Oracle: - return Oracle(token_ids, text, logprob) - - -def _manifest(**changes: object) -> SnapshotManifest: - manifest = SnapshotManifest( - schema_version=1, - boundary="post-engine-init-reloadable-state-released", - created_at="2026-08-06T00:00:00Z", - artifact_bytes=8, - vllm_version="0.1.dev1+snapshot", - python_version="3.12.3", - torch_version="2.9.0", - cuda_runtime="12.9", - driver_version="575.57.08", - criu_version="4.1", - cuda_checkpoint_sha256="a" * 64, - kernel_release="6.8.0", - host_id="host-a", - gpu_name="NVIDIA A10", - gpu_uuid="GPU-abc", - model="Qwen/Qwen3-0.6B", - served_model_name="Qwen/Qwen3-0.6B", - model_revision="model-sha", - tokenizer_revision="tokenizer-sha", - engine_argv=("Qwen/Qwen3-0.6B",), - environment=(("VLLM_USE_V1", "1"),), - process_tree=(100, 101), - cuda_holders=(101,), - oracle_token_ids=(12095,), - oracle_text=" Paris", - oracle_sampled_token_logprob=-0.125, - ) - return SnapshotManifest.model_validate( - {**manifest.model_dump(mode="python"), **changes}, strict=True - ) - - -def _runtime_identity(**changes: object) -> Any: - fields = SnapshotRuntimeIdentity.model_fields.keys() - payload = _manifest().model_dump(mode="python", include=fields) - return SnapshotRuntimeIdentity.model_validate({**payload, **changes}, strict=True) - - -def _fake_snapshot_tools() -> Any: - tools = create_autospec(LocalSnapshotTools, instance=True) - tools.launch_child.return_value = 100 - tools.wait_ready.return_value = _oracle() - tools.inventory.return_value = ProcessInventory(100, (100, 101), (101,), "GPU-abc") - - tools.current_identity.side_effect = lambda _gpu_uuid: _runtime_identity() - tools.restore.return_value = 100 - tools.request_oracle.return_value = _oracle() - return tools - - -@pytest.mark.parametrize("value", ["0", "-1", "nan", "inf"]) -def test_snapshot_timeout_must_be_positive_and_finite( - value: str, monkeypatch: pytest.MonkeyPatch -): - monkeypatch.setenv("VLLM_SNAPSHOT_TIMEOUT_S", value) - with pytest.raises(ValueError, match="positive finite number"): - LocalSnapshotTools() - - -def test_snapshot_create_abort_reaps_after_process_group_cleanup( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - tools = LocalSnapshotTools() - process = subprocess.Popen( - [sys.executable, "-c", "raise SystemExit(1)"], - stdin=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - start_new_session=True, - ) - tools._children[process.pid] = process - (tmp_path / "child.log").write_text("root failed") - killed: list[tuple[int, signal.Signals]] = [] - monkeypatch.setattr( - os, - "killpg", - lambda pid, sig: killed.append((pid, sig)), - ) - monkeypatch.setattr(os, "waitid", lambda *_args: object(), raising=False) - - try: - with pytest.raises(SnapshotCreateError, match="root failed"): - tools.wait_ready(tmp_path, process.pid) - assert process.returncode is None - tools.abort_create(process.pid) - finally: - if process.pid in tools._children: - tools.abort_create(process.pid) - - assert killed == [(process.pid, signal.SIGKILL)] - assert process.pid not in tools._children - - -def _supply_waitid(monkeypatch: pytest.MonkeyPatch) -> None: - """`wait_ready` binds Linux-only `os.waitid` before it reads the marker.""" - monkeypatch.setattr(os, "waitid", lambda *_args: None, raising=False) - - -def test_snapshot_wait_ready_reads_back_the_oracle_the_child_wrote( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - _supply_waitid(monkeypatch) - snapshot_server.write_ready_atomic(tmp_path / "ready.json", _oracle()) - - assert LocalSnapshotTools().wait_ready(tmp_path, os.getpid()) == _oracle() - - -@pytest.mark.parametrize( - "payload", - [ - {"token_ids": [12095], "sampled_token_logprob": -0.125}, - {"token_ids": ["12095"], "text": " Paris", "sampled_token_logprob": -0.125}, - ], -) -def test_snapshot_wait_ready_rejects_a_malformed_marker( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, payload: dict[str, object] -): - _supply_waitid(monkeypatch) - (tmp_path / "ready.json").write_text(json.dumps(payload)) - - with pytest.raises(SnapshotCreateError, match="ready marker invalid"): - LocalSnapshotTools().wait_ready(tmp_path, os.getpid()) - - -def test_snapshot_create_abort_rejects_surviving_process_group( - monkeypatch: pytest.MonkeyPatch, -): - tools = LocalSnapshotTools() - process = create_autospec(subprocess.Popen, instance=True) - process.wait.side_effect = subprocess.TimeoutExpired("snapshot child", 10) - tools._children[100] = process - monkeypatch.setattr(os, "killpg", lambda _pid, _sig: None) - - with pytest.raises(SnapshotCreateError, match="survived SIGKILL"): - tools.abort_create(100) - - -@pytest.mark.asyncio -async def test_snapshot_child_rejects_mismatch_before_publication( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - from vllm.entrypoints.launchers.api_server import entry as api_server - - canaries = iter((_oracle(), _oracle((42,), " Lyon"))) - - @asynccontextmanager - async def engine_context(): - yield object() - - async def next_canary(_engine: object) -> Oracle: - return next(canaries) - - async def return_none(_engine: object) -> None: - return None - - monkeypatch.setattr( - api_server, "build_async_engine_client", lambda _args: engine_context() - ) - monkeypatch.setattr(snapshot_server, "run_engine_canary", next_canary) - monkeypatch.setattr(snapshot_server, "_release_reloadable_state", return_none) - monkeypatch.setattr(snapshot_server, "_restore_reloadable_state", return_none) - monkeypatch.setattr(snapshot_server, "detach_snapshot_streams", lambda: None) - monkeypatch.setattr( - snapshot_server, - "write_ready_atomic", - lambda *_args: pytest.fail("mismatched snapshot was published"), - ) - control = snapshot_server.ControlArgs( - ready_file=tmp_path / "ready.json", - release_file=tmp_path / "release.json", - release_timeout_s=1, - ) - with pytest.raises(SnapshotCanaryError, match="rehearsal changed"): - await snapshot_server.run_vllm_snapshot_child(control, argparse.Namespace()) - - -def test_snapshot_release_marker_ignores_unknown_keys(tmp_path: Path): - marker = tmp_path / "release.json" - marker.write_text(json.dumps({"release": True, "port": 8000, "unknown": 1})) - - listener = snapshot_server.read_release_marker(marker) - - assert listener == snapshot_server.ListenerConfig(host=None, port=8000) - - -@pytest.mark.parametrize( - "payload", - [ - ["release"], - {"port": 8000}, - {"release": "true", "port": 8000}, - {"release": 1, "port": 8000}, - {"release": True, "port": 0}, - {"release": True, "port": 70000}, - {"release": True, "port": True}, - {"release": True, "port": "8000"}, - {"release": True, "port": 8000, "host": 5}, - ], -) -def test_snapshot_release_marker_rejects_an_invalid_payload( - tmp_path: Path, payload: object -): - marker = tmp_path / "release.json" - marker.write_text(json.dumps(payload)) - - with pytest.raises(snapshot_server.SnapshotBarrierError, match="release marker"): - snapshot_server.read_release_marker(marker) - - -def test_snapshot_create_rolls_back_failed_dump(tmp_path: Path): - target = tmp_path / "snapshot" - tools = _fake_snapshot_tools() - tools.dump.side_effect = RuntimeError("dump failed") - with pytest.raises(RuntimeError, match="dump failed"): - create_snapshot( - argparse.Namespace(snapshot_dir=str(target), model_tag="Qwen/Qwen3-0.6B"), - tools=tools, - ) - tools.abort_create.assert_called_once_with(100) - assert not target.exists() - - -@pytest.mark.parametrize( - "failure_kind", ["called-process", "timeout", "missing", "unreadable"] -) -def test_snapshot_create_preserves_bounded_criu_diagnostics_after_cleanup( - tmp_path: Path, failure_kind: str, caplog: pytest.LogCaptureFixture -): - target = tmp_path / "snapshot" - tools = _fake_snapshot_tools() - criu_tools = LocalSnapshotTools() - criu_tools.plugin_dir = tmp_path - original_run = criu_tools._run - criu_tools._privileged = lambda: [] # type: ignore[method-assign] - captured = "captured stderr " + "x" * 9000 - if failure_kind == "timeout": - primary: subprocess.SubprocessError = subprocess.TimeoutExpired( - ["criu", "dump"], 1, output=b"captured stdout\xff", stderr=captured.encode() - ) - else: - primary = subprocess.CalledProcessError( - 1, ["criu", "dump"], output="captured stdout", stderr=captured - ) - - def run(command: list[str], **_kwargs: object): - if "tail" not in command: - raise primary - if failure_kind == "unreadable": - raise PermissionError("log is root-owned") - if failure_kind == "timeout": - return subprocess.CompletedProcess(command, 0, "timeout log tail", "") - return original_run(command, **_kwargs) - - criu_tools._run = run # type: ignore[method-assign] - - def fail_dump(workdir: Path, _inventory: ProcessInventory) -> None: - images = workdir / "images" - images.mkdir() - if failure_kind != "missing": - (images / "dump.log").write_text( - "omitted dump prefix" + "d" * 9000 + "distinctive dump diagnostic" - ) - criu_tools._criu("dump", workdir, []) - - tools.dump.side_effect = fail_dump - with pytest.raises(type(primary)) as excinfo: - create_snapshot( - argparse.Namespace(snapshot_dir=str(target), model_tag="Qwen/Qwen3-0.6B"), - tools=tools, - ) - - assert excinfo.value is primary - assert not target.exists() - diagnostics = "\n".join(caplog.messages) - assert "captured stdout" in diagnostics - assert "captured stderr" not in diagnostics - assert "x" * 100 in diagnostics - assert len(diagnostics) < 17000 - if failure_kind in {"missing", "unreadable"}: - assert "dump.log unavailable" in diagnostics - expected_error = ( - "CalledProcessError" if failure_kind == "missing" else "PermissionError" - ) - assert expected_error in diagnostics - elif failure_kind == "called-process": - assert "distinctive dump diagnostic" in diagnostics - assert "omitted dump prefix" not in diagnostics - else: - assert "timeout log tail" in diagnostics - if failure_kind == "timeout": - assert "captured stdout�" in diagnostics - - -def test_snapshot_criu_failure_selects_restore_log( - tmp_path: Path, caplog: pytest.LogCaptureFixture -): - tools = LocalSnapshotTools() - tools.plugin_dir = tmp_path - tools._privileged = lambda: [] # type: ignore[method-assign] - artifact = tmp_path / "snapshot" - images = artifact / "images" - images.mkdir(parents=True) - (images / "restore.log").write_text( - "omitted restore prefix" + "r" * 9000 + "distinctive restore diagnostic" - ) - primary = subprocess.CalledProcessError(1, ["criu", "restore"]) - original_run = tools._run - tailed: list[str] = [] - - def run(command: list[str], **_kwargs: object): - if "tail" not in command: - raise primary - tailed.append(command[-1]) - return original_run(command, **_kwargs) - - tools._run = run # type: ignore[method-assign] - - with pytest.raises(subprocess.CalledProcessError) as excinfo: - tools._criu("restore", artifact, []) - - assert excinfo.value is primary - assert tailed == [str(images / "restore.log")] - diagnostics = "\n".join(caplog.messages) - assert "distinctive restore diagnostic" in diagnostics - assert "omitted restore prefix" not in diagnostics - - -def _restore_fixture(tmp_path: Path) -> argparse.Namespace: - artifact = tmp_path / "snapshot" - artifact.mkdir(mode=0o700) - write_manifest_atomic(artifact, _manifest()) - return argparse.Namespace(snapshot_dir=str(artifact), host="127.0.0.1", port=9000) - - -def test_snapshot_restore_rejects_identity_mismatch(tmp_path: Path): - args = _restore_fixture(tmp_path) - - identity_mismatch = _fake_snapshot_tools() - identity_mismatch.current_identity.side_effect = lambda _gpu_uuid: ( - _runtime_identity(vllm_version="different") - ) - with pytest.raises(SnapshotCompatibilityError, match="vllm_version"): - restore_snapshot(args, tools=identity_mismatch) - identity_mismatch.restore.assert_not_called() - - -def test_snapshot_identity_mismatch_names_every_differing_field(): - with pytest.raises(SnapshotCompatibilityError) as mismatch: - validate_identity( - _manifest(), - _runtime_identity(vllm_version="different", gpu_uuid="GPU-other"), - ) - - assert str(mismatch.value) == "snapshot mismatch: vllm_version, gpu_uuid" - - -def test_snapshot_restore_cleans_oracle_mismatch(tmp_path: Path): - args = _restore_fixture(tmp_path) - oracle_mismatch = _fake_snapshot_tools() - oracle_mismatch.request_oracle.return_value = _oracle(logprob=-0.5) - with pytest.raises(SnapshotRestoreError, match="oracle mismatch"): - restore_snapshot(args, tools=oracle_mismatch) - oracle_mismatch.cleanup.assert_called_once() - oracle_mismatch.complete_restore.assert_not_called() - assert oracles_match(_oracle(logprob=-100.0), _oracle(logprob=-100.0005)) - assert not oracles_match(_oracle(logprob=-100.0), _oracle(logprob=-100.002)) - - -@pytest.mark.parametrize( - ("invalid_update", "diagnostic"), - [ - ({"unexpected": True}, "unexpected"), - ({"schema_version": True}, "schema_version"), - ({"boundary": "unknown"}, "boundary"), - ({"process_tree": [100, 100]}, "root"), - ({"cuda_holders": [999]}, "root"), - ({"oracle_token_ids": [1, 2]}, "oracle_token_ids"), - ({"oracle_sampled_token_logprob": 0}, "oracle_sampled_token_logprob"), - ( - {"oracle_sampled_token_logprob": float("nan")}, - "oracle_sampled_token_logprob", - ), - ], -) -def test_snapshot_manifest_validation( - tmp_path: Path, - invalid_update: dict[str, object], - diagnostic: str, -): - artifact = tmp_path / "snapshot" - artifact.mkdir(mode=0o700) - write_manifest_atomic(artifact, _manifest()) - invalid = _manifest().model_dump(mode="json") - invalid.update(invalid_update) - (artifact / "manifest.json").write_text(json.dumps(invalid)) - with pytest.raises( - SnapshotCompatibilityError, - match=f"invalid snapshot manifest: {diagnostic}", - ): - read_manifest(artifact) - - -def _local_restore_artifact(tmp_path: Path, name: str = "snapshot") -> Path: - artifact = tmp_path / name - artifact.mkdir(mode=0o700) - (artifact / "child.log").write_bytes(b"startup log") - (artifact / "child.log.snapshot-size").write_text("11\n") - return artifact - - -def _mark_snapshot_pids_free(monkeypatch: pytest.MonkeyPatch) -> list[int]: - probed_pids: list[int] = [] - - def missing(pid: int, _signal: int) -> None: - probed_pids.append(pid) - raise ProcessLookupError - - monkeypatch.setattr(snapshot_runtime.os, "kill", missing) - return probed_pids - - -@pytest.mark.parametrize( - ("probe_error", "diagnostic"), - [ - (None, "already occupied: 100"), - (PermissionError(), "already occupied: 100"), - (OSError("probe failed"), "availability probe failed: 100"), - ], -) -def test_snapshot_restore_rejects_unavailable_pid_before_mutation( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - probe_error: OSError | None, - diagnostic: str, -): - args = _restore_fixture(tmp_path) - artifact = Path(args.snapshot_dir) - (artifact / "child.log").write_bytes(b"startup log") - (artifact / "release.json").write_bytes(b"release sentinel") - before = {path.name: path.read_bytes() for path in artifact.iterdir()} - tools = LocalSnapshotTools() - - def fail_probe(_pid: int, _signal: int) -> None: - if probe_error is not None: - raise probe_error - - monkeypatch.setattr(snapshot_runtime.os, "kill", fail_probe) - monkeypatch.setattr(tools, "preflight", lambda *_args: None) - monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity()) - monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) - - with pytest.raises(SnapshotRestoreError, match=diagnostic): - restore_snapshot(args, tools=tools) - assert {path.name: path.read_bytes() for path in artifact.iterdir()} == before - - -@pytest.mark.parametrize("mode", ["pinned", "unpinned", "cleanup-fails"]) -def test_snapshot_restore_terminates_only_a_pinned_tree( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - mode: str, -): - _mark_snapshot_pids_free(monkeypatch) - artifact = _local_restore_artifact(tmp_path) - manifest = _manifest() - root_pid = manifest.process_tree[0] - tools = LocalSnapshotTools() - terminated: list[int] = [] - - def terminate_and_wait(pid: int) -> None: - terminated.append(pid) - if mode == "cleanup-fails": - raise SnapshotRestoreError("restored process cleanup is incomplete") - tools._restored_processes.pop(pid) - - def pin( - _artifact: Path, process_tree: tuple[int, ...], _holders: tuple[int, ...] - ) -> None: - if mode == "unpinned": - raise SnapshotRestoreError( - "restored session does not match the captured process tree" - ) - tools._restored_processes[process_tree[0]] = () - - monkeypatch.setattr(tools, "cleanup", terminate_and_wait) - monkeypatch.setattr(tools, "_pin_restored_tree", pin) - monkeypatch.setattr(tools, "_criu", lambda *_args: None) - monkeypatch.setattr( - tools, - "_run", - lambda command, **_kwargs: subprocess.CompletedProcess(command, 0, "333\n", ""), - ) - - with pytest.raises(SnapshotRestoreError, match="does not match") as failure: - tools.restore(artifact, manifest) - - # A pinned tree is ours: it must be terminated and waited for, not left to - # honour the abort marker. An unpinned PID is never signalled. - assert terminated == ([] if mode == "unpinned" else [root_pid]) - if mode == "cleanup-fails": - # The failure to terminate is reported beside the primary error, and a - # tree that may still be running stays pinned. - assert "restore cleanup failed" in str(failure.value) - assert tools._restored_processes == {root_pid: ()} - else: - assert tools._restored_processes == {} - assert json.loads((artifact / "release.json").read_text())["release"] is False - - -def test_snapshot_private_path_and_link_remap_security( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - _mark_snapshot_pids_free(monkeypatch) - real = tmp_path / "real" - real.mkdir(mode=0o700) - linked = tmp_path / "linked" - linked.symlink_to(real, target_is_directory=True) - with pytest.raises(SnapshotSecurityError, match="symlink"): - validate_artifact_root(linked, creating=False) - - unsafe = tmp_path / "unsafe" - unsafe.mkdir() - unsafe.chmod(0o777) - with pytest.raises(SnapshotSecurityError, match="world-writable"): - validate_artifact_root(unsafe / "snapshot", creating=True) - - artifact = _local_restore_artifact(tmp_path) - saved_remaps = artifact / "link-remaps" - saved_remaps.mkdir() - (saved_remaps / "link_remap.270").write_bytes(b"semaphore state") - shm_dir = tmp_path / "shm" - shm_dir.mkdir() - target = shm_dir / "link_remap.270" - replacement = tmp_path / "replacement" - replacement.write_bytes(b"semaphore state") - target.symlink_to(replacement) - tools = LocalSnapshotTools() - tools.shm_dir = shm_dir - monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) - with pytest.raises(SnapshotRestoreError, match="conflicting CRIU link remap"): - tools.restore(artifact, _manifest()) - - -@pytest.mark.parametrize("inventory_error", [None, PermissionError("denied")]) -def test_snapshot_dump_rejects_unavailable_link_remap_inventory_before_criu( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - inventory_error: PermissionError | None, -): - tools = LocalSnapshotTools() - workdir = tmp_path / "snapshot" - workdir.mkdir() - tools.shm_dir = tmp_path / "missing-shm" - if inventory_error is not None: - tools.shm_dir.mkdir() - - def fail_inventory(_path: Path): - raise inventory_error - - monkeypatch.setattr(Path, "iterdir", fail_inventory) - monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) - - with pytest.raises(SnapshotCreateError, match="inventory CRIU link remaps"): - tools.dump(workdir, ProcessInventory(100, (100,), (100,), "GPU")) - - -def test_snapshot_restore_rejects_saved_remap_without_prefix( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - _mark_snapshot_pids_free(monkeypatch) - artifact = _local_restore_artifact(tmp_path) - saved = artifact / "link-remaps" - saved.mkdir() - (saved / "270").write_bytes(b"invalid remap name") - tools = LocalSnapshotTools() - tools.shm_dir = tmp_path / "shm" - tools.shm_dir.mkdir() - monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) - - with pytest.raises(SnapshotRestoreError, match="invalid saved link remap"): - tools.restore(artifact, _manifest()) - - assert not tuple(tools.shm_dir.iterdir()) - - -def test_snapshot_link_remap_inventory_matches_only_numeric_suffixes(tmp_path: Path): - for name in ( - "link_remap.0", - "link_remap.270", - "link_remap.", - "link_remap.12a", - "link-remap.12", - "prefix.link_remap.12", - "270", - ): - (tmp_path / name).touch() - tools = LocalSnapshotTools() - tools.shm_dir = tmp_path - - assert tools._link_remap_names() == {"link_remap.0", "link_remap.270"} - - -def test_snapshot_dump_preserves_primary_failure_when_cleanup_inventory_fails( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - tools = LocalSnapshotTools() - workdir = tmp_path / "snapshot" - workdir.mkdir() - inventory_calls = 0 - - def inventory_names() -> set[str]: - nonlocal inventory_calls - inventory_calls += 1 - if inventory_calls > 1: - raise SnapshotCreateError( - "could not inventory CRIU link remaps: PermissionError: denied" - ) - return set() - - monkeypatch.setattr(tools, "_link_remap_names", inventory_names) - - def fail_dump(*_args): - raise RuntimeError("dump failed") - - monkeypatch.setattr(tools, "_criu", fail_dump) - - with pytest.raises(SnapshotCreateError) as failure: - tools.dump(workdir, ProcessInventory(100, (100,), (100,), "GPU")) - - assert "RuntimeError: dump failed" in str(failure.value) - assert "cleanup inventory failed" in str(failure.value) - - -@pytest.mark.parametrize( - ("blocked_state", "message"), - [("tcp", "external established TCP"), ("io_uring", "kernel.io_uring_disabled=1")], -) -def test_snapshot_rejects_unsafe_process_state_before_criu( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, blocked_state: str, message: str -): - tools = LocalSnapshotTools() - monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101)) - monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",)) - monkeypatch.setattr( - tools, - "_descriptor_inventory", - lambda _tree: ( - (101,) if blocked_state == "io_uring" else (), - {41: 101}, - ), - ) - tcp_record = _TcpSocketRecord( - family="AF_INET", - local_raw="3C00000A:D6F6", - remote_raw="01010101:01BB" if blocked_state == "tcp" else "00000000:0000", - inode=41, - ) - monkeypatch.setattr(tools, "_tcp_records", lambda: (tcp_record,)) - with pytest.raises(SnapshotCreateError, match=message) as excinfo: - tools.inventory(100, tmp_path) - if blocked_state == "tcp": - assert "pid 101" in str(excinfo.value) - assert "1.1.1.1:443" in str(excinfo.value) - - -def test_snapshot_rejects_open_infiniband_device_before_criu( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - tools = LocalSnapshotTools() - monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101)) - monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",)) - monkeypatch.setattr( - tools, - "_descriptor_targets", - lambda pid: ("/dev/infiniband/uverbs7",) if pid == 101 else (), - ) - monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) - - with pytest.raises(SnapshotCreateError) as failure: - tools.inventory(100, tmp_path) - - message = str(failure.value) - assert "pid 101" in message - assert "/dev/infiniband/uverbs7" in message - assert "NCCL_IB_DISABLE=1" in message - - -def test_decode_endpoint_families(): - assert _decode_endpoint("AF_INET", "3C00000A:D6F6") == "10.0.0.60:55030" - assert ( - _decode_endpoint("AF_INET6", "00000000000000000000000001000000:01BB") - == "[::1]:443" - ) - - -def test_tcp_records_skips_missing_tables(tmp_path, monkeypatch): - tcp, tcp6 = tmp_path / "tcp", tmp_path / "tcp6" - contents = "header\n0: 0100007F:1F90 00000000:0000 01 0 0 0 0 0 41\n" - tcp.write_text(contents) - tables = (("AF_INET", tcp), ("AF_INET6", tcp6)) - monkeypatch.setattr(snapshot_runtime, "_TCP_TABLES", tables) - - tools = LocalSnapshotTools() - assert tuple(record.family for record in tools._tcp_records()) == ("AF_INET",) - - tcp6.write_text(contents) - families = tuple(record.family for record in tools._tcp_records()) - assert families == ("AF_INET", "AF_INET6") - - -def test_snapshot_manifest_records_external_cache_files(tmp_path: Path): - artifact = tmp_path / "snapshot" - artifact.mkdir(mode=0o700) - recorded = ( - (str(tmp_path / "kernel.so"), "b" * 64), - (str(tmp_path / "launcher.lock"), "locked"), - ) - write_manifest_atomic(artifact, _manifest(external_cache_files=recorded)) - - assert read_manifest(artifact).external_cache_files == recorded - - # A manifest written before this field existed still loads. - legacy = _manifest().model_dump(mode="json") - del legacy["external_cache_files"] - (artifact / "manifest.json").write_text(json.dumps(legacy)) - - assert read_manifest(artifact).external_cache_files == () - - -def _pin_generated_cache_roots(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: - """Keep the resolver off whatever cache roots the host running this has.""" - monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "vllm_cache")) - monkeypatch.setenv("TRITON_CACHE_DIR", str(tmp_path / "triton_cache_default")) - monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", str(tmp_path / "inductor_cache")) - - -def _unreadable(_path: Path) -> str: - raise OSError("cache file is unreadable") - - -def test_snapshot_inventory_records_open_generated_cache_files( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - cache_root = tmp_path / "triton_cache" - cache_root.mkdir() - kernel = cache_root / "kernel.so" - kernel.write_bytes(b"compiled kernel") - lock = cache_root / "launcher.lock" - lock.write_bytes(b"held by pid 100") - weights = tmp_path / "model.safetensors" - weights.write_bytes(b"weights") - artifact = tmp_path / "snapshot" - artifact.mkdir(mode=0o700) - (artifact / "child.log.lock").write_bytes(b"") - _pin_generated_cache_roots(monkeypatch, tmp_path) - monkeypatch.setenv("TRITON_CACHE_DIR", str(cache_root)) - - targets = { - 100: (str(kernel), str(weights), "socket:[41]"), - 101: (str(lock), str(cache_root), str(artifact / "child.log.lock")), - 102: (f"{cache_root / 'evicted.so'} (deleted)",), - } - tools = LocalSnapshotTools() - monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101, 102)) - monkeypatch.setattr(tools, "_descriptor_targets", lambda pid: targets[pid]) - monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",)) - monkeypatch.setattr(tools, "_tcp_records", lambda: ()) - monkeypatch.setattr(tools, "current_identity", lambda _uuid: _runtime_identity()) - - manifest = tools.make_manifest( - argparse.Namespace( - model_tag="Qwen/Qwen3-0.6B", - revision=_MODEL_REVISION, - served_model_name=None, - tokenizer_revision=None, - ), - ("Qwen/Qwen3-0.6B",), - tools.inventory(100, artifact), - _oracle(), - artifact, - ) - - # Weights, sockets, the cache directory itself, files inside the artifact, - # and deleted targets that CRIU remaps are all left out. - assert manifest.external_cache_files == ( - (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()), - (str(lock), "locked"), - ) - - -def test_snapshot_inventory_marks_a_cache_file_it_cannot_read( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - cache_root = tmp_path / "triton_cache" - cache_root.mkdir() - kernel = cache_root / "kernel.so" - kernel.write_bytes(b"compiled kernel") - lock = cache_root / "launcher.lock" - lock.write_bytes(b"held by pid 100") - artifact = tmp_path / "snapshot" - artifact.mkdir(mode=0o700) - _pin_generated_cache_roots(monkeypatch, tmp_path) - monkeypatch.setenv("TRITON_CACHE_DIR", str(cache_root)) - - tools = LocalSnapshotTools() - monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100,)) - monkeypatch.setattr( - tools, "_descriptor_targets", lambda _pid: (str(kernel), str(lock)) - ) - monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("100, GPU-abc",)) - monkeypatch.setattr(tools, "_tcp_records", lambda: ()) - monkeypatch.setattr(tools, "_sha256", _unreadable) - - # A file that cannot be digested keeps its own marker, so the manifest says - # which entries restore can check for existence only. - assert tools.inventory(100, artifact).external_cache_files == ( - (str(kernel), "unreadable"), - (str(lock), "locked"), - ) - - -def test_snapshot_inventory_matches_resolved_cache_roots_not_siblings( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - # Descriptor targets come back kernel-canonical, so the fixture must be too. - base = Path(os.path.realpath(tmp_path)) - real_cache = base / "mnt_cache" - real_cache.mkdir() - kernel = real_cache / "kernel.so" - kernel.write_bytes(b"compiled kernel") - linked_cache = base / "linked_cache" - linked_cache.symlink_to(real_cache, target_is_directory=True) - sibling = base / "vllm_cache-backup" - sibling.mkdir() - stale = sibling / "stale.so" - stale.write_bytes(b"stale kernel") - artifact = base / "snapshot" - artifact.mkdir(mode=0o700) - - _pin_generated_cache_roots(monkeypatch, base) - monkeypatch.setenv("TRITON_CACHE_DIR", str(linked_cache)) - - tools = LocalSnapshotTools() - monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100,)) - monkeypatch.setattr( - tools, "_descriptor_targets", lambda _pid: (str(kernel), str(stale)) - ) - monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("100, GPU-abc",)) - monkeypatch.setattr(tools, "_tcp_records", lambda: ()) - - # A cache root reached through a symlink still covers the file the kernel - # reports by its real path, and a directory that merely shares the root's - # name prefix is not part of the cache at all. - assert tools.inventory(100, artifact).external_cache_files == ( - (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()), - ) - - -def test_snapshot_restore_rejects_rotated_cache_file_before_mutation( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - _mark_snapshot_pids_free(monkeypatch) - kernel = tmp_path / "kernel.so" - kernel.write_bytes(b"compiled kernel") - lock = tmp_path / "launcher.lock" - lock.write_bytes(b"held by pid 100") - manifest = _manifest( - external_cache_files=( - (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()), - (str(lock), "locked"), - ) - ) - artifact = _local_restore_artifact(tmp_path) - before = {path.name: path.read_bytes() for path in artifact.iterdir()} - tools = LocalSnapshotTools() - monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) - - kernel.write_bytes(b"recompiled kernel") - with pytest.raises(SnapshotRestoreError, match="changed since capture") as changed: - tools.restore(artifact, manifest) - assert str(kernel) in str(changed.value) - - # Restoring the digest gets past the first entry, so the lock file is only - # checked for existence: the second failure proves the first one passed. - kernel.write_bytes(b"compiled kernel") - lock.unlink() - with pytest.raises(SnapshotRestoreError, match="vanished since capture") as gone: - tools.restore(artifact, manifest) - assert str(lock) in str(gone.value) - - assert {path.name: path.read_bytes() for path in artifact.iterdir()} == before - - -def test_snapshot_restore_rejects_a_cache_file_it_cannot_read( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - _mark_snapshot_pids_free(monkeypatch) - kernel = tmp_path / "kernel.so" - kernel.write_bytes(b"compiled kernel") - manifest = _manifest( - external_cache_files=( - (str(kernel), hashlib.sha256(b"compiled kernel").hexdigest()), - ) - ) - artifact = _local_restore_artifact(tmp_path) - tools = LocalSnapshotTools() - monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) - monkeypatch.setattr(tools, "_sha256", _unreadable) - - # A recorded digest that cannot be read again is a failed check, not a file - # to wave through on existence. - with pytest.raises( - SnapshotRestoreError, match="unreadable since capture" - ) as blocked: - tools.restore(artifact, manifest) - assert str(kernel) in str(blocked.value) - - -def test_snapshot_runtime_installer_logs_the_failing_exit_code(): - installer = Path(__file__).parents[3] / "tools" / "install_snapshot_runtime.sh" - if not installer.is_file(): - pytest.skip("installer script is not shipped in this test image") - timing = [ - line - for line in installer.read_text().splitlines() - if line.startswith(("_snapshot_install_started=", "trap ")) - ] - assert len(timing) == 2 - - result = subprocess.run( - ["bash", "-c", "\n".join([*timing, "exit 7"])], - capture_output=True, - text=True, - ) - - assert result.returncode == 7 - assert "(exit 7)" in result.stdout diff --git a/tests/model_executor/test_jit_warmup_cutedsl_launcher.py b/tests/model_executor/test_jit_warmup_cutedsl_launcher.py deleted file mode 100644 index 76ec4d1200..0000000000 --- a/tests/model_executor/test_jit_warmup_cutedsl_launcher.py +++ /dev/null @@ -1,111 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from dataclasses import dataclass -from typing import Any - -from vllm.model_executor.warmup import jit_warmup_cutedsl_helper -from vllm.model_executor.warmup.jit_warmup import kernel_launcher -from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( - CuTeDSLLaunchSpec, - VllmCuTeDSLJitKernel, -) - - -class _TestCuTeDSLKernel(VllmCuTeDSLJitKernel["_TestCuTeDSLKernel.CompileKey"]): - @dataclass(frozen=True) - class CompileKey: - variant: int - - @staticmethod - def kernel(compile_key: CompileKey) -> tuple[str, int]: - return "entry", compile_key.variant - - def dispatch(self, *, variant: int) -> CompileKey: - return self.CompileKey(variant=variant) - - def get_warmup_keys(self) -> list[CompileKey]: - return self._trace_dispatch(self.dispatch)(variant=(1, 2)) - - def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: - return (f"fake-{compile_key.variant}",) - - @kernel_launcher - def __call__( - self, - payload: str, - *, - variant: int, - ) -> CuTeDSLLaunchSpec[CompileKey]: - compile_key = self.dispatch(variant=variant) - return compile_key, (payload,) - - -def test_cutedsl_launcher_reuses_compiled_executor(monkeypatch) -> None: - compile_calls: list[tuple[Any, tuple[Any, ...]]] = [] - launch_calls: list[tuple[int, tuple[Any, ...]]] = [] - - def compile_cutedsl(entry: tuple[str, int], *args: Any) -> Any: - compile_calls.append((entry, args)) - variant = entry[1] - - def executor(*runtime_args: Any) -> tuple[int, tuple[Any, ...]]: - launch_calls.append((variant, runtime_args)) - return variant, runtime_args - - return executor - - monkeypatch.setattr( - jit_warmup_cutedsl_helper, - "compile_cutedsl", - compile_cutedsl, - ) - kernel = _TestCuTeDSLKernel() - compile_key = kernel.CompileKey(variant=1) - - kernel.compile(compile_key) - kernel.compile(compile_key) - assert compile_calls == [(("entry", 1), ("fake-1",))] - - assert kernel("runtime", variant=1) == (1, ("runtime",)) - assert kernel("other", variant=2) == (2, ("other",)) - assert compile_calls == [ - (("entry", 1), ("fake-1",)), - (("entry", 2), ("fake-2",)), - ] - assert launch_calls == [(1, ("runtime",)), (2, ("other",))] - - -def test_cutedsl_launcher_runs_post_launch_finalizer(monkeypatch) -> None: - events: list[str] = [] - - def compile_cutedsl(entry: tuple[str, int], *args: Any) -> Any: - def executor(*runtime_args: Any) -> None: - events.append("launch") - - return executor - - class FinalizedKernel(_TestCuTeDSLKernel): - @kernel_launcher - def __call__( - self, - payload: str, - *, - variant: int, - ) -> CuTeDSLLaunchSpec[_TestCuTeDSLKernel.CompileKey]: - compile_key = self.dispatch(variant=variant) - - def finalize() -> str: - events.append("finalize") - return "output" - - return compile_key, (payload,), "unused", finalize - - monkeypatch.setattr( - jit_warmup_cutedsl_helper, - "compile_cutedsl", - compile_cutedsl, - ) - - assert FinalizedKernel()("runtime", variant=1) == "output" - assert events == ["launch", "finalize"] diff --git a/tests/models/transformers/test_layer_registry_e2e.py b/tests/models/transformers/test_layer_registry_e2e.py deleted file mode 100644 index 13d4b80a16..0000000000 --- a/tests/models/transformers/test_layer_registry_e2e.py +++ /dev/null @@ -1,113 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""End-to-end check of the Transformers backend's hw-agnostic layers. - -Serves a tiny Llama through `vllm_runner` with and without -`VLLM_USE_HW_AGNOSTIC` and checks the layer providers and logprobs. Needs a -real engine, so it runs in the GPU job. The CPU-only unit tests for the -resolution logic (`layers._resolve`) are in `test_layer_registry.py`. -""" - -import pytest -import torch - - -@pytest.fixture(scope="module") -def tiny_llama_path(tmp_path_factory): - """A randomly-initialized microscopic Llama saved to disk (with an ungated - tokenizer) so vLLM can load it like any local checkpoint.""" - from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM - - tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer") - config = LlamaConfig( - vocab_size=tokenizer.vocab_size, - hidden_size=128, - intermediate_size=256, - num_hidden_layers=2, - num_attention_heads=4, - rms_norm_eps=1e-6, - hidden_act="silu", - ) - torch.manual_seed(0) - model = LlamaForCausalLM(config) - - path = tmp_path_factory.mktemp("tiny_llama") - model.save_pretrained(path) - tokenizer.save_pretrained(path) - return str(path) - - -# Registered names of the layers the backend can -# currently route to hw-agnostic implementations. -_COVERED_LAYERS = ("rms_norm", "silu_and_mul") - - -def _layer_providers(model) -> dict[str, str]: - """Map each covered layer type present in the model to the provider its - implementation came from (``hw_agnostic`` or ``vllm``). - """ - - def provider_of(module) -> str | None: - for cls in type(module).__mro__: - if "hw_agnostic.layers" in cls.__module__: - return "hw_agnostic" - if ".model_executor.layers." in cls.__module__: - return "vllm" - return None - - providers: dict[str, str] = {} - for module in model.modules(): - name = getattr(module, "name", None) - if ( - name in _COVERED_LAYERS - and name not in providers - and (prov := provider_of(module)) is not None - ): - providers[name] = prov - return providers - - -def _serve(vllm_runner, model_path, prompts): - """Serve the model through the backend; return (layer_providers, logprobs).""" - with vllm_runner( - model_path, - model_impl="transformers", - max_model_len=64, - enforce_eager=True, - gpu_memory_utilization=0.3, - ) as runner: - assert runner.llm.llm_engine.model_config.using_transformers_backend() - providers = runner.apply_model(_layer_providers)[0] - outputs = runner.generate_greedy_logprobs( - prompts, max_tokens=32, num_logprobs=5 - ) - return providers, outputs - - -def test_hw_agnostic_matches_vllm_end_to_end(monkeypatch, vllm_runner, tiny_llama_path): - """Serving the tiny model with hw-agnostic layers matches the vLLM baseline.""" - # spawn: `layers` resolves its symbols at import time, so the worker must - # re-import it with VLLM_USE_HW_AGNOSTIC already set (fork would inherit - # the parent's already-resolved module). - monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") - # apply_model pickles the introspection function. - monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") - from ..utils import check_logprobs_close - - prompts = ["The capital of France is", "vLLM is"] - - monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "0") - vllm_providers, vllm_outputs = _serve(vllm_runner, tiny_llama_path, prompts) - # Both replaceable layers present in a Llama block must be vLLM's here. - assert vllm_providers == {"rms_norm": "vllm", "silu_and_mul": "vllm"} - - monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1") - hw_providers, hw_outputs = _serve(vllm_runner, tiny_llama_path, prompts) - assert hw_providers == {"rms_norm": "hw_agnostic", "silu_and_mul": "hw_agnostic"} - - check_logprobs_close( - outputs_0_lst=vllm_outputs, - outputs_1_lst=hw_outputs, - name_0="vllm", - name_1="hw_agnostic", - ) diff --git a/tests/v1/ec_connector/__init__.py b/tests/v1/ec_connector/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/v1/ec_connector/integration/README.md b/tests/v1/ec_connector/integration/README.md deleted file mode 100644 index a89bbfd347..0000000000 --- a/tests/v1/ec_connector/integration/README.md +++ /dev/null @@ -1,233 +0,0 @@ -# EPD Correctness Test - -This test verifies that EPD (Encoder-Prefill-Decode) disaggregation produces identical outputs to a baseline single instance. - -## What It Tests - -- **Baseline**: Single vLLM instance serving a multimodal model -- **EPD (1E+1PD)**: 1 Encoder + 1 Prefill-Decode instance -- **Baseline (1P+1D)**: 1 Prefill + 1 Decode instance -- **EPD (1E+1P+1D)**: 1 Encoder + 1 Prefill + 1 Decode instance - -The test ensures that disaggregated encoding produces **identical** outputs to the baseline. - -Note that currently PD disaggregation set up may give slightly different results from a single instance. Therefore, we need the result from 1P+1D as the baseline for 1E+1P+1D - -Please refer to [Disaggregated Encoder Feature](../../../../docs/features/disagg_encoder.md) for the detailed explanation for the EPD features. - -## Files - -- `run_epd_correctness_test.sh` - Main test script (starts all instances and runs tests) -- `test_epd_correctness.py` - Python test script (compares outputs) - -## Usage - -### Multimodal Prompts (Default) - -```bash -cd vllm -./tests/v1/ec_connector/integration/run_epd_correctness_test.sh -``` - -This runs the test with actual multimodal (image) prompts. - -### Text-Only Prompts - -```bash -cd vllm -USE_MM_PROMPTS=0 ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh -``` - -This runs a quick test with text-only prompts to verify the setup works. - -### Custom Configuration - -```bash -# Use specific GPUs -GPU_E=0 GPU_PD=1 GPU_P=1 GPU_D=2 bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh - -# Use specific ports -ENDPOINT_PORT=10001 bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh - -# Use specific model -MODEL="Qwen/Qwen2.5-VL-3B-Instruct" bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh - -# Use specific storage path -EC_SHARED_STORAGE_PATH="/tmp/my_ec_cache" bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh -``` - -## How It Works - -### NIXL EC failure isolation (1E + 1PD) - -```bash -PATH="$PWD/.venv/bin:$PATH" CUDA_VISIBLE_DEVICES=0,1 \ - .venv/bin/python -m pytest \ - tests/v1/ec_connector/integration/test_nixl_failure.py -v -s -``` - -Requires two CUDA GPUs and NIXL. `MODEL` overrides the default -`Qwen/Qwen2.5-VL-3B-Instruct`. The test starts real E/PD servers with CUDA -graphs enabled and uses the proxy's rewrite helper. Only the failing request's -control endpoint is replaced: a ZMQ peer returns `NACK_MISSING` and verifies -that the consumer actually requested the encoding. - -Metadata-only requests must return a request-level error; requests retaining -the image must succeed through local fallback. After each case, a fresh -metadata-only image request must succeed through the real NIXL transfer path. -No server restart, proxy retry, cache hit, or timeout race can mask the result. -This is a correctness test, not a throughput benchmark. - -### Mooncake EC (1E + 1PD) - -```bash -PYTHON_BIN="$PWD/.venv/bin/python" \ - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh -``` - -Requires two GPUs and Mooncake TransferEngine. TCP is the default transport; -no RDMA-capable network hardware is required. -The script sets `MC_FORCE_TCP=1` in TCP mode so Mooncake cannot auto-select RDMA. -The baseline runs first on GPU 0, followed by E on GPU 0 and PD on GPU 1. -`MOONCAKE_EC_PROTOCOL=rdma` selects RDMA instead; host-specific transport -environment variables should be set by the caller. - -Three black-box cases check fixed short answers and compare with the baseline: - -- One image: read the STOP sign. -- Two different images (including a local file): identify flowers and birds. -- The same image twice: read both STOP signs. - -By default, all three requests run concurrently for two rounds, exercising -shared hashes across requests and reuse after completion. Every response is -compared, not just the final round. Set `CONCURRENCY` and `REPEAT` to override. -Prefix caching is disabled and CUDA graphs remain enabled. This is a small -correctness suite, not a performance benchmark or failure-injection suite. - -`LOG_PATH` and `BASELINE_FILE` select the log directory and reference output. -`SKIP_BASELINE=1` reuses a reference generated with the same model and test -configuration. `USE_MM_PROMPTS=0` only checks text routing, not EC transfer. - -Buildkite runs this script as `mooncake-ec-tcp-e2e-2-gpus` on two L4 GPUs. -The job is defined in `.buildkite/test_areas/disaggregated_mooncake.yaml` and -selected for changes to EC, multimodal processing, scheduler/model-runner -integration, the proxy, or these tests (subject to normal PR CI approval). -It reuses the CI image's Python packages through a system-site-packages venv -and installs the CUDA-compatible Mooncake wheel. It uses loopback networking -and TCP only, without RDMA devices or peer-memory setup. - -The job fails on startup errors, request errors, or answer mismatches. On -failure, the script prints the last 100 lines of each server/proxy log to -the CI job log; full files remain under `LOG_PATH` while the container exists. - -### Step 1: Baseline - -1. Start single vLLM instance on GPU -2. Run test prompts (multimodal or text-only) -3. Save outputs to `.vllm_epd_baseline.txt` -4. Shutdown instance - -### Step 2: EPD (1E + 1PD) - -1. Clear encoder cache storage -2. Start instances and proxy -3. Run same test prompts -4. Assert outputs match baseline exactly -5. Shutdown instances - -### Step 3: EPD (1E + 1P + 1D) - -1. Clear encoder cache storage -2. Start instances and proxy -3. Run same test prompts -4. Assert outputs match baseline exactly -5. Shutdown instances - -## Test Scenarios - -### Multimodal Prompts (--use_mm_prompts) - -Tests encoder cache transfer: - -- Single image query -- Multiple images in one request -- Mixed image and text -- Image with detailed questions - -### Text-Only Prompts (default) - -Quick sanity check: - -- Simple text queries -- Text-only explanations -- Verifies proxy routing works - -## Expected Behavior - -### ✅ Test Passes When - -- All disagg outputs match baseline outputs exactly -- No errors during instance startup -- Encoder cache is properly saved and loaded -- Proxy correctly routes requests - -### ❌ Test Fails When - -- Outputs differ between baseline and disagg -- Server startup fails -- Encoder cache not found (should fall back to local execution) -- Proxy routing errors - -## Notes - -- The test uses deterministic generation (`temperature=0.0`, `seed=42`) -- Encoder cache should enable exact output reproduction -- Test cleans up all instances and cache files after completion -- Safe to run multiple times (idempotent) -- We setup the PD disagg part with NixlConnector. Please read details about EPD in `examples/disaggregated/disaggregated_encoder/README.md` - -## Requirements - -- Multiple GPUs (3 for 1E+1P+1D, 2 for 1E+1PD, 1 for baseline) - - 1E+1P+1D is runnable with 2 GPU by assign E and P on the same GPU now. -- Multimodal model (e.g., Qwen2.5-VL-3B-Instruct) -- Internet access (for accessing vllm test images) - -## Debugging - -### Check Logs - -Logs and baseline output are saved in `/tmp/` by default. -Can be customized by changing the environment variables. - -### Check Encoder Cache - -```bash -# Verify cache files are created -ls -la $EC_SHARED_STORAGE_PATH/ - -# Should see directories with mm_hash names -# Each containing encoder_cache.safetensors -``` - -### Manual Testing - -Run individual components: - -```bash -# Baseline only -python test_epd_correctness.py \ - --service_url http://localhost:8000 \ - --model_name Qwen/Qwen2.5-VL-3B-Instruct \ - --mode baseline \ - --baseline_file test_output.txt \ - --use_mm_prompts - -# Disagg only (requires baseline output file!) -python test_epd_correctness.py \ - --service_url http://localhost:8000 \ - --model_name Qwen/Qwen2.5-VL-3B-Instruct \ - --mode disagg \ - --baseline_file test_output.txt \ - --use_mm_prompts -``` diff --git a/tests/v1/ec_connector/integration/hato.jpg b/tests/v1/ec_connector/integration/hato.jpg deleted file mode 100644 index 9c7e390e7d..0000000000 Binary files a/tests/v1/ec_connector/integration/hato.jpg and /dev/null differ diff --git a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh b/tests/v1/ec_connector/integration/run_epd_correctness_test.sh deleted file mode 100644 index b48949f922..0000000000 --- a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh +++ /dev/null @@ -1,505 +0,0 @@ -#!/bin/bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# -# EPD (Encoder-Prefill-Decode) Correctness Test -# -# This script tests that EPD disaggregation produces the same outputs as baseline. -# It runs: -# 1. Baseline: Single vLLM instance -# 2. EPD: 1E + 1PD setup -# 3. Baseline for (E + P + D): 1P + 1D vLLM instances disagg -# 4. EPD: 1E + 1P + 1D setup - -# For GPU usage - -# set -xe - -# Resolve the repository root from the script location instead of `.git`. -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" -GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" - -# Model to test -MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}" -MAX_MODEL_LEN="${MAX_MODEL_LEN:-10240}" -GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.7}" -MAX_NUM_SEQS="${MAX_NUM_SEQS:-128}" - -# Set 1 to use multimodal prompts; else to use text-only -USE_MM_PROMPTS="${USE_MM_PROMPTS:-1}" -USE_TWO_IMAGE_PROMPT="${USE_TWO_IMAGE_PROMPT:-1}" -TEST_FLAGS=() -if [[ "$USE_MM_PROMPTS" == "1" ]]; then - TEST_FLAGS+=(--use_mm_prompts) -fi -if [[ "$USE_TWO_IMAGE_PROMPT" != "1" ]]; then - TEST_FLAGS+=(--skip_two_image_prompt) -fi - -# GPU configuration -GPU_E="${GPU_E:-0}" -GPU_P="${GPU_P:-1}" -GPU_D="${GPU_D:-2}" -GPU_SINGLE="${GPU_SINGLE:-$GPU_P}" -GPU_PD="${GPU_PD:-$GPU_P}" - -# Device platform and affinity environment variable -DEVICE_PLATFORM="${DEVICE_PLATFORM:-cuda}" -if [[ -z "${DEVICE_AFFINITY_ENV:-}" ]]; then - if [[ "${DEVICE_PLATFORM,,}" == "xpu" ]]; then - DEVICE_AFFINITY_ENV="ZE_AFFINITY_MASK" - else - DEVICE_AFFINITY_ENV="CUDA_VISIBLE_DEVICES" - fi -fi - -# Port -ENCODE_PORT="${ENCODE_PORT:-19534}" -PREFILL_PORT="${PREFILL_PORT:-19535}" -DECODE_PORT="${DECODE_PORT:-19536}" -PREFILL_DECODE_PORT="${PREFILL_DECODE_PORT:-19537}" -ENDPOINT_PORT="${ENDPOINT_PORT:-10001}" - -# Storage path for encoder cache -EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/tmp/ec_cache_test}" -TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-600}" - -# Output file for baseline comparison and logs -LOG_PATH="${LOG_PATH:-/tmp}" -BASELINE_FILE="${BASELINE_FILE:-/tmp/vllm_baseline.txt}" -BASELINE_PD_FILE="${BASELINE_PD_FILE:-/tmp/vllm_epd_baseline.txt}" - -mkdir -p "$LOG_PATH" - -# Trap the SIGINT signal (triggered by Ctrl+C) -trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT - -# Wait for server to be ready -wait_for_server() { - local port=$1 - timeout "$TIMEOUT_SECONDS" bash -c " - until curl -s localhost:${port}/v1/chat/completions > /dev/null; do - sleep 1 - done" && return 0 || return 1 -} - -# Cleanup function -cleanup_instances() { - echo "Cleaning up any running vLLM instances..." - pkill -f "vllm serve" || true - pkill -f "disagg_epd_proxy.py" || true - sleep 2 -} - -# Function to run baseline (single instance) -run_baseline() { - echo "================================" - echo "Running BASELINE (single instance)" - echo "================================" - - cleanup_instances - rm -rf "$EC_SHARED_STORAGE_PATH" - - local PORT=$ENDPOINT_PORT - - # Start baseline instance - echo "Starting baseline instance on GPU $GPU_SINGLE, port $PORT" - env "$DEVICE_AFFINITY_ENV=$GPU_SINGLE" vllm serve "$MODEL" \ - --port "$PORT" \ - --max-model-len "$MAX_MODEL_LEN" \ - --enforce-eager \ - --gpu-memory-utilization 0.9 \ - --max-num-seqs "$MAX_NUM_SEQS" \ - --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ - > "$LOG_PATH"/baseline.log 2>&1 & - - local BASELINE_PID=$! - - # Wait for baseline to start - echo "Waiting for baseline instance to start..." - wait_for_server "$PORT" - - curl http://127.0.0.1:"$PORT"/v1/models - echo "" - - # Run test in baseline mode - echo "Running baseline..." - - python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \ - --service_url "http://localhost:$PORT" \ - --model_name "$MODEL" \ - --mode baseline \ - --baseline_file "$BASELINE_FILE" \ - "${TEST_FLAGS[@]}" - - # Cleanup baseline - echo "Stopping baseline instance..." - kill $BASELINE_PID 2>/dev/null || true - sleep 2 - cleanup_instances -} - -# Function to run EPD with 1E + 1PD -run_epd_1e_1pd() { - echo "================================" - echo "Running EPD (1E + 1PD)" - echo "================================" - - cleanup_instances - rm -rf "$EC_SHARED_STORAGE_PATH" - mkdir -p "$EC_SHARED_STORAGE_PATH" - - local ENCODE_PORT=$ENCODE_PORT - local PREFILL_DECODE_PORT=$PREFILL_DECODE_PORT - local PROXY_PORT=$ENDPOINT_PORT - - declare -a PIDS=() - - # Start encoder instance - echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT" - env "$DEVICE_AFFINITY_ENV=$GPU_E" vllm serve "$MODEL" \ - --port "$ENCODE_PORT" \ - --max-model-len "$MAX_MODEL_LEN" \ - --enforce-eager \ - --gpu-memory-utilization 0.01 \ - --enable-request-id-headers \ - --no-enable-prefix-caching \ - --max-num-batched-tokens 114688 \ - --max-num-seqs "$MAX_NUM_SEQS" \ - --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ - --ec-transfer-config '{ - "ec_connector": "ECExampleConnector", - "ec_role": "ec_producer", - "ec_connector_extra_config": { - "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'" - } - }' \ - > "$LOG_PATH"/1e1pd_encoder.log 2>&1 & - PIDS+=($!) - - # Start prefill+decode instance - echo "Starting PD instance on GPU $GPU_PD, port $PREFILL_DECODE_PORT" - env "$DEVICE_AFFINITY_ENV=$GPU_PD" vllm serve "$MODEL" \ - --port "$PREFILL_DECODE_PORT" \ - --max-model-len "$MAX_MODEL_LEN" \ - --enforce-eager \ - --enable-mm-embeds \ - --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ - --enable-request-id-headers \ - --max-num-seqs "$MAX_NUM_SEQS" \ - --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ - --ec-transfer-config '{ - "ec_connector": "ECExampleConnector", - "ec_role": "ec_consumer", - "ec_connector_extra_config": { - "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'" - } - }' \ - > "$LOG_PATH"/1e1pd_pd.log 2>&1 & - PIDS+=($!) - - # Wait for instances to start - echo "Waiting for encoder instance..." - wait_for_server "$ENCODE_PORT" - echo "Waiting for PD instance..." - wait_for_server "$PREFILL_DECODE_PORT" - - # Start proxy - echo "Starting EPD proxy on port $PROXY_PORT" - python "${GIT_ROOT}/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py" \ - --host "0.0.0.0" \ - --port "$PROXY_PORT" \ - --encode-servers-urls "http://localhost:$ENCODE_PORT" \ - --prefill-servers-urls "disable" \ - --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \ - > "$LOG_PATH"/1e1pd_proxy.log 2>&1 & - PIDS+=($!) - - # Wait for proxy - echo "Waiting for proxy..." - wait_for_server "$PROXY_PORT" - - curl http://127.0.0.1:"$PROXY_PORT"/v1/models - curl http://127.0.0.1:"$PROXY_PORT"/health - echo "" - - echo "All EPD (1E+1PD) services are up!" - - # Run test in disagg mode - echo "Running EPD (1E+1PD) correctness test..." - - python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \ - --service_url "http://localhost:$PROXY_PORT" \ - --model_name "$MODEL" \ - --mode disagg \ - --baseline_file "$BASELINE_FILE" \ - "${TEST_FLAGS[@]}" - - # Cleanup - echo "✓✓ 1E+1PD Correctness Test finished" - echo "Stopping EPD (1E+1PD) instances..." - for pid in "${PIDS[@]}"; do - kill "$pid" 2>/dev/null || true - done - sleep 2 - cleanup_instances -} - -# Function to run baseline for 1E + 1P + 1D (PD disagg) -run_baseline_1p_1d() { - echo "================================" - echo "Running PD BASELINE (1P + 1D)" - echo "================================" - - cleanup_instances - rm -rf "$EC_SHARED_STORAGE_PATH" - mkdir -p "$EC_SHARED_STORAGE_PATH" - - local PREFILL_PORT=$PREFILL_PORT - local DECODE_PORT=$DECODE_PORT - local PROXY_PORT=$ENDPOINT_PORT - - declare -a PIDS=() - - # Start prefill instance - echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT" - env "$DEVICE_AFFINITY_ENV=$GPU_P" \ - VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \ - vllm serve "$MODEL" \ - --port "$PREFILL_PORT" \ - --max-model-len "$MAX_MODEL_LEN" \ - --enforce-eager \ - --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ - --enable-request-id-headers \ - --max-num-seqs "$MAX_NUM_SEQS" \ - --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ - --kv-transfer-config '{ - "kv_connector": "NixlConnector", - "kv_role": "kv_producer" - }' \ - > "$LOG_PATH"/1p1d_prefill.log 2>&1 & - PIDS+=($!) - - # Start decode instance - echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT" - env "$DEVICE_AFFINITY_ENV=$GPU_D" \ - VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \ - vllm serve "$MODEL" \ - --port "$DECODE_PORT" \ - --max-model-len "$MAX_MODEL_LEN" \ - --enforce-eager \ - --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ - --enable-request-id-headers \ - --max-num-seqs "$MAX_NUM_SEQS" \ - --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ - --kv-transfer-config '{ - "kv_connector": "NixlConnector", - "kv_role": "kv_consumer" - }' \ - > "$LOG_PATH"/1p1d_decode.log 2>&1 & - PIDS+=($!) - - # Wait for instances to start - echo "Waiting for prefill instance..." - wait_for_server "$PREFILL_PORT" - echo "Waiting for decode instance..." - wait_for_server "$DECODE_PORT" - - # Start proxy - echo "Starting EPD proxy on port $PROXY_PORT" - python "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py" \ - --host "0.0.0.0" \ - --port "$PROXY_PORT" \ - --prefiller-ports "$PREFILL_PORT" \ - --decoder-ports "$DECODE_PORT" \ - > "$LOG_PATH"/1p1d_proxy.log 2>&1 & - PIDS+=($!) - - # Wait for proxy - echo "Waiting for proxy..." - wait_for_server "$PROXY_PORT" - - curl http://127.0.0.1:"$PROXY_PORT"/healthcheck - echo "" - - echo "All PD (1P+1D) services are up!" - - # Run test in baseline mode - echo "Running PD disagg baseline..." - - python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \ - --service_url "http://localhost:$PROXY_PORT" \ - --model_name "$MODEL" \ - --mode baseline_pd \ - --baseline_file "$BASELINE_PD_FILE" \ - "${TEST_FLAGS[@]}" - - # Cleanup - echo "Stopping PD (1P+1D) instances..." - for pid in "${PIDS[@]}"; do - kill "$pid" 2>/dev/null || true - done - sleep 2 - cleanup_instances -} - -# Function to run EPD with 1E + 1P + 1D -run_epd_1e_1p_1d() { - echo "================================" - echo "Running EPD (1E + 1P + 1D)" - echo "================================" - - cleanup_instances - rm -rf "$EC_SHARED_STORAGE_PATH" - mkdir -p "$EC_SHARED_STORAGE_PATH" - - local ENCODE_PORT=$ENCODE_PORT - local PREFILL_PORT=$PREFILL_PORT - local DECODE_PORT=$DECODE_PORT - local PROXY_PORT=$ENDPOINT_PORT - - declare -a PIDS=() - - # Start encoder instance - echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT" - env "$DEVICE_AFFINITY_ENV=$GPU_E" vllm serve "$MODEL" \ - --port "$ENCODE_PORT" \ - --max-model-len "$MAX_MODEL_LEN" \ - --enforce-eager \ - --gpu-memory-utilization 0.01 \ - --enable-request-id-headers \ - --no-enable-prefix-caching \ - --max-num-batched-tokens 114688 \ - --max-num-seqs "$MAX_NUM_SEQS" \ - --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ - --ec-transfer-config '{ - "ec_connector": "ECExampleConnector", - "ec_role": "ec_producer", - "ec_connector_extra_config": { - "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'" - } - }' \ - > "$LOG_PATH"/1e1p1d_encoder.log 2>&1 & - PIDS+=($!) - - # Start prefill instance - echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT" - env "$DEVICE_AFFINITY_ENV=$GPU_P" \ - VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \ - vllm serve "$MODEL" \ - --port "$PREFILL_PORT" \ - --max-model-len "$MAX_MODEL_LEN" \ - --enforce-eager \ - --enable-mm-embeds \ - --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ - --enable-request-id-headers \ - --max-num-seqs "$MAX_NUM_SEQS" \ - --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ - --ec-transfer-config '{ - "ec_connector": "ECExampleConnector", - "ec_role": "ec_consumer", - "ec_connector_extra_config": { - "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'" - } - }' \ - --kv-transfer-config '{ - "kv_connector": "NixlConnector", - "kv_role": "kv_producer" - }' \ - > "$LOG_PATH"/1e1p1d_prefill.log 2>&1 & - PIDS+=($!) - - # Start decode instance - echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT" - env "$DEVICE_AFFINITY_ENV=$GPU_D" \ - VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \ - vllm serve "$MODEL" \ - --port "$DECODE_PORT" \ - --max-model-len "$MAX_MODEL_LEN" \ - --enforce-eager \ - --enable-mm-embeds \ - --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ - --enable-request-id-headers \ - --max-num-seqs "$MAX_NUM_SEQS" \ - --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ - --kv-transfer-config '{ - "kv_connector": "NixlConnector", - "kv_role": "kv_consumer" - }' \ - > "$LOG_PATH"/1e1p1d_decode.log 2>&1 & - PIDS+=($!) - - # Wait for instances to start - echo "Waiting for encoder instance..." - wait_for_server "$ENCODE_PORT" - echo "Waiting for prefill instance..." - wait_for_server "$PREFILL_PORT" - echo "Waiting for decode instance..." - wait_for_server "$DECODE_PORT" - - # Start proxy - echo "Starting EPD proxy on port $PROXY_PORT" - python "${GIT_ROOT}/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py" \ - --host "0.0.0.0" \ - --port "$PROXY_PORT" \ - --encode-servers-urls "http://localhost:$ENCODE_PORT" \ - --prefill-servers-urls "http://localhost:$PREFILL_PORT" \ - --decode-servers-urls "http://localhost:$DECODE_PORT" \ - > "$LOG_PATH"/1e1p1d_proxy.log 2>&1 & - PIDS+=($!) - - # Wait for proxy - echo "Waiting for proxy..." - wait_for_server "$PROXY_PORT" - - curl http://127.0.0.1:"$PROXY_PORT"/v1/models - curl http://127.0.0.1:"$PROXY_PORT"/health - echo "" - - echo "All EPD (1E+1P+1D) services are up!" - - # Run test in disagg mode - echo "Running EPD (1E+1P+1D) correctness test..." - - python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \ - --service_url "http://localhost:$PROXY_PORT" \ - --model_name "$MODEL" \ - --mode disagg \ - --baseline_file "$BASELINE_PD_FILE" \ - "${TEST_FLAGS[@]}" - - # Cleanup - echo "✓✓ 1E+1P+1D Correctness Test finished" - echo "Stopping EPD (1E+1P+1D) instances..." - for pid in "${PIDS[@]}"; do - kill "$pid" 2>/dev/null || true - done - sleep 2 - cleanup_instances -} - -# Main execution -echo "================================" -echo "EPD Correctness Test Suite" -echo "Model: $MODEL" -echo "================================" - -# Step 1: Run baseline -run_baseline - -# Step 2: Test 1E + 1PD -run_epd_1e_1pd - -# Step 3: Test baseline 1P + 1D -run_baseline_1p_1d - -# # Step 4: Test 1E + 1P + 1D -run_epd_1e_1p_1d - -# Cleanup output file -rm -f "$BASELINE_FILE" -rm -f "$BASELINE_PD_FILE" - -echo "================================" -echo "✓✓ All EPD correctness tests finished!" -echo "================================" diff --git a/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh b/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh deleted file mode 100755 index c4e9b0d069..0000000000 --- a/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh +++ /dev/null @@ -1,323 +0,0 @@ -#!/usr/bin/env bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# -# Full-stack EPD validation with ECMooncakeConnector (Mooncake TransferEngine): -# 1) Single-GPU baseline (multimodal) -> saves baseline JSON -# 2) 1 Encoder + 1 PD with Mooncake EC -> compare outputs to baseline -# -# Usage (from repo root): -# ./tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh -# -# Env: -# MODEL HF model id (default: Qwen/Qwen2.5-VL-3B-Instruct) -# GPU_SINGLE / GPU_E / GPU_PD GPU ids (defaults 0 / 0 / 1) -# ENDPOINT_PORT, ENCODE_PORT, PREFILL_DECODE_PORT -# EC_MOONCAKE_RESERVATION_HOST consumer control host (default 127.0.0.1) -# MOONCAKE_EC_PROTOCOL tcp | rdma (default tcp) -# USE_MM_PROMPTS 1 (default) or 0 for text-only quick sanity -# TIMEOUT_SECONDS wait_for_server timeout (default 1200) -# SKIP_BASELINE set to 1 to reuse existing BASELINE_FILE -# CONCURRENCY / REPEAT concurrent requests / rounds (defaults 3 / 2) -# MAX_MODEL_LEN context length (default 16384) -# E_CUDAGRAPH_MM_ENCODER enable encoder CUDA graphs (default 0) -# E_MM_ENCODER_ATTN_BACKEND optional encoder attention backend override - -set -euo pipefail - -GIT_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd) -cd "$GIT_ROOT" || exit 1 -export PYTHONPATH="${GIT_ROOT}:${PYTHONPATH:-}" -PYTHON_BIN="${PYTHON_BIN:-${GIT_ROOT}/.venv/bin/python}" - -MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}" -USE_MM_PROMPTS="${USE_MM_PROMPTS:-1}" -TEST_ARGS=(--concurrency "${CONCURRENCY:-3}" --repeat "${REPEAT:-2}") -if [[ "$USE_MM_PROMPTS" == "1" ]]; then - TEST_ARGS+=(--use_mm_prompts --mm_smoke_test) -fi -MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}" -E_ATTN_ARGS=() -if [[ -n "${E_MM_ENCODER_ATTN_BACKEND:-}" ]]; then - E_ATTN_ARGS+=(--mm-encoder-attn-backend "$E_MM_ENCODER_ATTN_BACKEND") -fi -case "${E_CUDAGRAPH_MM_ENCODER:-0}" in - 0) E_COMPILATION_CONFIG='{"cudagraph_mm_encoder":false}' ;; - 1) E_COMPILATION_CONFIG='{"cudagraph_mm_encoder":true}' ;; - *) echo "E_CUDAGRAPH_MM_ENCODER must be 0 or 1"; exit 1 ;; -esac - -GPU_SINGLE="${GPU_SINGLE:-0}" -GPU_E="${GPU_E:-0}" -GPU_PD="${GPU_PD:-1}" - -ENCODE_PORT="${ENCODE_PORT:-19534}" -PREFILL_DECODE_PORT="${PREFILL_DECODE_PORT:-19537}" -ENDPOINT_PORT="${ENDPOINT_PORT:-10002}" -BASELINE_PORT="${BASELINE_PORT:-10003}" - -EC_MOONCAKE_RESERVATION_HOST="${EC_MOONCAKE_RESERVATION_HOST:-127.0.0.1}" -EC_MOONCAKE_RESERVATION_PORT="${EC_MOONCAKE_RESERVATION_PORT:-19019}" -MOONCAKE_EC_PROTOCOL="${MOONCAKE_EC_PROTOCOL:-tcp}" -export EC_MOONCAKE_RESERVATION_HOST -export EC_MOONCAKE_RESERVATION_PORT -export MOONCAKE_EC_PROTOCOL -if [[ "$MOONCAKE_EC_PROTOCOL" == "tcp" ]]; then - # TransferEngine may otherwise auto-select RDMA on hosts with an HCA. - export MC_FORCE_TCP=1 -else - unset MC_FORCE_TCP -fi - -LOG_PATH="${LOG_PATH:-/tmp}" -BASELINE_FILE="${BASELINE_FILE:-/tmp/vllm_epd_mooncake_baseline.txt}" -TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-1200}" -PIDS=() - -mkdir -p "$LOG_PATH" - -VLLM_SERVE=("$PYTHON_BIN" -m vllm.entrypoints.cli.main serve) - -ENC_EC_JSON=$("$PYTHON_BIN" </dev/null || return 1 - if curl --max-time 2 -fsS "http://localhost:${port}/health" >/dev/null 2>&1; then - return 0 - fi - sleep 2 - done - return 1 -} - -cleanup_instances() { - if ((${#PIDS[@]} == 0)); then - return - fi - echo "Cleaning up tracked vLLM / proxy processes..." - for pid in "${PIDS[@]}"; do - kill "$pid" 2>/dev/null || true - done - for _ in {1..10}; do - local alive=0 - for pid in "${PIDS[@]}"; do - if kill -0 "$pid" 2>/dev/null; then - alive=1 - fi - done - ((alive == 0)) && break - sleep 1 - done - for pid in "${PIDS[@]}"; do - if kill -0 "$pid" 2>/dev/null; then - kill -KILL "$pid" 2>/dev/null || true - fi - wait "$pid" 2>/dev/null || true - done - PIDS=() -} - -finish() { - local status=$? - cleanup_instances - if ((status != 0)); then - for log in "${LOG_PATH}"/mooncake_epd_*.log; do - [[ -f "$log" ]] || continue - echo "=== $log (last 100 lines) ===" - tail -100 "$log" - done - fi - exit "$status" -} - -trap finish EXIT -trap 'exit 130' INT -trap 'exit 143' TERM - -run_baseline() { - echo "================================" - echo "BASELINE (single vLLM, MM if enabled)" - echo "================================" - cleanup_instances - local PORT=$BASELINE_PORT - echo "Starting baseline on GPU $GPU_SINGLE port $PORT" - CUDA_VISIBLE_DEVICES="$GPU_SINGLE" "${VLLM_SERVE[@]}" "$MODEL" \ - --port "$PORT" \ - --gpu-memory-utilization 0.75 \ - --max-num-seqs 32 \ - --max-model-len "$MAX_MODEL_LEN" \ - --no-enable-prefix-caching \ - --limit-mm-per-prompt '{"image":2,"video":0}' \ - --allowed-local-media-path "${GIT_ROOT}/tests/v1/ec_connector/integration" \ - >"${LOG_PATH}/mooncake_epd_baseline.log" 2>&1 & - local BASELINE_PID=$! - PIDS+=("$BASELINE_PID") - echo "Waiting for baseline..." - wait_for_server "$PORT" "$BASELINE_PID" || { echo "Baseline failed to start"; return 1; } - curl -s "http://127.0.0.1:${PORT}/v1/models" | head -c 200 || true - echo "" - "$PYTHON_BIN" "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \ - --service_url "http://localhost:$PORT" \ - --model_name "$MODEL" \ - --mode baseline \ - --baseline_file "$BASELINE_FILE" \ - "${TEST_ARGS[@]}" - cleanup_instances -} - -run_epd_mooncake() { - echo "================================" - echo "EPD 1E + 1PD with ECMooncakeConnector" - echo "Reservation port on consumer: $EC_MOONCAKE_RESERVATION_PORT" - echo "Mooncake protocol: $MOONCAKE_EC_PROTOCOL" - echo "================================" - cleanup_instances - - echo "Starting ENCODER on GPU $GPU_E port $ENCODE_PORT" - VLLM_SERVER_DEV_MODE=1 CUDA_VISIBLE_DEVICES="$GPU_E" "${VLLM_SERVE[@]}" "$MODEL" \ - --port "$ENCODE_PORT" \ - --gpu-memory-utilization 0.35 \ - --mm-tensor-ipc torch_shm \ - --enable-request-id-headers \ - --no-enable-prefix-caching \ - --max-num-batched-tokens "$MAX_MODEL_LEN" \ - --max-num-seqs 32 \ - --max-model-len "$MAX_MODEL_LEN" \ - --limit-mm-per-prompt '{"image":2,"video":0}' \ - --allowed-local-media-path "${GIT_ROOT}/tests/v1/ec_connector/integration" \ - --ec-transfer-config "$ENC_EC_JSON" \ - --compilation-config "$E_COMPILATION_CONFIG" \ - "${E_ATTN_ARGS[@]}" \ - --worker-extension-cls tests.v1.ec_connector.integration.test_epd_correctness.EncoderGraphTestWorkerExtension \ - >"${LOG_PATH}/mooncake_epd_encoder.log" 2>&1 & - local ENCODER_PID=$! - PIDS+=("$ENCODER_PID") - - echo "Starting PD on GPU $GPU_PD port $PREFILL_DECODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_PD" "${VLLM_SERVE[@]}" "$MODEL" \ - --port "$PREFILL_DECODE_PORT" \ - --gpu-memory-utilization 0.75 \ - --mm-tensor-ipc torch_shm \ - --enable-mm-embeds \ - --enable-request-id-headers \ - --max-num-seqs 32 \ - --max-model-len "$MAX_MODEL_LEN" \ - --no-enable-prefix-caching \ - --limit-mm-per-prompt '{"image":2,"video":0}' \ - --allowed-local-media-path "${GIT_ROOT}/tests/v1/ec_connector/integration" \ - --ec-transfer-config "$PD_EC_JSON" \ - >"${LOG_PATH}/mooncake_epd_pd.log" 2>&1 & - local PD_PID=$! - PIDS+=("$PD_PID") - - echo "Waiting for encoder..." - wait_for_server "$ENCODE_PORT" "$ENCODER_PID" || { echo "Encoder failed to start"; return 1; } - echo "Waiting for PD..." - wait_for_server "$PREFILL_DECODE_PORT" "$PD_PID" || { echo "PD failed to start"; return 1; } - - echo "Starting EPD proxy on $ENDPOINT_PORT" - "$PYTHON_BIN" "${GIT_ROOT}/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py" \ - --host "0.0.0.0" \ - --port "$ENDPOINT_PORT" \ - --encode-servers-urls "http://localhost:$ENCODE_PORT" \ - --prefill-servers-urls "disable" \ - --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \ - --ec-consumer-zmq-addrs \ - "$EC_MOONCAKE_RESERVATION_ADDR" \ - >"${LOG_PATH}/mooncake_epd_proxy.log" 2>&1 & - local PROXY_PID=$! - PIDS+=("$PROXY_PID") - - echo "Waiting for proxy..." - wait_for_server "$ENDPOINT_PORT" "$PROXY_PID" || { echo "Proxy failed to start"; return 1; } - curl -s "http://127.0.0.1:${ENDPOINT_PORT}/health" || true - echo "" - - "$PYTHON_BIN" "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \ - --service_url "http://localhost:$ENDPOINT_PORT" \ - --model_name "$MODEL" \ - --mode disagg \ - --baseline_file "$BASELINE_FILE" \ - "${TEST_ARGS[@]}" - - curl -fsS "http://localhost:${ENCODE_PORT}/collective_rpc" \ - -H 'Content-Type: application/json' \ - -d '{"method":"encoder_graph_stats"}' \ - >"${LOG_PATH}/encoder_graph_stats.json" - "$PYTHON_BIN" - "${LOG_PATH}/encoder_graph_stats.json" "${E_CUDAGRAPH_MM_ENCODER:-0}" <<'PY' -import json -import sys - -with open(sys.argv[1]) as f: - results = json.load(f)["results"] -assert results, "No encoder workers returned graph statistics" -for stats in results: - if sys.argv[2] == "1": - assert stats["is_captured"], stats - assert stats["graph_hits"] > 0, stats - else: - assert stats["graph_hits"] == 0, stats -print("Encoder graph statistics:", results) -PY - - cleanup_instances -} - -echo "================================" -echo "1E + 1PD ECMooncake end-to-end correctness" -echo "MODEL=$MODEL" -echo "================================" - -if [[ "${SKIP_BASELINE:-0}" != "1" ]]; then - run_baseline -else - echo "SKIP_BASELINE=1 -> using existing $BASELINE_FILE" - [[ -f "$BASELINE_FILE" ]] || { echo "Missing baseline file"; exit 1; } -fi - -run_epd_mooncake - -echo "================================" -echo "PASS: Mooncake EC EPD matches baseline" -echo "Logs: ${LOG_PATH}/mooncake_epd_*.log" -echo "================================" diff --git a/tests/v1/ec_connector/integration/test_epd_correctness.py b/tests/v1/ec_connector/integration/test_epd_correctness.py deleted file mode 100644 index 4e37637ea4..0000000000 --- a/tests/v1/ec_connector/integration/test_epd_correctness.py +++ /dev/null @@ -1,395 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""EPD Correctness Test. - -Tests that EPD (Encoder-Prefill-Decode) disaggregation produces the same -outputs as a baseline single instance. - -Usage: - # Baseline mode (saves outputs): - python test_epd_correctness.py \ - --service_url http://localhost:8000 \ - --model_name Qwen/Qwen2.5-VL-3B-Instruct \ - --mode baseline \ - --baseline_file .vllm_epd_baseline.txt - - # Disagg mode (compares outputs): - python test_epd_correctness.py \ - --service_url http://localhost:8000 \ - --model_name Qwen/Qwen2.5-VL-3B-Instruct \ - --mode disagg \ - --baseline_file .vllm_epd_baseline.txt -""" - -import argparse -import json -import os -import time -from concurrent.futures import ThreadPoolExecutor -from typing import TYPE_CHECKING - -import openai -import requests - -from vllm.assets.image import ImageAsset -from vllm.multimodal.utils import encode_image_url - -if TYPE_CHECKING: - from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner - - -class EncoderGraphTestWorkerExtension: - """Expose encoder replay evidence through the existing test RPC endpoint.""" - - model_runner: "MMEncoderModelRunner" - - def encoder_graph_stats(self): - from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner - - runner = self.model_runner - assert isinstance(runner, MMEncoderModelRunner) - manager = runner.model_state.encoder_runner.cudagraph_manager - if manager is None: - return {"is_captured": False, "graph_hits": 0, "graph_misses": 0} - return {"is_captured": manager.is_captured(), **manager.get_cumulative_stats()} - - -MAX_OUTPUT_LEN = 256 - -# Sample prompts with multimodal content -image_1 = ImageAsset("stop_sign").pil_image.resize((1280, 720)) -image_2 = ImageAsset("cherry_blossom").pil_image.resize((1280, 720)) - -image_local_path = f"{os.path.dirname(os.path.abspath(__file__))}/hato.jpg" - -SAMPLE_PROMPTS_MM: list[dict] = [ - { - "messages": [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": {"url": encode_image_url(image_1)}, - }, - {"type": "text", "text": "What's in this image?"}, - ], - } - ], - "description": "Single image query", - }, - { - "messages": [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": {"url": encode_image_url(image_2)}, - }, - { - "type": "image_url", - "image_url": {"url": f"file://{image_local_path}"}, - }, - {"type": "text", "text": "Describe these 2 images in detail."}, - ], - } - ], - "description": "2 images with detailed query", - }, -] - -# Text-only prompts for mixed testing -SAMPLE_PROMPTS_TEXT: list[dict] = [ - { - "messages": [{"role": "user", "content": "What is the capital of France?"}], - "description": "Simple text-only query", - }, - { - "messages": [ - {"role": "user", "content": "Explain quantum computing in simple terms."} - ], - "description": "Text-only explanation request", - }, -] - - -def check_vllm_server(url: str, timeout=5, retries=10) -> bool: - """Check if the vLLM server is ready. - - Args: - url: The URL to check (usually /health or /healthcheck endpoint) - timeout: Timeout in seconds for each request - retries: Number of retries if the server is not ready - - Returns: - True if the server is ready, False otherwise - - """ - for attempt in range(retries): - try: - response = requests.get(url, timeout=timeout) - if response.status_code == 200: - print(f"Server is ready at {url}") - return True - else: - print( - f"Attempt {attempt + 1}/{retries}: Server returned " - f"status code {response.status_code}" - ) - except requests.exceptions.RequestException as e: - print(f"Attempt {attempt + 1}/{retries}: Error connecting: {e}") - time.sleep(2) # Wait before retrying - return False - - -def run_chat_completion( - base_url: str, - model_name: str, - messages: list, - max_tokens: int = MAX_OUTPUT_LEN, -) -> str: - """Run a chat completion request. - - Args: - base_url: Base URL of the vLLM server - model_name: Name of the model - messages: Messages for chat completion - max_tokens: Maximum tokens to generate - - Returns: - Generated text content - - """ - with openai.OpenAI( - api_key="EMPTY", base_url=base_url, timeout=120, max_retries=0 - ) as client: - completion = client.chat.completions.create( - model=model_name, - messages=messages, - max_tokens=max_tokens, - temperature=0.0, - seed=42, - ) - - content = completion.choices[0].message.content - assert content, "Expected a nonempty completion" - return content - - -def main(): - """Main test function.""" - parser = argparse.ArgumentParser( - description="EPD correctness test - compare disagg vs baseline" - ) - - parser.add_argument( - "--service_url", - type=str, - required=True, - help="The vLLM service URL (e.g., http://localhost:8000)", - ) - - parser.add_argument( - "--model_name", - type=str, - required=True, - help="Model name", - ) - - parser.add_argument( - "--mode", - type=str, - default="baseline", - choices=["baseline", "baseline_pd", "disagg"], - help="Mode: baseline/baseline_pd (saves outputs) or disagg (compares outputs)", - ) - - parser.add_argument( - "--baseline_file", - type=str, - default=".vllm_epd_baseline.txt", - help="File to save/load baseline outputs", - ) - - parser.add_argument( - "--use_mm_prompts", - action="store_true", - help="Use multimodal prompts (default: use text-only for quick testing)", - ) - - parser.add_argument( - "--skip_two_image_prompt", - action="store_true", - help="Skip the two-image multimodal prompt", - ) - - parser.add_argument( - "--mm_smoke_test", - action="store_true", - help="Use three short, fixed-answer image cases, including duplicate images", - ) - parser.add_argument("--concurrency", type=int, default=1) - parser.add_argument("--repeat", type=int, default=1) - - args = parser.parse_args() - if args.concurrency < 1 or args.repeat < 1: - parser.error("--concurrency and --repeat must be positive") - if args.mm_smoke_test and (not args.use_mm_prompts or args.skip_two_image_prompt): - parser.error("--mm_smoke_test requires all multimodal prompts") - - print(f"Service URL: {args.service_url}") - print(f"Model: {args.model_name}") - print(f"Mode: {args.mode}") - print(f"Output file: {args.baseline_file}") - print(f"Use MM prompts: {args.use_mm_prompts}") - - # Determine health check endpoint - if args.mode == "baseline": - health_check_url = f"{args.service_url}/health" - elif args.mode == "baseline_pd": - # Nixl toy proxy use /healthcheck - health_check_url = f"{args.service_url}/healthcheck" - else: - # Disagg EPD proxy uses /health - health_check_url = f"{args.service_url}/health" - if not os.path.exists(args.baseline_file): - raise ValueError( - f"In disagg mode, the output file {args.baseline_file} from " - "baseline does not exist. Run baseline mode first." - ) - - # Check if server is ready - if not check_vllm_server(health_check_url): - raise RuntimeError(f"vLLM server at {args.service_url} is not ready!") - - # Select prompts to use - if args.use_mm_prompts: - test_prompts = ( - SAMPLE_PROMPTS_MM[:1] if args.skip_two_image_prompt else SAMPLE_PROMPTS_MM - ) - print("Using multimodal prompts") - else: - test_prompts = SAMPLE_PROMPTS_TEXT - print("Using text-only prompts for quick testing") - - if args.mm_smoke_test: - image = SAMPLE_PROMPTS_MM[0]["messages"][0]["content"][0] - image_pair = SAMPLE_PROMPTS_MM[1]["messages"][0]["content"][:2] - cases = [ - ( - "Single image", - [image], - "What word is on the red road sign? Reply with only that word " - "in uppercase.", - "STOP", - ), - ( - "Two different images", - image_pair, - "Do these pictures contain both flowers and birds? " - "Reply with only YES or NO.", - "YES", - ), - ( - "Same image twice", - [image, image], - "Read the red road sign in each image, in image order. Reply " - "with only the two uppercase words separated by a comma and a space.", - "STOP, STOP", - ), - ] - test_prompts = [ - { - "description": description, - "expected": expected, - "messages": [ - { - "role": "user", - "content": [*images, {"type": "text", "text": question}], - } - ], - } - for description, images, question, expected in cases - ] - - # Run completions - service_url = f"{args.service_url}/v1" - output_strs = {} - - def complete(prompt_data): - output = run_chat_completion( - base_url=service_url, - model_name=args.model_name, - messages=prompt_data["messages"], - max_tokens=16 if args.mm_smoke_test else MAX_OUTPUT_LEN, - ) - if args.mm_smoke_test: - output = output.strip() - assert output == prompt_data["expected"], ( - f"{prompt_data['description']}: expected {prompt_data['expected']!r}, " - f"got {output!r}" - ) - return output - - # Each round includes concurrent requests sharing image hashes; later rounds - # exercise reuse after previous requests have finished. - with ThreadPoolExecutor(max_workers=args.concurrency) as executor: - for repeat in range(args.repeat): - outputs = executor.map(complete, test_prompts) - for prompt_data, output_str in zip(test_prompts, outputs): - key = prompt_data["description"] - if args.repeat > 1: - key = f"{key} (round {repeat + 1})" - output_strs[key] = output_str - print(f"{key}: {output_str}") - - if args.mode in ("baseline", "baseline_pd"): - # Baseline mode: Save outputs - print(f"\nSaving baseline outputs to {args.baseline_file}") - try: - with open(args.baseline_file, "w") as json_file: - json.dump(output_strs, json_file, indent=4) - print("✅ Baseline outputs saved successfully") - except OSError as e: - print(f"Error writing to file: {e}") - raise - else: - # Disagg mode: Load and compare outputs - print(f"\nLoading baseline outputs from {args.baseline_file}") - baseline_outputs = None - try: - with open(args.baseline_file) as json_file: - baseline_outputs = json.load(json_file) - except OSError as e: - print(f"Error reading from file: {e}") - raise - - # Verify outputs match - print("\nComparing disagg outputs with baseline...") - assert isinstance(baseline_outputs, dict), "Baseline outputs should be a dict" - assert len(baseline_outputs) == len(output_strs), ( - f"Length mismatch: baseline has {len(baseline_outputs)}, " - f"disagg has {len(output_strs)}" - ) - - all_match = True - for key, baseline_output in baseline_outputs.items(): - assert key in output_strs, f"{key} not in disagg outputs" - - disagg_output = output_strs[key] - if baseline_output == disagg_output: - print(f"✅ {key}: MATCH") - else: - print(f"❌ {key}: MISMATCH") - print(f" Baseline: {baseline_output}") - print(f" Disagg: {disagg_output}") - all_match = False - - assert all_match, "❌❌Disagg outputs do not match baseline!❌❌" - if all_match: - print("\n✅ All outputs match! Test PASSED") - - -if __name__ == "__main__": - main() diff --git a/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py b/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py deleted file mode 100644 index 8447532909..0000000000 --- a/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py +++ /dev/null @@ -1,389 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for EmbeddingCache. - -Pure integer bookkeeping — no mmap, torch, or CUDA — so these run anywhere. -""" - -import threading - -import pytest - -from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import ( - EmbeddingCache, -) - - -def _cache(num_blocks: int = 8) -> EmbeddingCache: - return EmbeddingCache(num_blocks) - - -# ── alloc ──────────────────────────────────────────────────────────────────── - - -def test_alloc_returns_unique_ascending_block_ids(): - cache = _cache() - entry = cache.alloc("a", 3) - assert entry is not None - assert len(entry.block_ids) == 3 - assert len(set(entry.block_ids)) == 3 - assert all(0 <= i < 8 for i in entry.block_ids) - # Ascending order lets consumers detect which blocks are adjacent. - assert list(entry.block_ids) == sorted(entry.block_ids) - - -def test_alloc_returns_tuple(): - cache = _cache() - entry = cache.alloc("a", 2) - assert isinstance(entry.block_ids, tuple) - - -def test_alloc_returns_none_when_full(): - cache = _cache(num_blocks=4) - cache.alloc("a", 4) - assert cache.alloc("b", 1) is None - - -def test_alloc_asserts_duplicate_key(): - cache = _cache() - cache.alloc("a", 2) - with pytest.raises(AssertionError, match="duplicate alloc"): - cache.alloc("a", 2) - - -def test_alloc_asserts_n_blocks_exceeds_capacity(): - cache = _cache(num_blocks=4) - with pytest.raises(AssertionError, match="capacity"): - cache.alloc("a", 5) - - -def test_alloc_zero_blocks(): - cache = _cache() - entry = cache.alloc("a", 0) - assert entry is not None - assert entry.block_ids == () - - -def test_alloc_starts_not_ready(): - cache = _cache() - entry = cache.alloc("a", 2) - assert not entry.ready - assert not entry.evictable - - -# ── get ────────────────────────────────────────────────────────────────────── - - -def test_get_missing_returns_none(): - cache = _cache() - assert cache.get("a") is None - - -def test_get_returns_entry(): - cache = _cache() - alloc_entry = cache.alloc("a", 2) - got = cache.get("a") - assert got is alloc_entry - - -def test_get_returns_not_ready_entry(): - cache = _cache() - cache.alloc("a", 2) - entry = cache.get("a") - assert entry is not None - assert not entry.ready - - -# ── mark_ready ─────────────────────────────────────────────────────────────── - - -def test_mark_ready_transitions(): - cache = _cache() - cache.alloc("a", 2) - entry = cache.get("a") - assert not entry.ready - cache.mark_ready("a") - assert entry.ready - assert entry.evictable - - -def test_mark_ready_asserts_already_ready(): - cache = _cache() - cache.alloc("a", 2) - cache.mark_ready("a") - with pytest.raises(AssertionError, match="already-ready"): - cache.mark_ready("a") - - -def test_mark_ready_asserts_missing_key(): - cache = _cache() - with pytest.raises(KeyError): - cache.mark_ready("nonexistent") - - -# ── pin / unpin ────────────────────────────────────────────────────────────── - - -def test_pin_asserts_not_ready(): - cache = _cache() - cache.alloc("a", 2) - with pytest.raises(AssertionError, match="not-ready"): - cache.pin("a") - - -def test_pin_unpin_basic(): - cache = _cache() - cache.alloc("a", 2) - cache.mark_ready("a") - cache.pin("a") - entry = cache.get("a") - assert not entry.evictable - cache.unpin("a") - assert entry.evictable - - -def test_nested_pin(): - cache = _cache() - cache.alloc("a", 2) - cache.mark_ready("a") - cache.pin("a") - cache.pin("a") - cache.unpin("a") - entry = cache.get("a") - assert not entry.evictable # still pinned (count=1) - cache.unpin("a") - assert entry.evictable - - -def test_unpin_asserts_on_unpinned(): - cache = _cache() - cache.alloc("a", 2) - cache.mark_ready("a") - with pytest.raises(AssertionError, match="unpinned"): - cache.unpin("a") - - -# ── has_held_entries ───────────────────────────────────────────────────────── - - -def test_has_held_entries_empty_cache(): - cache = _cache() - assert cache.has_held_entries() is False - - -def test_has_held_entries_true_while_not_ready(): - cache = _cache() - cache.alloc("a", 2) # not ready → held - assert cache.has_held_entries() is True - cache.mark_ready("a") # ready + unpinned → settled - assert cache.has_held_entries() is False - - -def test_has_held_entries_true_while_pinned(): - cache = _cache() - cache.alloc("a", 2) - cache.mark_ready("a") - assert cache.has_held_entries() is False - cache.pin("a") # pinned → held - assert cache.has_held_entries() is True - cache.unpin("a") - assert cache.has_held_entries() is False - - -# ── eviction ───────────────────────────────────────────────────────────────── - - -def test_eviction_frees_ready_unpinned_fifo(): - cache = _cache(num_blocks=4) - cache.alloc("a", 2) - cache.mark_ready("a") - cache.alloc("b", 2) - cache.mark_ready("b") - # Pool full. Alloc should evict "a" (oldest). - entry_c = cache.alloc("c", 2) - assert entry_c is not None - assert cache.get("a") is None # evicted - assert cache.get("b") is not None # still there - - -def test_eviction_skips_pinned(): - cache = _cache(num_blocks=4) - cache.alloc("a", 2) - cache.mark_ready("a") - cache.pin("a") - cache.alloc("b", 2) - cache.mark_ready("b") - # "a" is pinned, so "b" should be evicted instead. - entry_c = cache.alloc("c", 2) - assert entry_c is not None - assert cache.get("a") is not None # pinned, survived - assert cache.get("b") is None # evicted - cache.unpin("a") - - -def test_eviction_skips_not_ready(): - cache = _cache(num_blocks=4) - cache.alloc("a", 2) # not ready - cache.alloc("b", 2) - cache.mark_ready("b") - # "a" is not ready (non-evictable), "b" is evictable. - entry_c = cache.alloc("c", 2) - assert entry_c is not None - assert cache.get("a") is not None # not ready, survived - assert cache.get("b") is None # evicted - cache.mark_ready("a") # clean up - - -def test_eviction_returns_none_if_all_non_evictable(): - cache = _cache(num_blocks=4) - cache.alloc("a", 2) - cache.mark_ready("a") - cache.pin("a") - cache.alloc("b", 2) # not ready - # Both entries are non-evictable (a=pinned, b=not-ready). - assert cache.alloc("c", 1) is None - # Fail-fast: nothing was evicted. - assert cache.get("a") is not None - assert cache.get("b") is not None - cache.unpin("a") - - -def test_fail_fast_no_wasted_eviction(): - cache = _cache(num_blocks=6) - cache.alloc("a", 2) - cache.mark_ready("a") # evictable: 2 blocks - cache.alloc("b", 2) # not ready - cache.alloc("c", 2) # not ready - # free=0, evictable=2, need 3 -> cannot succeed - assert cache.alloc("d", 3) is None - # "a" was NOT evicted despite being evictable - assert cache.get("a") is not None - - -def test_eviction_multiple_entries(): - cache = _cache(num_blocks=6) - cache.alloc("a", 2) - cache.mark_ready("a") - cache.alloc("b", 2) - cache.mark_ready("b") - cache.alloc("c", 2) - cache.mark_ready("c") - # Need 4 blocks: must evict both "a" and "b". - entry_d = cache.alloc("d", 4) - assert entry_d is not None - assert cache.get("a") is None - assert cache.get("b") is None - assert cache.get("c") is not None - - -# ── thread safety ──────────────────────────────────────────────────────────── - - -def test_concurrent_pin_unpin(): - cache = _cache(num_blocks=8) - cache.alloc("a", 4) - cache.mark_ready("a") - errors: list[Exception] = [] - barrier = threading.Barrier(8) - - def worker(): - barrier.wait() - try: - for _ in range(500): - cache.pin("a") - cache.unpin("a") - except Exception as exc: - errors.append(exc) - - threads = [threading.Thread(target=worker) for _ in range(8)] - for t in threads: - t.start() - for t in threads: - t.join() - - assert not errors, errors - entry = cache.get("a") - assert entry.evictable # all pins balanced - - -def test_concurrent_alloc_evict(): - cache = _cache(num_blocks=4) - errors: list[Exception] = [] - barrier = threading.Barrier(4) - counter = {"i": 0} - counter_lock = threading.Lock() - - def worker(): - barrier.wait() - try: - for _ in range(200): - with counter_lock: - key = f"k{counter['i']}" - counter["i"] += 1 - entry = cache.alloc(key, 1) - if entry is not None: - cache.mark_ready(key) - except Exception as exc: - errors.append(exc) - - threads = [threading.Thread(target=worker) for _ in range(4)] - for t in threads: - t.start() - for t in threads: - t.join() - - assert not errors, errors - - -# ── discard ────────────────────────────────────────────────────────────────── - - -def test_discard_not_ready_returns_blocks(): - cache = EmbeddingCache(num_blocks=4) - entry = cache.alloc("h", 3) - assert entry is not None - cache.discard("h") - assert cache.get("h") is None - # All blocks reclaimed: a fresh 4-block alloc now fits. - assert cache.alloc("h2", 4) is not None - - -def test_discard_rejects_ready_entry(): - cache = EmbeddingCache(num_blocks=4) - cache.alloc("h", 1) - cache.mark_ready("h") - with pytest.raises(AssertionError): - cache.discard("h") - - -def test_discard_rejects_absent_key(): - cache = EmbeddingCache(num_blocks=4) - with pytest.raises(KeyError): - cache.discard("nope") - - -# ── pin_if_ready ───────────────────────────────────────────────────────────── - - -def test_pin_if_ready_pins_and_returns_entry(): - cache = EmbeddingCache(num_blocks=4) - cache.alloc("h", 2) - cache.mark_ready("h") - entry = cache.pin_if_ready("h") - assert entry is not None and entry.ready and len(entry.block_ids) == 2 - assert cache.alloc("h2", 4) is None # pinned ⇒ not evictable - - -def test_pin_if_ready_none_when_absent(): - cache = EmbeddingCache(num_blocks=4) - assert cache.pin_if_ready("nope") is None - - -def test_pin_if_ready_returns_unpinned_entry_when_not_ready(): - """The producer needs "present but not ready" told apart from "absent": - the first is NACKed as retryable, the second as a miss.""" - cache = EmbeddingCache(num_blocks=4) - cache.alloc("h", 1) # not-ready - entry = cache.pin_if_ready("h") - assert entry is not None and not entry.ready - assert cache.alloc("h2", 4) is None # not-ready ⇒ still not evictable diff --git a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py b/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py deleted file mode 100644 index ef86505057..0000000000 --- a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py +++ /dev/null @@ -1,466 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import uuid - -import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod -from tests.v1.ec_connector.unit.utils import create_ec_vllm_config -from vllm.config.ec_transfer import ECRole -from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( - ECCPUWorkerMetadata, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( - ECSharedRegion, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler - -_N = 16 -_BS = 64 - - -class _Pos: - def __init__(self, offset, length): - self.offset = offset - self.length = length - - def get_num_embeds(self): - return self.length - - -class _Feature: - def __init__(self, mm_hash, length=1, identifier=None): - self.mm_hash = mm_hash - self.identifier = identifier if identifier is not None else mm_hash - self.mm_position = _Pos(0, length) - self.modality = "image" - - -class _Request: - _counter = 0 - - def __init__(self, features, request_id=None): - self.mm_features = features - if request_id is None: - _Request._counter += 1 - self.request_id = f"req_{_Request._counter}" - else: - self.request_id = request_id - - -class _WorkerOutput: - """Stand-in for ECConnectorOutput carrying a completion report. - - `saves` are mm_hashes; `loads` are transfer ids. One entry per reporting - rank, so a list holding the same id twice is two ranks reporting it. - """ - - def __init__(self, *, saves=None, loads=None): - self.ec_connector_worker_meta = ECCPUWorkerMetadata( - completed_saves=list(saves or []), - completed_loads=list(loads or []), - ) - - -def _make_scheduler( - monkeypatch, - num_blocks=_N, - *, - ec_role: ECRole = "ec_both", - tensor_parallel_size=1, -) -> ECCPUScheduler: - region = ECSharedRegion( - engine_id=str(uuid.uuid4()), - num_blocks=num_blocks, - block_size_bytes=_BS, - ) - monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region) - - return ECCPUScheduler( - create_ec_vllm_config( - ec_role=ec_role, - tensor_parallel_size=tensor_parallel_size, - ) - ) - - -def _load_ids(meta) -> list[int]: - """Transfer ids the scheduler dispatched in this step's metadata.""" - return [transfer_id for transfer_id, _, _ in meta.loads.values()] - - -def _load_blocks(meta, mm_hash: str) -> list[int]: - """Block ids the scheduler dispatched for one load.""" - return meta.loads[mm_hash][1] - - -def _seed_cached(s: ECCPUScheduler, mm_hash: str, n_blocks: int): - """Pre-populate a ready cache entry backed by real blocks.""" - s._cache.alloc(mm_hash, n_blocks, (n_blocks, _BS // s._dtype.itemsize)) - s._cache.mark_ready(mm_hash) - - -def test_offload_reuse_cycle(monkeypatch): - s = _make_scheduler(monkeypatch) - # Two blocks, so the reload also pins down that every allocated block id - # comes back in meta.loads, not just the first. - req = _Request([_Feature("h1", length=2)]) - - # Step A: first sight — allocate and save. - assert s.has_cache_item("h1") is False - assert s.ensure_cache_available(req, 0) is True - s.update_state_after_alloc(req, 0) - meta_a = s.build_connector_meta(scheduler_output=None) - assert "h1" in meta_a.saves - assert meta_a.loads == {} - # Not ready until the worker reports the save memcpy complete. - assert s.has_cache_item("h1") is False - - # Worker reports the save copy finished → entry becomes ready. - s.update_connector_output(_WorkerOutput(saves=["h1"])) - assert s.has_cache_item("h1") is True - - # Step C: reload now works — entry is ready at update_state_after_alloc time. - s.update_state_after_alloc(req, 0) - meta_c = s.build_connector_meta(scheduler_output=None) - assert "h1" in meta_c.loads - assert _load_blocks(meta_c, "h1") == meta_a.saves["h1"] - - s.shutdown() - - -def test_has_cache_item_false_when_not_consumer(monkeypatch): - s = _make_scheduler(monkeypatch, ec_role="ec_producer") - assert s.has_cache_item("anything") is False - s.shutdown() - - -def test_connector_keys_on_identifier_not_mm_hash(monkeypatch): - """The connector must key the encoder cache on feature.identifier (what - has_cache_item is called with), NOT feature.mm_hash.""" - s = _make_scheduler(monkeypatch) - req = _Request([_Feature("PROC_KEY", length=1, identifier="ENC_KEY")]) - s.update_state_after_alloc(req, 0) - meta = s.build_connector_meta(scheduler_output=None) - assert "ENC_KEY" in meta.saves - assert "PROC_KEY" not in meta.saves - # Not ready until the completion report. - assert s.has_cache_item("ENC_KEY") is False - # After the report it becomes ready — keyed on identifier. - s.update_connector_output(_WorkerOutput(saves=["ENC_KEY"])) - assert s.has_cache_item("ENC_KEY") is True - assert s.has_cache_item("PROC_KEY") is False - s.shutdown() - - -def test_cpu_region_fifo_eviction(monkeypatch): - # Region holds exactly 2 one-block encodings. - s = _make_scheduler(monkeypatch, num_blocks=2) - - # Fill with A and B, marking each ready via a completion report. - for h in ("A", "B"): - s.update_state_after_alloc(_Request([_Feature(h, length=1)]), 0) - s.build_connector_meta(scheduler_output=None) - s.update_connector_output(_WorkerOutput(saves=[h])) - assert s.has_cache_item("A") is True - assert s.has_cache_item("B") is True - - # Allocate C: must evict A (oldest) to make room. - s.update_state_after_alloc(_Request([_Feature("C", length=1)]), 0) - s.build_connector_meta(scheduler_output=None) - assert s.has_cache_item("A") is False # evicted - assert s.has_cache_item("B") is True - - -def test_multiple_mm_items_per_request(monkeypatch): - s = _make_scheduler(monkeypatch) - req = _Request([_Feature("h1", length=1), _Feature("h2", length=1)]) - s.update_state_after_alloc(req, 0) - s.update_state_after_alloc(req, 1) - meta = s.build_connector_meta(scheduler_output=None) - assert set(meta.saves) == {"h1", "h2"} - - # Worker reports both saves complete. - s.update_connector_output(_WorkerOutput(saves=["h1", "h2"])) - assert s.has_cache_item("h1") is True - assert s.has_cache_item("h2") is True - - # Now reload both. - s.update_state_after_alloc(req, 0) - s.update_state_after_alloc(req, 1) - meta2 = s.build_connector_meta(scheduler_output=None) - assert set(meta2.loads) == {"h1", "h2"} - s.shutdown() - - -def test_loads_only_serves_hashes_touched_this_step(monkeypatch): - s = _make_scheduler(monkeypatch) - _seed_cached(s, "a", n_blocks=2) - _seed_cached(s, "b", n_blocks=2) - - s.update_state_after_alloc(_Request([_Feature("a")]), 0) - meta = s.build_connector_meta(scheduler_output=None) - assert "a" in meta.loads - assert "b" not in meta.loads - s.shutdown() - - -def test_repeated_reload_same_step_loads_once(monkeypatch): - """The same mm_hash requested twice in one step must appear in - meta.loads exactly once.""" - s = _make_scheduler(monkeypatch) - _seed_cached(s, "a", n_blocks=2) - - s.update_state_after_alloc(_Request([_Feature("a")]), 0) - s.update_state_after_alloc(_Request([_Feature("a")]), 0) - meta = s.build_connector_meta(scheduler_output=None) - assert list(meta.loads.keys()).count("a") == 1 - s.shutdown() - - -def test_load_not_emitted_for_uncached_entry(monkeypatch): - """A feature that was never saved must not appear in meta.loads.""" - s = _make_scheduler(monkeypatch) - s.update_state_after_alloc(_Request([_Feature("missing")]), 0) - meta = s.build_connector_meta(scheduler_output=None) - assert meta.loads == {} - s.shutdown() - - -def test_load_pin_protects_blocks_until_completion(monkeypatch): - """Loaded blocks stay pinned until the worker reports the load memcpy - complete. Observable: a save that needs eviction cannot reclaim a - still-pinned loaded entry, but can once the completion report arrives.""" - # 2 blocks, both occupied by "a" (ready). - s = _make_scheduler(monkeypatch, num_blocks=2) - _seed_cached(s, "a", n_blocks=2) - - # Load "a" — pins it. - s.update_state_after_alloc(_Request([_Feature("a")]), 0) - meta = s.build_connector_meta(scheduler_output=None) - assert "a" in meta.loads - (transfer_id,) = _load_ids(meta) - - # Try to save "b" (needs 2 blocks, but "a" is pinned) — must fail. - s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0) - meta2 = s.build_connector_meta(scheduler_output=None) - assert "b" not in meta2.saves # can't evict pinned "a" - - # The one participant reports the load complete → "a" unpins. Now "b" can - # evict it. - s.update_connector_output(_WorkerOutput(loads=[transfer_id])) - s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0) - meta3 = s.build_connector_meta(scheduler_output=None) - assert "b" in meta3.saves - s.shutdown() - - -def _can_evict(s: ECCPUScheduler, mm_hash: str, n_blocks: int) -> bool: - """Whether a new save of `mm_hash` can claim `n_blocks`, i.e. whether the - blocks currently held by other entries are reclaimable.""" - s.update_state_after_alloc(_Request([_Feature(mm_hash, length=n_blocks)]), 0) - meta = s.build_connector_meta(scheduler_output=None) - return mm_hash in meta.saves - - -def test_load_needs_every_participant_before_unpin(monkeypatch): - """With two participating ranks, one report must not release the pin. - - This is the multi-rank case: each rank copies the same CPU blocks into its - own GPU memory, so the blocks stay in use until the last rank is done. - """ - s = _make_scheduler(monkeypatch, num_blocks=2, tensor_parallel_size=2) - _seed_cached(s, "a", n_blocks=2) - - s.update_state_after_alloc(_Request([_Feature("a")]), 0) - (transfer_id,) = _load_ids(s.build_connector_meta(scheduler_output=None)) - - # First rank reports: still in use by the second, so "a" must hold. - s.update_connector_output(_WorkerOutput(loads=[transfer_id])) - assert _can_evict(s, "b", 2) is False - assert s.has_pending_push_work() is True - - # Second rank reports: now it is genuinely free. - s.update_connector_output(_WorkerOutput(loads=[transfer_id])) - assert s.has_pending_push_work() is False - assert _can_evict(s, "c", 2) is True - s.shutdown() - - -def test_late_report_does_not_release_a_later_load_of_same_hash(monkeypatch): - """A straggling report from one dispatch must not release the pin taken by - a subsequent dispatch of the same mm_hash. - - This is what the transfer id is for: mm_hash identifies a cache entry, not - a transfer, so it cannot tell the two dispatches apart. - """ - s = _make_scheduler(monkeypatch, num_blocks=2, tensor_parallel_size=2) - _seed_cached(s, "a", n_blocks=2) - - # First dispatch, fully reported by both ranks → pin released. - s.update_state_after_alloc(_Request([_Feature("a")]), 0) - (first_id,) = _load_ids(s.build_connector_meta(scheduler_output=None)) - s.update_connector_output(_WorkerOutput(loads=[first_id, first_id])) - assert s.has_pending_push_work() is False - - # Second dispatch of the same hash takes a fresh pin. - s.update_state_after_alloc(_Request([_Feature("a")]), 0) - (second_id,) = _load_ids(s.build_connector_meta(scheduler_output=None)) - assert second_id != first_id - - # A replayed report for the first dispatch must be ignored entirely. - s.update_connector_output(_WorkerOutput(loads=[first_id, first_id])) - assert s.has_pending_push_work() is True - assert _can_evict(s, "b", 2) is False - s.shutdown() - - -def test_region_full_skips_save_and_never_blocks(monkeypatch): - """When the region is fully occupied by pinned entries, new saves are - silently skipped and ensure_cache_available never blocks.""" - s = _make_scheduler(monkeypatch, num_blocks=1) - _seed_cached(s, "pinned", n_blocks=1) - s._cache.pin("pinned") - - req = _Request([_Feature("new", length=1)]) - assert s.ensure_cache_available(req, 0) is True - - s.update_state_after_alloc(req, 0) - meta = s.build_connector_meta(scheduler_output=None) - assert "new" not in meta.saves - assert s.has_cache_item("new") is False - - s._cache.unpin("pinned") - s.shutdown() - - -def test_producer_only_never_emits_loads(monkeypatch): - """A producer-only scheduler must never populate meta.loads, even when - entries are ready.""" - s = _make_scheduler(monkeypatch, ec_role="ec_producer") - req = _Request([_Feature("h1", length=1)]) - - s.update_state_after_alloc(req, 0) - meta = s.build_connector_meta(scheduler_output=None) - assert "h1" in meta.saves - assert meta.loads == {} - - s.update_connector_output(_WorkerOutput(saves=["h1"])) - - # Even if we re-encounter the same feature, no load. - s.update_state_after_alloc(req, 0) - meta2 = s.build_connector_meta(scheduler_output=None) - assert meta2.loads == {} - s.shutdown() - - -def test_consumer_only_never_emits_saves(monkeypatch): - """A consumer-only scheduler must never populate meta.saves.""" - s = _make_scheduler(monkeypatch, ec_role="ec_consumer") - _seed_cached(s, "a", n_blocks=1) - - req = _Request([_Feature("a", length=1)]) - s.update_state_after_alloc(req, 0) - meta = s.build_connector_meta(scheduler_output=None) - assert meta.saves == {} - assert "a" in meta.loads - s.shutdown() - - -def test_consumer_only_has_cache_item(monkeypatch): - """Consumer-only: has_cache_item works normally for ready entries.""" - s = _make_scheduler(monkeypatch, ec_role="ec_consumer") - assert s.has_cache_item("a") is False - _seed_cached(s, "a", n_blocks=1) - assert s.has_cache_item("a") is True - s.shutdown() - - -def test_save_not_emitted_for_already_cached_entry(monkeypatch): - """An entry already in the cache (from a prior save) must not trigger - a second allocation or appear in meta.saves again.""" - s = _make_scheduler(monkeypatch) - req = _Request([_Feature("h1", length=1)]) - - # First encounter → save. - s.update_state_after_alloc(req, 0) - meta1 = s.build_connector_meta(scheduler_output=None) - assert "h1" in meta1.saves - - # Second encounter (same step or later) → no save. - s.update_state_after_alloc(req, 0) - meta2 = s.build_connector_meta(scheduler_output=None) - assert "h1" not in meta2.saves - s.shutdown() - - -# ── completion reporting (update_connector_output) ────────────────────────── - - -def test_update_connector_output_ignores_foreign_meta(monkeypatch): - """A worker output whose payload is not an ECCPUWorkerMetadata is a - no-op — no cache mutation, no crash.""" - s = _make_scheduler(monkeypatch) - _seed_cached(s, "a", n_blocks=1) - - class _Foreign: - ec_connector_worker_meta = object() - - s.update_connector_output(_Foreign()) # must not raise - assert s.has_cache_item("a") is True - s.shutdown() - - -def test_update_connector_output_ignores_unknown_report(monkeypatch): - """Reports for a never-seen save hash or a transfer id this scheduler never - dispatched are dropped rather than mutating unrelated state.""" - s = _make_scheduler(monkeypatch) - s.update_connector_output(_WorkerOutput(saves=["gone"], loads=[4242])) - assert s.has_cache_item("gone") is False - s.shutdown() - - -# ── has_pending_push_work ─────────────────────────────────────────────────── - - -def test_has_pending_push_work_tracks_inflight_save(monkeypatch): - """A dispatched-but-unconfirmed save keeps push work pending until the - completion report marks it ready.""" - s = _make_scheduler(monkeypatch) - assert s.has_pending_push_work() is False - - s.update_state_after_alloc(_Request([_Feature("h1", length=1)]), 0) - s.build_connector_meta(scheduler_output=None) - assert s.has_pending_push_work() is True # save not yet confirmed - - s.update_connector_output(_WorkerOutput(saves=["h1"])) - assert s.has_pending_push_work() is False - s.shutdown() - - -def test_has_pending_push_work_tracks_inflight_load(monkeypatch): - """A dispatched-but-unconfirmed load keeps push work pending until the - unpin completion report arrives.""" - s = _make_scheduler(monkeypatch) - _seed_cached(s, "a", n_blocks=1) - assert s.has_pending_push_work() is False # ready + unpinned - - s.update_state_after_alloc(_Request([_Feature("a")]), 0) - meta = s.build_connector_meta(scheduler_output=None) - assert s.has_pending_push_work() is True # pinned, awaiting unpin - - s.update_connector_output(_WorkerOutput(loads=_load_ids(meta))) - assert s.has_pending_push_work() is False - s.shutdown() - - -# ── shutdown ──────────────────────────────────────────────────────────────── - - -def test_shutdown_disables_roles_and_cleans_region(monkeypatch): - """After shutdown the scheduler serves no items and the region is - cleaned up.""" - s = _make_scheduler(monkeypatch, num_blocks=2) - _seed_cached(s, "a", n_blocks=2) - assert s.has_cache_item("a") is True - - s.shutdown() - assert s.has_cache_item("a") is False # consumer role disabled diff --git a/tests/v1/ec_connector/unit/cpu/test_connector.py b/tests/v1/ec_connector/unit/cpu/test_connector.py deleted file mode 100644 index cd5ba91969..0000000000 --- a/tests/v1/ec_connector/unit/cpu/test_connector.py +++ /dev/null @@ -1,71 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from unittest.mock import MagicMock - -import pytest - -from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole -from vllm.distributed.ec_transfer.ec_connector.cpu.connector import ECCPUConnector -from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory - - -def _cfg(): - ec = MagicMock() - ec.is_ec_producer = True - ec.is_ec_consumer = True - cfg = MagicMock() - cfg.ec_transfer_config = ec - return cfg - - -def test_scheduler_role_builds_only_scheduler(monkeypatch): - fake_sched = MagicMock() - monkeypatch.setattr(ECCPUConnector, "_make_scheduler", lambda self, cfg: fake_sched) - c = ECCPUConnector(_cfg(), ECConnectorRole.SCHEDULER) - assert c.connector_scheduler is fake_sched - assert c.connector_worker is None - c.has_cache_item("x") - fake_sched.has_cache_item.assert_called_once_with("x") - - -def test_worker_role_builds_only_worker(monkeypatch): - fake_worker = MagicMock() - monkeypatch.setattr(ECCPUConnector, "_make_worker", lambda self, cfg: fake_worker) - c = ECCPUConnector(_cfg(), ECConnectorRole.WORKER) - assert c.connector_worker is fake_worker - assert c.connector_scheduler is None - - -@pytest.mark.parametrize("role", [ECConnectorRole.SCHEDULER, ECConnectorRole.WORKER]) -def test_v1_model_runner_is_rejected(monkeypatch, role): - """The V1 runner never reports save completions, so refuse to start.""" - monkeypatch.setattr( - ECCPUConnector, "_make_scheduler", lambda self, cfg: MagicMock() - ) - monkeypatch.setattr(ECCPUConnector, "_make_worker", lambda self, cfg: MagicMock()) - cfg = _cfg() - cfg.use_v2_model_runner = False - - with pytest.raises(ValueError, match="V2 model runner"): - ECCPUConnector(cfg, role) - - -def test_factory_registered(): - cls = ECConnectorFactory._registry["ECCPUConnector"]() - assert cls is ECCPUConnector - - -def test_request_finished_forwards_to_scheduler(monkeypatch): - from unittest.mock import MagicMock - - import vllm.distributed.ec_transfer.ec_connector.cpu.connector as conn_mod # noqa: F401 - from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole - from vllm.distributed.ec_transfer.ec_connector.cpu.connector import ECCPUConnector - - fake_sched = MagicMock() - fake_sched.request_finished.return_value = (False, {"h1": {"peer_port": 1}}) - monkeypatch.setattr(ECCPUConnector, "_make_scheduler", lambda self, cfg: fake_sched) - c = ECCPUConnector(_cfg(), ECConnectorRole.SCHEDULER) - req = MagicMock() - assert c.request_finished(req) == (False, {"h1": {"peer_port": 1}}) - fake_sched.request_finished.assert_called_once_with(req) diff --git a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py deleted file mode 100644 index 0b64d68ca6..0000000000 --- a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py +++ /dev/null @@ -1,396 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for ECSharedRegion (the mmap substrate).""" - -import contextlib -import ctypes -import errno -import mmap -import os -import uuid -from unittest.mock import MagicMock, patch - -import pytest -import torch - -from vllm.config import ParallelConfig -from vllm.distributed.ec_transfer.ec_connector.cpu import common -from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( - ECSharedRegion, - _wait_for_file_size, -) - - -def _make_region(num_blocks: int = 8, block_size_bytes: int = 64) -> ECSharedRegion: - return ECSharedRegion( - engine_id=str(uuid.uuid4()), - num_blocks=num_blocks, - block_size_bytes=block_size_bytes, - ) - - -@pytest.fixture -def region() -> ECSharedRegion: - r = _make_region() - yield r - r.cleanup() - - -@contextlib.contextmanager -def _region(**kwargs): - """Context manager: create one region, clean up on exit.""" - r = _make_region(**kwargs) - try: - yield r - finally: - r.cleanup() - - -def _page_residency(mmap_obj: mmap.mmap, length: int) -> list[bool]: - """Return Linux page-residency bits for a writable mmap.""" - try: - libc = ctypes.CDLL(None, use_errno=True) - mincore = libc.mincore - except AttributeError: - pytest.skip("mincore is unavailable") - - page_count = (length + mmap.PAGESIZE - 1) // mmap.PAGESIZE - mincore.argtypes = [ - ctypes.c_void_p, - ctypes.c_size_t, - ctypes.POINTER(ctypes.c_ubyte), - ] - mincore.restype = ctypes.c_int - vector = (ctypes.c_ubyte * page_count)() - address = ctypes.addressof(ctypes.c_ubyte.from_buffer(mmap_obj)) - result = mincore(ctypes.c_void_p(address), length, vector) - if result != 0: - raise OSError(ctypes.get_errno()) - return [bool(value & 1) for value in vector] - - -# ── mmap sharing between two instances ─────────────────────────────────────── - - -def test_second_instance_opens_existing_file_and_shares_memory(): - instance_id = str(uuid.uuid4()) - r1 = ECSharedRegion( - engine_id=instance_id, - num_blocks=4, - block_size_bytes=64, - ) - try: - r2 = ECSharedRegion( - engine_id=instance_id, - num_blocks=4, - block_size_bytes=64, - ) - try: - assert r1._is_creator - assert not r2._is_creator - # Both map the same physical pages; writes via r1 are visible via r2. - r1.blocks[0, :4] = torch.tensor([10, 20, 30, 40], dtype=torch.int8) - assert r2.blocks[0, :4].tolist() == [10, 20, 30, 40] - finally: - r2.cleanup() - finally: - r1.cleanup() - - -def test_dense_dp_engines_get_distinct_regions(): - """Dense DP engines reset data_parallel_rank to 0 but share instance_id, - so the region must be keyed by data_parallel_index.""" - instance_id = str(uuid.uuid4()) - regions = [] - try: - for dp_index in range(2): - parallel_config = ParallelConfig( - data_parallel_size=2, data_parallel_rank=dp_index - ) - parallel_config.reconfigure_for_independent_dp_rank() - cfg = MagicMock() - cfg.instance_id = instance_id - cfg.parallel_config = parallel_config - cfg.model_config.dtype = torch.float16 - cfg.model_config.get_inputs_embeds_size.return_value = 8 - cfg.ec_transfer_config.ec_connector_extra_config = {"ec_cpu_bytes": 64} - regions.append(common.create_ec_shared_region(cfg)) - - assert all(r._is_creator for r in regions) - assert regions[0]._mmap_path != regions[1]._mmap_path - finally: - for r in regions: - r.cleanup() - - -def test_only_creator_unlinks_file_on_cleanup(): - """Critical contract: if the non-creator unlinks, the creator's mmap path - becomes a dangling backing file and a third opener would create a new one - out from under the creator.""" - instance_id = str(uuid.uuid4()) - r1 = ECSharedRegion( - engine_id=instance_id, - num_blocks=4, - block_size_bytes=64, - ) - path = r1._mmap_path - r2 = ECSharedRegion( - engine_id=instance_id, - num_blocks=4, - block_size_bytes=64, - ) - - # Non-creator goes away first — file must still be on disk for r1. - r2.cleanup() - assert os.path.exists(path), "non-creator cleanup must not unlink the file" - - # Creator goes away — file is removed. - r1.cleanup() - assert not os.path.exists(path), "creator cleanup must unlink the file" - - -# ── _wait_for_file_size ────────────────────────────────────────────────────── - - -def test_wait_for_file_size_returns_when_already_big_enough(tmp_path): - """The fast path: file already at expected size — return immediately.""" - p = tmp_path / "f.bin" - p.write_bytes(b"\x00" * 128) - fd = os.open(str(p), os.O_RDONLY) - try: - _wait_for_file_size(fd, expected_size=128, timeout=1.0) # must not raise - finally: - os.close(fd) - - -def test_wait_for_file_size_times_out_when_file_stays_empty(tmp_path): - p = tmp_path / "f.bin" - p.write_bytes(b"") - fd = os.open(str(p), os.O_RDONLY) - try: - with pytest.raises(TimeoutError): - _wait_for_file_size(fd, expected_size=4096, timeout=0.05) - finally: - os.close(fd) - - -# ── MADV_POPULATE_WRITE pre-faulting / fallback ────────────────────────────── - - -def test_madvise_success_selects_madvise_population(monkeypatch): - """A successful probe must keep using MADV_POPULATE_WRITE — the fallback - helper must never be invoked on a kernel that accepts the advice. - - Spies on both helpers so we can verify call counts: 1 probe call over a - single page + 1 populate call over the whole flat region, fallback 0 times. - """ - from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr - - madvise_calls: list[tuple[int, int, int]] = [] - fallback_calls: list[tuple[int, int]] = [] - - def _spy_madvise(mm, off, ln): - madvise_calls.append((off, ln, id(mm))) - - def _spy_fallback(mm, off, ln): - fallback_calls.append((off, ln)) - - monkeypatch.setattr(esr, "_madvise_populate_write", _spy_madvise) - monkeypatch.setattr(esr, "_fallback_populate_write", _spy_fallback) - - num_blocks, block_size_bytes = 4, mmap.PAGESIZE - with _region(num_blocks=num_blocks, block_size_bytes=block_size_bytes): - total = num_blocks * block_size_bytes - mmap_id = madvise_calls[0][2] - assert madvise_calls == [ - (0, mmap.PAGESIZE, mmap_id), # probe - (0, total, mmap_id), # flat pre-fault of the whole region - ] - assert fallback_calls == [], ( - "native-path constructor must not invoke the fallback helper" - ) - - -def test_madvise_einval_selects_fallback_for_whole_region(monkeypatch): - """An EINVAL probe must select the fallback for the whole flat region. - - This is the regression under test: on Linux < 5.14 the constant is absent - from the `mmap` module, so advice=23 reaches a kernel that rejects it and - every ECSharedRegion construction used to abort with EINVAL. - """ - from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr - - fallback_calls: list[tuple[int, int]] = [] - real_fallback = esr._fallback_populate_write - - def _raise_einval(mm, off, ln): - raise OSError(errno.EINVAL, "simulated unsupported kernel") - - def _spy_fallback(mm, off, ln): - fallback_calls.append((off, ln)) - return real_fallback(mm, off, ln) - - monkeypatch.setattr(esr, "_madvise_populate_write", _raise_einval) - monkeypatch.setattr(esr, "_fallback_populate_write", _spy_fallback) - - num_blocks, block_size_bytes = 4, mmap.PAGESIZE - with _region(num_blocks=num_blocks, block_size_bytes=block_size_bytes) as r: - assert fallback_calls == [(0, num_blocks * block_size_bytes)] - # Region is still fully usable after taking the fallback path. - r.blocks[0, :4] = torch.tensor([1, 2, 3, 4], dtype=torch.int8) - assert r.blocks[0, :4].tolist() == [1, 2, 3, 4] - - -def test_non_creator_does_not_probe_or_populate(monkeypatch): - """Only the creator pre-faults. A second opener must not probe the advice - nor re-touch pages the creator already populated.""" - from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr - - instance_id = str(uuid.uuid4()) - r1 = ECSharedRegion( - engine_id=instance_id, num_blocks=4, block_size_bytes=mmap.PAGESIZE - ) - try: - calls: list[str] = [] - monkeypatch.setattr( - esr, - "_madvise_populate_write", - lambda mm, off, ln: calls.append("madvise"), - ) - monkeypatch.setattr( - esr, - "_fallback_populate_write", - lambda mm, off, ln: calls.append("fallback"), - ) - r2 = ECSharedRegion( - engine_id=instance_id, num_blocks=4, block_size_bytes=mmap.PAGESIZE - ) - try: - assert not r2._is_creator - assert calls == [] - finally: - r2.cleanup() - finally: - r1.cleanup() - - -def test_fallback_populate_write_preserves_bytes_and_faults_pages(): - """Fallback preserves existing bytes and touches every target page. - - `|= 0` rather than `= 0` matters here: a peer worker may already have - written EC data into the shared mmap before we pre-fault it. - """ - from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr - - size = 3 * mmap.PAGESIZE - if not hasattr(mmap, "MADV_DONTNEED"): - pytest.skip("MADV_DONTNEED is unavailable") - - mmap_obj = mmap.mmap( - -1, - size, - flags=mmap.MAP_SHARED, - prot=mmap.PROT_READ | mmap.PROT_WRITE, - ) - try: - mmap_obj.madvise(mmap.MADV_DONTNEED, 0, size) - if any(_page_residency(mmap_obj, size)): - pytest.skip("kernel did not discard anonymous mmap pages") - - mmap_obj[0] = 0xAB - assert _page_residency(mmap_obj, size) == [True, False, False] - - esr._fallback_populate_write(mmap_obj, 0, size) - - assert _page_residency(mmap_obj, size) == [True, True, True] - assert mmap_obj[0] == 0xAB - finally: - mmap_obj.close() - - -def test_madvise_unexpected_oserror_propagates(monkeypatch): - """Only EINVAL triggers the fallback. Other OSErrors (e.g. EIO) must - propagate out of __init__, not be silently masked by the fallback branch. - """ - from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr - - monkeypatch.setattr( - esr, - "_madvise_populate_write", - lambda mm, off, ln: (_ for _ in ()).throw( - OSError(errno.EIO, "simulated I/O failure") - ), - ) - engine_id = str(uuid.uuid4()) - try: - with pytest.raises(OSError) as exc_info: - ECSharedRegion( - engine_id=engine_id, num_blocks=4, block_size_bytes=mmap.PAGESIZE - ) - assert exc_info.value.errno == errno.EIO - finally: - # __init__ raised past the unlink-owning cleanup(), so drop the - # backing file here rather than leaking it into /dev/shm. - with contextlib.suppress(FileNotFoundError): - os.unlink(f"/dev/shm/vllm_ec_{engine_id}.mmap") - - -# ── pin_memory ─────────────────────────────────────────────────────────────── - - -def test_pin_memory_success_sets_flag(region): - """When cudaHostRegister returns 0, _is_pinned flips to True - and cleanup will correspondingly call cudaHostUnregister.""" - fake_cudart = MagicMock() - success = MagicMock() - success.value = 0 - fake_cudart.cudaHostRegister.return_value = success - fake_cudart.cudaHostUnregister.return_value = success - - with ( - patch("torch.cuda.is_available", return_value=True), - patch("torch.cuda.cudart", return_value=fake_cudart), - ): - region.pin_memory() - assert region._is_pinned is True - # cleanup must pair with cudaHostUnregister exactly once. - region.cleanup() - fake_cudart.cudaHostUnregister.assert_called_once() - - -def test_pin_memory_failure_leaves_flag_false(): - """If cudaHostRegister fails (non-zero), don't pretend it succeeded — - cleanup must NOT call cudaHostUnregister on memory we never registered.""" - r = _make_region() - try: - fake_cudart = MagicMock() - fail = MagicMock() - fail.value = 1 # non-zero == error - fake_cudart.cudaHostRegister.return_value = fail - - with ( - patch("torch.cuda.is_available", return_value=True), - patch("torch.cuda.cudart", return_value=fake_cudart), - ): - r.pin_memory() - assert r._is_pinned is False - # Now run cleanup and verify cudaHostUnregister was NOT called. - r.cleanup() - fake_cudart.cudaHostUnregister.assert_not_called() - finally: - r.cleanup() - - -def test_pin_memory_noop_without_cuda(region): - """pin_memory is a no-op when CUDA is not available.""" - with patch("torch.cuda.is_available", return_value=False): - region.pin_memory() - assert region._is_pinned is False - - -# ── cleanup idempotency ─────────────────────────────────────────────────────── - - -def test_cleanup_is_idempotent(region): - region.cleanup() - region.cleanup() # fixture calls a third time — must not raise diff --git a/tests/v1/ec_connector/unit/test_control.py b/tests/v1/ec_connector/unit/test_control.py deleted file mode 100644 index 93701738eb..0000000000 --- a/tests/v1/ec_connector/unit/test_control.py +++ /dev/null @@ -1,285 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for control-layer ABCs, ZmqClientConnection, and ECSession.""" - -from unittest.mock import MagicMock, patch - -import pytest -import zmq - -from vllm.distributed.ec_transfer.ec_connector.cpu.control.base import ( - ControlConnection, -) - - -class _ConcreteConn(ControlConnection): - """Minimal concrete subclass for ABC-instantiation tests.""" - - def __init__(self): - self._alive = True - self._sent = [] - self._inbox = [] - - @property - def alive(self) -> bool: - return self._alive - - def send(self, msg: bytes) -> None: - self._sent.append(msg) - - def recv(self) -> list[bytes]: - msgs, self._inbox = self._inbox, [] - return msgs - - def mark_dead(self) -> None: - self._alive = False - - def close(self) -> None: - self._alive = False - - -def test_control_connection_abc_cannot_be_instantiated_directly(): - with pytest.raises(TypeError): - ControlConnection() # type: ignore[abstract] - - -def test_concrete_connection_send_recv(): - conn = _ConcreteConn() - assert conn.alive - conn.send(b"hello") - conn._inbox.append(b"world") - assert conn.recv() == [b"world"] - assert conn._sent == [b"hello"] - - -def test_mark_dead_sets_alive_false(): - conn = _ConcreteConn() - conn.mark_dead() - assert not conn.alive - - -def test_close_sets_alive_false(): - conn = _ConcreteConn() - conn.close() - assert not conn.alive - - -# ── ZmqClientConnection ─────────────────────────────────────────────────────── - - -from vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq import ( # noqa: E402 - ZmqClientConnection, - ZmqClientTransport, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.utils import PeerAddr # noqa: E402 - - -def _conn(messages: list[list[bytes]], monitor=None) -> ZmqClientConnection: - """Build a ZmqClientConnection whose dealer yields the given frame-lists.""" - dealer = MagicMock() - dealer.recv_multipart.side_effect = [*messages, zmq.Again()] - return ZmqClientConnection(dealer=dealer, monitor=monitor) - - -def test_zmq_client_connection_recv_extracts_payloads_from_envelopes(): - conn = _conn([[b"", b"msg1"], [b"", b"msg2"]]) - assert conn.recv() == [b"msg1", b"msg2"] - - -def test_zmq_client_connection_recv_skips_frames_without_empty_delimiter(): - # A frame that does not start with b"" is malformed and must be dropped. - conn = _conn([[b"", b"ok"], [b"no-delimiter", b"bad"]]) - assert conn.recv() == [b"ok"] - - -def test_zmq_client_connection_recv_breaks_on_non_again_exception(): - # A socket error should stop the drain and return whatever was collected. - dealer = MagicMock() - dealer.recv_multipart.side_effect = [[b"", b"first"], OSError("broken")] - conn = ZmqClientConnection(dealer=dealer, monitor=None) - assert conn.recv() == [b"first"] - - -def test_zmq_client_connection_send_prefixes_empty_delimiter(): - dealer = MagicMock() - conn = ZmqClientConnection(dealer=dealer, monitor=None) - conn.send(b"payload") - dealer.send_multipart.assert_called_once_with([b"", b"payload"]) - - -def test_zmq_client_connection_mark_dead_does_not_close_socket(): - # mark_dead signals liveness only; resource teardown is close()'s job. - dealer = MagicMock() - conn = ZmqClientConnection(dealer=dealer, monitor=None) - conn.mark_dead() - assert not conn.alive - dealer.close.assert_not_called() - - -def test_zmq_client_connection_close_without_monitor(): - dealer = MagicMock() - conn = ZmqClientConnection(dealer=dealer, monitor=None) - conn.close() - dealer.close.assert_called_once_with(linger=0) - assert not conn.alive - - -def test_zmq_client_connection_close_with_monitor_disables_then_closes_both(): - # The disable_monitor → monitor.close → dealer.close sequence must be - # respected: closing the dealer before disabling the monitor can cause - # a zmq assertion in some versions. - dealer = MagicMock() - monitor = MagicMock() - conn = ZmqClientConnection(dealer=dealer, monitor=monitor) - conn.close() - dealer.disable_monitor.assert_called_once() - monitor.close.assert_called_once_with(linger=0) - dealer.close.assert_called_once_with(linger=0) - assert not conn.alive - - -def test_zmq_client_connection_close_is_idempotent(): - dealer = MagicMock() - conn = ZmqClientConnection(dealer=dealer, monitor=None) - conn.close() - conn.close() # must not raise - assert not conn.alive - - -# ── ZmqClientTransport ──────────────────────────────────────────────────────── - - -def _make_transport() -> tuple[ZmqClientTransport, MagicMock]: - """Return a ZmqClientTransport with a mocked ZMQ context.""" - ctx = MagicMock() - ctx.socket.return_value = MagicMock() - with ( - patch( - "vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context", - return_value=ctx, - ), - patch( - "vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket", - return_value=MagicMock(), - ), - patch( - "vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path", - return_value="tcp://host:0", - ), - ): - transport = ZmqClientTransport() - transport._ctx = ctx - return transport, ctx - - -def _inject_conn(transport: ZmqClientTransport, addr: PeerAddr) -> ZmqClientConnection: - """Inject a ZmqClientConnection directly into the transport pool.""" - conn = ZmqClientConnection(dealer=MagicMock(), monitor=MagicMock()) - transport._connections[addr] = conn - return conn - - -@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket") -@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path") -@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context") -def test_zmq_client_transport_connect_caches_connection( - mock_ctx_cls, mock_path, mock_sock -): - mock_ctx_cls.return_value = MagicMock() - mock_ctx_cls.return_value.socket.return_value = MagicMock() - mock_sock.return_value = MagicMock() - transport = ZmqClientTransport() - addr: PeerAddr = ("host", 1234) - assert transport.connect(addr) is transport.connect(addr) - - -@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket") -@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path") -@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context") -def test_zmq_client_transport_connect_creates_independent_conns_per_addr( - mock_ctx_cls, mock_path, mock_sock -): - mock_ctx_cls.return_value = MagicMock() - mock_ctx_cls.return_value.socket.return_value = MagicMock() - mock_sock.return_value = MagicMock() - transport = ZmqClientTransport() - conn_a = transport.connect(("host-a", 1111)) - conn_b = transport.connect(("host-b", 2222)) - assert conn_a is not conn_b - - -@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.recv_monitor_message") -def test_zmq_client_transport_poll_dead_detects_disconnected_peer(mock_recv_mon): - transport, _ = _make_transport() - addr: PeerAddr = ("host", 1234) - conn = _inject_conn(transport, addr) - mock_recv_mon.side_effect = [{"event": zmq.EVENT_DISCONNECTED}, zmq.Again()] - - dead = transport.poll_dead() - - assert addr in dead - assert not conn.alive - - -@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.recv_monitor_message") -def test_zmq_client_transport_poll_dead_skips_conn_without_monitor(mock_recv_mon): - transport, _ = _make_transport() - addr: PeerAddr = ("host", 9999) - conn = ZmqClientConnection(dealer=MagicMock(), monitor=None) - transport._connections[addr] = conn - - dead = transport.poll_dead() - - assert dead == [] - mock_recv_mon.assert_not_called() - - -@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.recv_monitor_message") -def test_zmq_client_transport_poll_dead_ignores_non_disconnect_events(mock_recv_mon): - transport, _ = _make_transport() - addr: PeerAddr = ("host", 1234) - conn = _inject_conn(transport, addr) - # A CONNECTED event should not trigger a dead notification. - mock_recv_mon.side_effect = [{"event": zmq.EVENT_CONNECTED}, zmq.Again()] - - dead = transport.poll_dead() - - assert dead == [] - assert conn.alive - - -def test_zmq_client_transport_remove_closes_and_evicts_conn(): - transport, _ = _make_transport() - addr: PeerAddr = ("host", 5678) - conn = _inject_conn(transport, addr) - - removed = transport.remove(addr) - - assert removed is conn - assert addr not in transport._connections - conn.dealer.close.assert_called() - - -def test_zmq_client_transport_remove_returns_none_for_unknown_addr(): - transport, _ = _make_transport() - assert transport.remove(("ghost", 0)) is None - - -@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket") -@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path") -@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context") -def test_zmq_client_transport_close_closes_all_connections( - mock_ctx_cls, mock_path, mock_sock -): - mock_ctx = MagicMock() - mock_ctx_cls.return_value = mock_ctx - mock_ctx.socket.return_value = MagicMock() - dealer_a, dealer_b = MagicMock(), MagicMock() - mock_sock.side_effect = [dealer_a, dealer_b] - transport = ZmqClientTransport() - transport.connect(("a", 1)) - transport.connect(("b", 2)) - transport.close() - dealer_a.close.assert_called() - dealer_b.close.assert_called() - assert transport._connections == {} diff --git a/tests/v1/ec_connector/unit/test_data.py b/tests/v1/ec_connector/unit/test_data.py deleted file mode 100644 index 33f5f11037..0000000000 --- a/tests/v1/ec_connector/unit/test_data.py +++ /dev/null @@ -1,184 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for DataTransport ABC and NixlDataTransport.""" - -from unittest.mock import MagicMock, patch - -import pytest - -from vllm.distributed.ec_transfer.ec_connector.cpu.data.base import ( - DataTransport, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.utils import ( - serialize_mem_descriptor, -) - -_NIXL_PATH = "vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl" - - -# ── helpers ─────────────────────────────────────────────────────────────────── - - -def _make_nixl_mock(prep_side_effect=None) -> MagicMock: - nixl = MagicMock() - nixl.get_agent_metadata.return_value = b"agent-meta" - nixl.get_reg_descs.return_value = MagicMock() - nixl.get_xfer_descs.return_value = MagicMock() - nixl.prep_xfer_dlist.side_effect = prep_side_effect or [77, 88] - nixl.add_remote_agent.return_value = "remote-agent-1" - nixl.get_new_notifs.return_value = {} - nixl.check_xfer_state.return_value = "PROC" - return nixl - - -def _make_transport(nixl_mock=None): - from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import ( - NixlDataTransport, - ) - - mock = nixl_mock or _make_nixl_mock() - with ( - patch(f"{_NIXL_PATH}.NixlWrapper", return_value=mock), - patch(f"{_NIXL_PATH}.nixl_agent_config", return_value=MagicMock()), - ): - t = NixlDataTransport("agent", 0, 4, 64, 256) - t._nixl = mock - return t - - -# ── ABC ─────────────────────────────────────────────────────────────────────── - - -def test_data_transport_abc_cannot_be_instantiated(): - with pytest.raises(TypeError): - DataTransport() # type: ignore[abstract] - - -# ── NixlDataTransport construction ─────────────────────────────────────────── - - -def test_nixl_data_transport_registers_memory_on_init(): - nixl = _make_nixl_mock() - _make_transport(nixl) - nixl.register_memory.assert_called_once() - - -def test_nixl_data_transport_unwinds_registration_if_init_fails(): - """A failure after register_memory must not leave the region pinned. - - __init__ propagates, so the caller never gets a transport to deregister. - """ - nixl = _make_nixl_mock(prep_side_effect=RuntimeError("no dlist")) - from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import ( - NixlDataTransport, - ) - - with ( - patch(f"{_NIXL_PATH}.NixlWrapper", return_value=nixl), - patch(f"{_NIXL_PATH}.nixl_agent_config", return_value=MagicMock()), - pytest.raises(RuntimeError, match="no dlist"), - ): - NixlDataTransport("agent", 0, 4, 64, 256) - - nixl.deregister_memory.assert_called_once_with(nixl.get_reg_descs.return_value) - - -def test_nixl_data_transport_raises_if_nixl_unavailable(): - from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import ( - NixlDataTransport, - ) - - with ( - patch(f"{_NIXL_PATH}.NixlWrapper", None), - patch(f"{_NIXL_PATH}.nixl_agent_config", None), - pytest.raises(RuntimeError, match="requires NIXL"), - ): - NixlDataTransport("a", 0, 1, 64, 64) - - -# ── public accessors ────────────────────────────────────────────────────────── - - -def test_get_agent_metadata_returns_nixl_blob(): - assert _make_transport().get_agent_metadata() == b"agent-meta" - - -def test_get_mem_descriptor_roundtrips_block_descs(): - t = _make_transport() - raw = t.get_mem_descriptor() - assert isinstance(raw, bytes) and len(raw) > 0 - - -# ── add_remote_peer ─────────────────────────────────────────────────────────── - - -def test_add_remote_peer_returns_agent_name_and_stores_handle_internally(): - nixl = _make_nixl_mock(prep_side_effect=[77, 88]) - t = _make_transport(nixl) - mem_desc = serialize_mem_descriptor([(0, 64, 0)]) - agent_name = t.add_remote_peer(b"meta", mem_desc) - assert agent_name == "remote-agent-1" - # Dlist handle stored internally; accessible via _peer_handles. - assert t._peer_handles["remote-agent-1"] == 88 - - -def test_add_remote_peer_registers_remote_agent(): - nixl = _make_nixl_mock(prep_side_effect=[77, 88]) - t = _make_transport(nixl) - mem_desc = serialize_mem_descriptor([(0, 64, 0)]) - t.add_remote_peer(b"fresh-meta", mem_desc) - nixl.add_remote_agent.assert_called_once_with(b"fresh-meta") - - -# ── remove_remote_peer ──────────────────────────────────────────────────────── - - -def test_remove_remote_peer_delegates_to_nixl(): - t = _make_transport() - t.remove_remote_peer("agent-x") - t._nixl.remove_remote_agent.assert_called_once_with("agent-x") - - -def test_remove_remote_peer_swallows_exceptions(): - t = _make_transport() - t._nixl.remove_remote_agent.side_effect = RuntimeError("gone") - t.remove_remote_peer("agent-x") # must not raise - - -# ── post_read ───────────────────────────────────────────────────────────────── - - -def test_post_read_raises_on_index_count_mismatch(): - nixl = _make_nixl_mock(prep_side_effect=[77, 88]) - t = _make_transport(nixl) - mem_desc = serialize_mem_descriptor([(0, 64, 0)]) - agent = t.add_remote_peer(b"meta", mem_desc) - with pytest.raises(ValueError, match="mismatch"): - t.post_read([0, 1], agent, [5], notif_msg=b"sid") - - -def test_post_read_calls_make_prepped_xfer_and_transfer(): - nixl = _make_nixl_mock(prep_side_effect=[77, 88]) - t = _make_transport(nixl) - mem_desc = serialize_mem_descriptor([(0, 64, 0)]) - agent = t.add_remote_peer(b"meta", mem_desc) - nixl.make_prepped_xfer.return_value = 999 - handle = t.post_read([0, 1], agent, [5, 6], notif_msg=b"sid") - nixl.make_prepped_xfer.assert_called_once() - nixl.transfer.assert_called_once_with(999) - assert handle == 999 - - -# ── release_xfer_handle / deregister ───────────────────────────────────────── - - -def test_release_xfer_handle_swallows_exceptions(): - t = _make_transport() - t._nixl.release_xfer_handle.side_effect = RuntimeError("stale") - t.release_xfer_handle(42) # must not raise - - -def test_deregister_swallows_exceptions(): - t = _make_transport() - t._nixl.deregister_memory.side_effect = RuntimeError("already gone") - t.deregister() # must not raise diff --git a/tests/v1/ec_connector/unit/test_ec_connector_metrics.py b/tests/v1/ec_connector/unit/test_ec_connector_metrics.py deleted file mode 100644 index 239f0267f0..0000000000 --- a/tests/v1/ec_connector/unit/test_ec_connector_metrics.py +++ /dev/null @@ -1,85 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for the connector-agnostic EC connector metrics plumbing. - -Covers only the base infra: no concrete connector type implements these -hooks yet, so every assertion here is about the documented no-op defaults -and the container contracts, mirroring the equivalent KV connector tests. -""" - -from unittest.mock import Mock - -import pytest - -from vllm.config import VllmConfig -from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole -from vllm.distributed.ec_transfer.ec_connector.example_connector import ( - ECExampleConnector, -) -from vllm.distributed.ec_transfer.ec_connector.metrics import ( - ECConnectorLogging, - ECConnectorProm, - ECConnectorStats, -) -from vllm.v1.outputs import ECConnectorOutput - -pytestmark = pytest.mark.cpu_test - - -@pytest.fixture -def example_connector(tmp_path): - config = Mock(spec=VllmConfig) - config.ec_transfer_config = Mock() - config.ec_transfer_config.get_from_extra_config = Mock(return_value=str(tmp_path)) - config.ec_transfer_config.is_ec_producer = True - config.ec_transfer_config.is_ec_consumer = False - return ECExampleConnector(vllm_config=config, role=ECConnectorRole.SCHEDULER) - - -def test_base_stats_methods_raise_not_implemented(): - """The base container has no aggregation logic of its own.""" - stats = ECConnectorStats() - - with pytest.raises(NotImplementedError): - stats.reset() - with pytest.raises(NotImplementedError): - stats.aggregate(stats) - with pytest.raises(NotImplementedError): - stats.reduce() - with pytest.raises(NotImplementedError): - stats.is_empty() - - -def test_connector_hooks_default_to_none(example_connector): - """No connector type overrides these yet, so every hook is a no-op.""" - assert example_connector.get_ec_connector_stats() is None - assert ECExampleConnector.build_ec_connector_stats() is None - assert ( - ECExampleConnector.build_prom_metrics( - vllm_config=Mock(spec=VllmConfig), - metric_types={}, - labelnames=[], - per_engine_labelvalues={}, - ) - is None - ) - - -def test_ec_connector_output_is_empty_reflects_stats_field(): - output = ECConnectorOutput() - assert output.is_empty() - - output.ec_connector_stats = ECConnectorStats(data={"x": 1}) - assert not output.is_empty() - - -def test_logging_and_prom_are_no_ops_without_ec_transfer_config(): - """A normal run with no EC connector configured must stay silent.""" - logging = ECConnectorLogging(ec_transfer_config=None) - logging.log() # Should not raise, nothing to log. - - vllm_config = Mock(spec=VllmConfig) - vllm_config.ec_transfer_config = None - prom = ECConnectorProm(vllm_config, labelnames=[], per_engine_labelvalues={}) - assert prom.prom_metrics is None - prom.observe({}) # Should not raise. diff --git a/tests/v1/ec_connector/unit/test_ec_example_connector.py b/tests/v1/ec_connector/unit/test_ec_example_connector.py deleted file mode 100644 index a940fd4fcf..0000000000 --- a/tests/v1/ec_connector/unit/test_ec_example_connector.py +++ /dev/null @@ -1,747 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for ECExampleConnector.""" - -import os -from unittest.mock import Mock, patch - -import pytest -import safetensors -import torch - -from vllm.config import VllmConfig -from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole -from vllm.distributed.ec_transfer.ec_connector.example_connector import ( - ECExampleConnector, - ECExampleConnectorMetadata, - MMMeta, -) -from vllm.distributed.ec_transfer.ec_connector.utils import collect_ec_item_metadata -from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange -from vllm.v1.core.sched.output import SchedulerOutput - - -# ------------------ Mock Classes ------------------ # -class MockRequest: - def __init__(self, request_id, mm_hashes: list[str], token_counts: list[int]): - assert len(mm_hashes) == len(token_counts) - self.request_id = request_id - self._token_counts = token_counts - self.mm_features = [] - for i, mm_hash in enumerate(mm_hashes): - feature = MultiModalFeatureSpec( - data=None, - modality="image", - identifier=mm_hash, - mm_position=PlaceholderRange(offset=0, length=self._token_counts[i]), - ) - self.mm_features.append(feature) - - def get_num_encoder_embeds(self, input_id: int) -> int: - assert input_id < len(self._token_counts) - return self._token_counts[input_id] - - -def test_metadata_preserves_repeated_item_positions(): - request = MockRequest("request", ["A", "B", "A"], [4, 4, 4]) - resolver = Mock() - resolver.fields_for.return_value = {"image_grid_thw"} - for feature in request.mm_features: - feature.data = Mock() - feature.data.get_data.return_value = {"image_grid_thw": torch.tensor([1, 2, 2])} - items = collect_ec_item_metadata(request.mm_features, resolver) - assert items == { - "A": {"metadata": {"image_grid_thw": [1, 2, 2]}, "item_indices": [0, 2]}, - "B": {"metadata": {"image_grid_thw": [1, 2, 2]}, "item_indices": [1]}, - } - - -def test_audio_metadata_survives_processor_cache_hit(): - resolver = Mock() - resolver.fields_for.return_value = {"audio_num_tokens"} - feature = MultiModalFeatureSpec( - data=None, - modality="audio", - identifier="audio", - mm_position=PlaceholderRange( - offset=0, length=4, is_embed=torch.tensor([False, True, True, False]) - ), - ) - assert collect_ec_item_metadata([feature], resolver) == { - "audio": {"metadata": {"audio_num_tokens": [2]}, "item_indices": [0]} - } - - -@pytest.fixture -def temp_storage(tmp_path): - """Fixture providing temporary storage path.""" - return str(tmp_path) - - -@pytest.fixture -def mock_vllm_config_producer(temp_storage): - """Fixture providing mock VllmConfig for producer role.""" - config = Mock(spec=VllmConfig) - config.ec_transfer_config = Mock() - config.ec_transfer_config.get_from_extra_config = Mock(return_value=temp_storage) - config.ec_transfer_config.is_ec_producer = True - return config - - -@pytest.fixture -def mock_vllm_config_consumer(temp_storage): - """Fixture providing mock VllmConfig for consumer role.""" - config = Mock(spec=VllmConfig) - config.ec_transfer_config = Mock() - config.ec_transfer_config.get_from_extra_config = Mock(return_value=temp_storage) - config.ec_transfer_config.is_ec_producer = False - return config - - -@pytest.fixture -def mock_request_with_3_mm(): - """Fixture providing mock Request with 3 multimodal items.""" - request_id = "test_req_123" - mm_hashes = ["img_hash_1", "img_hash_2", "img_hash_3"] - token_counts = [100, 150, 200] - - request = MockRequest(request_id, mm_hashes, token_counts) - return request - - -# ------------------ Unit Tests ------------------ # -class TestECExampleConnectorBasics: - """Test basic EC connector functionality.""" - - def test_initialization_producer(self, mock_vllm_config_producer, temp_storage): - """Test connector initializes correctly as producer.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.SCHEDULER, - ) - - assert connector.role == ECConnectorRole.SCHEDULER - assert connector.is_producer - assert connector._storage_path == temp_storage - assert connector._mm_datas_need_loads == {} - - def test_initialization_consumer(self, mock_vllm_config_consumer, temp_storage): - """Test connector initializes correctly as consumer.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.WORKER, - ) - - assert connector.role == ECConnectorRole.WORKER - assert not connector.is_producer - assert connector._storage_path == temp_storage - - def test_role_assignment(self, mock_vllm_config_producer): - """Test role is correctly assigned.""" - scheduler_connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.SCHEDULER, - ) - worker_connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.WORKER, - ) - - assert scheduler_connector.role == ECConnectorRole.SCHEDULER - assert worker_connector.role == ECConnectorRole.WORKER - - -class TestCacheExistence: - """Test cache existence checking using has_cache_item() API.""" - - def test_has_cache_item_all_exist_3_items( - self, - mock_vllm_config_producer, - mock_vllm_config_consumer, - mock_request_with_3_mm, - ): - """Test has_cache_item returns True when all 3 caches exist.""" - # Test for producer first - producer = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.SCHEDULER, - ) - - # Create cache files using save_caches (proper way) - encoder_cache: dict[str, torch.Tensor] = {} - - for mm_feature in mock_request_with_3_mm.mm_features: - mm_hash = mm_feature.identifier - encoder_cache[mm_hash] = torch.randn(10, 768) - producer.save_caches(encoder_cache, mm_hash) - - # Test using has_cache_item API - producer_result = [ - producer.has_cache_item(mm_feature.identifier) - for mm_feature in mock_request_with_3_mm.mm_features - ] - - # Assert - assert len(producer_result) == 3 - assert all(producer_result), f"Expected all True, got {producer_result}" - - # Also test consumer can check if cache exists - consumer = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.SCHEDULER, - ) - - # Test using has_cache_item API - consumer_result = [ - consumer.has_cache_item(mm_feature.identifier) - for mm_feature in mock_request_with_3_mm.mm_features - ] - - # Assert - assert len(consumer_result) == 3 - assert all(consumer_result), f"Expected all True, got {consumer_result}" - - def test_has_cache_item_none_exist( - self, mock_vllm_config_producer, mock_request_with_3_mm - ): - """Test has_caches returns False when no caches exist.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.SCHEDULER, - ) - - # Test without creating any files - result = [ - connector.has_cache_item(mm_feature.identifier) - for mm_feature in mock_request_with_3_mm.mm_features - ] - - # Assert - assert len(result) == 3 - assert not any(result), f"Expected all False, got {result}" - - def test_has_cache_item_partial_exist( - self, mock_vllm_config_producer, mock_request_with_3_mm - ): - """Test has_caches with some caches existing (1 of 3).""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.SCHEDULER, - ) - - # Create only the second cache file - mm_hash_second = mock_request_with_3_mm.mm_features[1].identifier - encoder_cache = {mm_hash_second: torch.randn(10, 768)} - connector.save_caches(encoder_cache, mm_hash_second) - - # Test - result = [ - connector.has_cache_item(mm_feature.identifier) - for mm_feature in mock_request_with_3_mm.mm_features - ] - - # Assert - assert len(result) == 3 - assert not result[0] # First doesn't exist - assert result[1] # Second exists - assert not result[2] # Third doesn't exist - - def test_has_cache_item_miss_creates_no_folder( - self, mock_vllm_config_consumer, temp_storage - ): - """Test a cache miss leaves no empty folder in the shared storage.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.SCHEDULER, - ) - - assert not connector.has_cache_item("missing_hash") - assert os.listdir(temp_storage) == [] - - -class TestStateManagement: - """Test connector state management.""" - - def test_update_state_after_alloc_3_items( - self, mock_vllm_config_producer, mock_request_with_3_mm - ): - """Test state update after allocation for 3 MM items.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.SCHEDULER, - ) - - # Initial state should be empty - assert len(connector._mm_datas_need_loads) == 0 - - # Update state for all 3 items (mock cache existence) - with patch.object(connector, "has_cache_item", return_value=True): - for i in range(3): - connector.update_state_after_alloc(mock_request_with_3_mm, index=i) - - # Check state updated for all 3 - assert len(connector._mm_datas_need_loads) == 3 - assert "img_hash_1" in connector._mm_datas_need_loads - assert "img_hash_2" in connector._mm_datas_need_loads - assert "img_hash_3" in connector._mm_datas_need_loads - assert connector._mm_datas_need_loads["img_hash_1"] == 100 - assert connector._mm_datas_need_loads["img_hash_2"] == 150 - assert connector._mm_datas_need_loads["img_hash_3"] == 200 - - def test_build_connector_meta_3_items( - self, mock_vllm_config_producer, mock_request_with_3_mm - ): - """Test metadata building for 3 MM items.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.SCHEDULER, - ) - - # Setup state for all 3 items (mock cache existence) - with patch.object(connector, "has_cache_item", return_value=True): - for i in range(3): - connector.update_state_after_alloc(mock_request_with_3_mm, index=i) - - # Build metadata - scheduler_output = Mock(spec=SchedulerOutput) - metadata = connector.build_connector_meta(scheduler_output) - - # Assert - assert isinstance(metadata, ECExampleConnectorMetadata) - assert len(metadata.mm_datas) == 3 - assert metadata.mm_datas[0].mm_hash == "img_hash_1" - assert metadata.mm_datas[0].num_token == 100 - assert metadata.mm_datas[1].mm_hash == "img_hash_2" - assert metadata.mm_datas[1].num_token == 150 - assert metadata.mm_datas[2].mm_hash == "img_hash_3" - assert metadata.mm_datas[2].num_token == 200 - - # State should be cleared after building - assert len(connector._mm_datas_need_loads) == 0 - - def test_build_connector_meta_empty(self, mock_vllm_config_producer): - """Test metadata building with empty state.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.SCHEDULER, - ) - - scheduler_output = Mock(spec=SchedulerOutput) - metadata = connector.build_connector_meta(scheduler_output) - - assert isinstance(metadata, ECExampleConnectorMetadata) - assert len(metadata.mm_datas) == 0 - - def test_state_cleared_after_metadata_build( - self, mock_vllm_config_producer, mock_request_with_3_mm - ): - """Test that state is properly cleared after building metadata.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.SCHEDULER, - ) - - # Add state (mock cache existence) - with patch.object(connector, "has_cache_item", return_value=True): - for i in range(3): - connector.update_state_after_alloc(mock_request_with_3_mm, index=i) - assert len(connector._mm_datas_need_loads) == 3 - - # Build metadata (should clear state) - scheduler_output = Mock(spec=SchedulerOutput) - connector.build_connector_meta(scheduler_output) - - # State should be empty - assert len(connector._mm_datas_need_loads) == 0 - - # Build again should return empty metadata - metadata2 = connector.build_connector_meta(scheduler_output) - assert len(metadata2.mm_datas) == 0 - - -class TestCacheSaving: - """Test encoder cache saving (producer only).""" - - def test_save_caches_producer_3_items( - self, mock_vllm_config_producer, mock_request_with_3_mm, temp_storage - ): - """Test cache saving as producer for 3 different MM items.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.WORKER, - ) - - # Create and save 3 different caches - mm_hashes = [f.identifier for f in mock_request_with_3_mm.mm_features] - encoder_cache: dict[str, torch.Tensor] = {} - - for mm_hash in mm_hashes: - encoder_cache[mm_hash] = torch.randn(10, 768) - connector.save_caches(encoder_cache, mm_hash) - - # Verify all files exist using has_cache_item - result = [ - connector.has_cache_item(mm_feature.identifier) - for mm_feature in mock_request_with_3_mm.mm_features - ] - assert all(result), f"Not all caches were saved: {result}" - - # Verify each file's content - for mm_hash in mm_hashes: - filename = connector._generate_filename_debug(mm_hash) - loaded = safetensors.torch.load_file(filename) - assert "ec_cache" in loaded - assert torch.allclose(loaded["ec_cache"], encoder_cache[mm_hash].cpu()) - - def test_save_caches_consumer_skips(self, mock_vllm_config_consumer): - """Test cache saving is skipped for consumer.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.WORKER, - ) - - mm_hash = "test_hash_consumer" - encoder_cache = {mm_hash: torch.randn(10, 768)} - - # Save should not raise but also not create file - connector.save_caches(encoder_cache, mm_hash) - - # Verify file doesn't exist using has_cache_item - result = connector.has_cache_item(mm_hash) - assert not result, "Consumer should not save caches" - - -class TestCacheLoading: - """Test encoder cache loading (consumer).""" - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_start_load_caches_consumer_3_items( - self, - mock_vllm_config_producer, - mock_vllm_config_consumer, - mock_request_with_3_mm, - temp_storage, - ): - """Test consumer loads 3 caches from storage.""" - # First, create producer to save caches - producer = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.WORKER, - ) - - # Producer saves 3 caches - mm_hashes = [f.identifier for f in mock_request_with_3_mm.mm_features] - saved_caches = {} - for mm_hash in mm_hashes: - saved_caches[mm_hash] = torch.randn(10, 768) - producer.save_caches(saved_caches, mm_hash) - - # Now consumer loads - consumer = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.WORKER, - ) - - # Setup metadata for all 3 - metadata = ECExampleConnectorMetadata() - for mm_hash in mm_hashes: - metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100)) - consumer.bind_connector_metadata(metadata) - - # Load - encoder_cache: dict[str, torch.Tensor] = {} - consumer.start_load_caches(encoder_cache=encoder_cache) - - # Verify all 3 loaded - assert len(encoder_cache) == 3 - for mm_hash in mm_hashes: - assert mm_hash in encoder_cache, f"{mm_hash} missing in encoder_cache" - assert encoder_cache[mm_hash].is_cuda, ( - f"{mm_hash} cache is in {encoder_cache[mm_hash].device}" - ) - assert torch.allclose( - encoder_cache[mm_hash].cpu(), saved_caches[mm_hash] - ), f"{mm_hash} cache saved and loaded tesnor are not the same" - - def test_start_load_caches_skip_existing( - self, mock_vllm_config_producer, mock_vllm_config_consumer, temp_storage - ): - """Test cache loading skips already cached items.""" - # Setup: producer saves cache - producer = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.WORKER, - ) - - mm_hash = "existing_hash" - saved_cache = torch.randn(10, 768) - producer.save_caches({mm_hash: saved_cache}, mm_hash) - - # Consumer setup - consumer = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.WORKER, - ) - - metadata = ECExampleConnectorMetadata() - metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100)) - consumer.bind_connector_metadata(metadata) - - # Pre-populate encoder_cache with different value - existing_cache = torch.randn(5, 512) - encoder_cache = {mm_hash: existing_cache} - - # Load (should skip since already exists) - with patch("safetensors.torch.load_file") as mock_load: - consumer.start_load_caches(encoder_cache=encoder_cache) - # Should not call load_file since cache exists - mock_load.assert_not_called() - - # Verify original cache unchanged - assert torch.equal(encoder_cache[mm_hash], existing_cache) - - def test_start_load_caches_empty_metadata(self, mock_vllm_config_consumer): - """Test loading with empty metadata does nothing.""" - consumer = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.WORKER, - ) - - # Setup empty metadata - metadata = ECExampleConnectorMetadata() - consumer.bind_connector_metadata(metadata) - - # Load (should not raise) - encoder_cache: dict[str, torch.Tensor] = {} - consumer.start_load_caches(encoder_cache=encoder_cache) - - # Cache should remain empty - assert len(encoder_cache) == 0 - - def test_start_load_caches_pins_device_to_current_rank( - self, mock_vllm_config_producer, mock_vllm_config_consumer - ): - """Load must pin to the selected device under tensor parallelism.""" - # Produce a real cache file so the load path has something to open. - producer = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.WORKER, - ) - mm_hash = "rank_pinning_hash" - producer.save_caches({mm_hash: torch.zeros(4, 8)}, mm_hash) - - consumer = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.WORKER, - ) - metadata = ECExampleConnectorMetadata() - metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100)) - consumer.bind_connector_metadata(metadata) - - mock_platform = Mock() - mock_platform.device_type = "cuda" - mock_platform.is_cuda_alike.return_value = True - mock_platform.current_device.return_value = 2 - - # Simulate running on TP rank 2 without requiring a multi-GPU host. - captured: dict = {} - - def fake_load_file(filename, device): - captured["device"] = device - return {"ec_cache": torch.zeros(4, 8)} - - with ( - patch("vllm.platforms.current_platform", mock_platform), - patch("safetensors.torch.load_file", side_effect=fake_load_file), - ): - encoder_cache: dict[str, torch.Tensor] = {} - consumer.start_load_caches(encoder_cache=encoder_cache) - - assert captured["device"] == "cuda:2", ( - f"Expected device='cuda:2' (pinned to current rank), " - f"got {captured['device']!r}. A bare 'cuda' string makes " - "safetensors load onto cuda:0, which breaks non-rank-0 TP workers." - ) - - def test_start_load_caches_non_cuda_platform_uses_device_type( - self, mock_vllm_config_producer, mock_vllm_config_consumer - ): - """Non-CUDA platforms pass `device_type` through unchanged (no rank suffix).""" - producer = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.WORKER, - ) - mm_hash = "cpu_device_hash" - producer.save_caches({mm_hash: torch.zeros(2, 4)}, mm_hash) - - consumer = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.WORKER, - ) - metadata = ECExampleConnectorMetadata() - metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100)) - consumer.bind_connector_metadata(metadata) - - mock_platform = Mock() - mock_platform.device_type = "cpu" - mock_platform.is_cuda_alike.return_value = False - - captured: dict = {} - - def fake_load_file(filename, device): - captured["device"] = device - return {"ec_cache": torch.zeros(2, 4)} - - with ( - patch("vllm.platforms.current_platform", mock_platform), - patch("safetensors.torch.load_file", side_effect=fake_load_file), - ): - encoder_cache: dict[str, torch.Tensor] = {} - consumer.start_load_caches(encoder_cache=encoder_cache) - - assert captured["device"] == "cpu" - - -class TestFilenameGeneration: - """Test filename and path generation.""" - - def test_generate_foldername(self, mock_vllm_config_producer, temp_storage): - """Test folder name generation.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.WORKER, - ) - - mm_hash = "test_folder_hash" - folder = connector._generate_foldername_debug(mm_hash) - - assert folder == os.path.join(temp_storage, mm_hash) - assert os.path.isdir(folder) # Should be created - - def test_generate_filename(self, mock_vllm_config_producer, temp_storage): - """Test filename generation.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.WORKER, - ) - - mm_hash = "test_file_hash" - filename = connector._generate_filename_debug(mm_hash) - - expected = os.path.join(temp_storage, mm_hash, "encoder_cache.safetensors") - assert filename == expected - assert os.path.isdir(os.path.dirname(filename)) # Folder created - - def test_generate_filename_consistency(self, mock_vllm_config_producer): - """Test filename generation is consistent.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.WORKER, - ) - - mm_hash = "consistency_hash" - filename1 = connector._generate_filename_debug(mm_hash) - filename2 = connector._generate_filename_debug(mm_hash) - - assert filename1 == filename2 - - -class TestMetadataBindingLifecycle: - """Test metadata binding and clearing lifecycle.""" - - def test_bind_connector_metadata(self, mock_vllm_config_consumer): - """Test binding connector metadata.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.WORKER, - ) - - metadata = ECExampleConnectorMetadata() - metadata.add_mm_data(MMMeta.make_meta("hash_1", 100)) - - connector.bind_connector_metadata(metadata) - - assert connector._connector_metadata is metadata - - def test_clear_connector_metadata(self, mock_vllm_config_consumer): - """Test clearing connector metadata.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.WORKER, - ) - - metadata = ECExampleConnectorMetadata() - connector.bind_connector_metadata(metadata) - - connector.clear_connector_metadata() - - assert connector._connector_metadata is None - - def test_get_connector_metadata(self, mock_vllm_config_consumer): - """Test getting connector metadata.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.WORKER, - ) - - metadata = ECExampleConnectorMetadata() - connector.bind_connector_metadata(metadata) - - retrieved = connector._get_connector_metadata() - - assert retrieved is metadata - - def test_get_connector_metadata_not_set(self, mock_vllm_config_consumer): - """Test getting metadata when not set raises.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.WORKER, - ) - - with pytest.raises(AssertionError): - connector._get_connector_metadata() - - -class TestEdgeCases: - """Test edge cases and error handling.""" - - def test_save_empty_cache(self, mock_vllm_config_producer): - """Test saving empty tensor.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.WORKER, - ) - - mm_hash = "empty_hash" - encoder_cache = {mm_hash: torch.empty(0)} - - # Should not raise - connector.save_caches(encoder_cache, mm_hash) - - def test_load_nonexistent_cache(self, mock_vllm_config_consumer): - """Test loading cache that doesn't exist raises error.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_consumer, - role=ECConnectorRole.WORKER, - ) - - metadata = ECExampleConnectorMetadata() - metadata.add_mm_data(MMMeta.make_meta("nonexistent_hash", 100)) - connector.bind_connector_metadata(metadata) - - encoder_cache: dict[str, torch.Tensor] = {} - - # Should raise FileNotFoundError - with pytest.raises(FileNotFoundError): - connector.start_load_caches(encoder_cache=encoder_cache) - - def test_has_cache_item_empty_request(self, mock_vllm_config_producer): - """Test has_cache_item with a nonexistent identifier.""" - connector = ECExampleConnector( - vllm_config=mock_vllm_config_producer, - role=ECConnectorRole.SCHEDULER, - ) - - result = connector.has_cache_item("nonexistent_hash") - - assert result is False diff --git a/tests/v1/ec_connector/unit/test_ec_output_aggregator.py b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py deleted file mode 100644 index df98ec364f..0000000000 --- a/tests/v1/ec_connector/unit/test_ec_output_aggregator.py +++ /dev/null @@ -1,131 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for ECOutputAggregator.""" - -import pytest - -from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata -from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats -from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator -from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator -from vllm.v1.outputs import ( - EMPTY_MODEL_RUNNER_OUTPUT, - ECConnectorOutput, - KVConnectorOutput, - ModelRunnerOutput, -) - -pytestmark = pytest.mark.cpu_test - - -class FakeWorkerMeta(ECConnectorWorkerMetadata): - """Records merge order. `aggregate` returns a new object, as the base class - declares: an aggregator discarding the return value would lose the merge. - """ - - def __init__(self, saves: list[str]): - self.saves = saves - - def aggregate(self, other: "FakeWorkerMeta") -> "FakeWorkerMeta": - return FakeWorkerMeta(self.saves + other.saves) - - -class FakeConnectorStats(ECConnectorStats): - """Records merge order, mirroring FakeWorkerMeta above.""" - - def __init__(self, saves: list[str]): - self.saves = saves - - def aggregate(self, other: "FakeConnectorStats") -> "FakeConnectorStats": - return FakeConnectorStats(self.saves + other.saves) - - -def _worker_output(ec_output: ECConnectorOutput | None) -> ModelRunnerOutput: - return ModelRunnerOutput( - req_ids=[], req_id_to_index={}, ec_connector_output=ec_output - ) - - -def test_aggregate_folds_every_rank_onto_output_rank(): - """EC work done on any rank reaches the scheduler via output_rank's output. - - The middle rank reports no worker metadata or stats: it must neither seed - nor clobber either accumulator. - """ - outputs = [ - _worker_output( - ECConnectorOutput( - finished_sending={"mm0"}, - ec_connector_stats=FakeConnectorStats(["mm0"]), - ec_connector_worker_meta=FakeWorkerMeta(["mm0"]), - ) - ), - _worker_output(ECConnectorOutput(finished_recving={"mm1"})), - _worker_output( - ECConnectorOutput( - ec_connector_stats=FakeConnectorStats(["mm2"]), - ec_connector_worker_meta=FakeWorkerMeta(["mm2"]), - ) - ), - ] - - result = ECOutputAggregator().aggregate(outputs, output_rank=2) - - assert result is outputs[2] - assert result.ec_connector_output.finished_sending == {"mm0"} - assert result.ec_connector_output.finished_recving == {"mm1"} - assert result.ec_connector_output.ec_connector_stats.saves == ["mm0", "mm2"] - assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0", "mm2"] - - -def test_aggregate_leaves_no_ec_output_when_no_worker_reported(): - """Empty per-worker reports must not reach the scheduler as an empty object.""" - outputs = [_worker_output(ECConnectorOutput()), _worker_output(ECConnectorOutput())] - - result = ECOutputAggregator().aggregate(outputs, output_rank=0) - - assert result is outputs[0] - assert result.ec_connector_output is None - assert ECOutputAggregator().aggregate([None], output_rank=0) is None - - -def test_aggregate_does_not_write_through_the_shared_empty_output(): - """A rank with nothing to report yields the shared empty output singleton. - - Folding another rank's metadata onto it must not write through to the - module-level object, which every later step would then carry. - """ - outputs = [ - _worker_output( - ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm0"])) - ), - EMPTY_MODEL_RUNNER_OUTPUT, - ] - - result = ECOutputAggregator().aggregate(outputs, output_rank=1) - - assert EMPTY_MODEL_RUNNER_OUTPUT.ec_connector_output is None - assert result is not EMPTY_MODEL_RUNNER_OUTPUT - assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0"] - - -def test_chaining_with_kv_aggregator_preserves_both_outputs(): - """MultiprocExecutor chains both aggregators and keeps only the last result, - so each must merge onto the same output_rank output rather than replace it. - """ - outputs = [ - _worker_output(ECConnectorOutput(finished_sending={"mm0"})), - _worker_output(None), - ] - outputs[1].kv_connector_output = KVConnectorOutput(invalid_block_ids={7}) - - result = None - for aggregator in ( - KVOutputAggregator(expected_finished_count=1), - ECOutputAggregator(), - ): - result = aggregator.aggregate(outputs, output_rank=1) - - assert result is outputs[1] - assert result.kv_connector_output.invalid_block_ids == {7} - assert result.ec_connector_output.finished_sending == {"mm0"} diff --git a/tests/v1/ec_connector/unit/test_ec_transfer_params.py b/tests/v1/ec_connector/unit/test_ec_transfer_params.py deleted file mode 100644 index 127a075b3d..0000000000 --- a/tests/v1/ec_connector/unit/test_ec_transfer_params.py +++ /dev/null @@ -1,129 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Sanity tests for ec_transfer_params protocol plumbing. - -No running engine required. -""" - -from unittest.mock import MagicMock - -import pytest - -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, -) -from vllm.outputs import CompletionOutput, RequestOutput -from vllm.sampling_params import SamplingParams -from vllm.v1.request import Request, RequestStatus - -# tests.v1.core.utils transitively imports the example KV connector, which -# requires the compiled CUDA flash-attention extensions; skip cleanly rather -# than erroring at collection when those extensions aren't built. -_core_utils = pytest.importorskip( - "tests.v1.core.utils", - reason="requires compiled CUDA flash-attention extensions", - exc_type=ImportError, -) -create_scheduler = _core_utils.create_scheduler - -EC_PARAMS: dict = {"mm_hash_abc": {"peer_host": "10.0.0.1", "peer_port": 5501}} - - -def test_ec_transfer_params_routed_to_sampling_params_extra_args(): - """ec_transfer_params on the request must land in SamplingParams.extra_args.""" - req = ChatCompletionRequest( - model="test-model", - messages=[{"role": "user", "content": "hi"}], - max_tokens=5, - ec_transfer_params=EC_PARAMS, - ) - sp = req.to_sampling_params(max_tokens=5, default_sampling_params={}) - assert sp.extra_args is not None - assert sp.extra_args.get("ec_transfer_params") == EC_PARAMS - - -def test_request_output_add_propagates_ec_transfer_params(): - """RequestOutput.add() must carry ec_transfer_params forward to the caller.""" - - def _out(ec_params): - return RequestOutput( - request_id="r1", - prompt="p", - prompt_token_ids=[1], - prompt_logprobs=None, - outputs=[ - CompletionOutput( - index=0, - text="", - token_ids=[], - cumulative_logprob=0.0, - logprobs=None, - finish_reason=None, - ) - ], - finished=False, - ec_transfer_params=ec_params, - ) - - accumulated = _out(None) - accumulated.add(_out(EC_PARAMS), aggregate=True) - assert accumulated.ec_transfer_params == EC_PARAMS - - -def test_request_reads_ec_transfer_params_from_extra_args(): - """v1 Request must pull ec_transfer_params out of SamplingParams.extra_args.""" - sp = SamplingParams(extra_args={"ec_transfer_params": EC_PARAMS}) - req = Request( - request_id="r1", - prompt_token_ids=[1, 2, 3], - sampling_params=sp, - pooling_params=None, - ) - assert req.ec_transfer_params == EC_PARAMS - - -def test_free_request_calls_ec_connector_and_surfaces_params(): - """_free_request must call ec_connector.request_finished() and return its params.""" - sp = SamplingParams(max_tokens=1) - sp.update_from_generation_config({}, 50256) - request = Request( - request_id="test-req", - prompt_token_ids=[1, 2, 3], - sampling_params=sp, - pooling_params=None, - client_index=0, - ) - scheduler = create_scheduler(use_ec_connector=True, ec_role="ec_producer") - scheduler.add_request(request) - request.status = RequestStatus.FINISHED_STOPPED - - mock_ec = MagicMock() - mock_ec.request_finished.return_value = (False, EC_PARAMS) - scheduler.ec_connector = mock_ec - - kv_params, ec_params = scheduler._free_request(request) - - mock_ec.request_finished.assert_called_once_with(request) - assert ec_params == EC_PARAMS - assert kv_params is None - - -def test_free_request_without_ec_connector_returns_none(): - """When no EC connector is configured, ec_transfer_params must be None.""" - sp = SamplingParams(max_tokens=1) - sp.update_from_generation_config({}, 50256) - request = Request( - request_id="test-req", - prompt_token_ids=[1, 2, 3], - sampling_params=sp, - pooling_params=None, - client_index=0, - ) - scheduler = create_scheduler(use_ec_connector=True, ec_role="ec_producer") - scheduler.add_request(request) - request.status = RequestStatus.FINISHED_STOPPED - - kv_params, ec_params = scheduler._free_request(request) - - assert ec_params is None - assert kv_params is None diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py deleted file mode 100644 index 8e1a730ca6..0000000000 --- a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py +++ /dev/null @@ -1,133 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Encoder fan-out fairness in the disaggregated EPD proxy. - -Exercises the REAL ``encoder_rr_assignment`` helper loaded from its -``examples/`` path -- no routing logic is re-implemented here, so a future -change to that function is what these tests exercise. - -Regression target: the fan-out cursor used to restart at e_urls[0] on every -incoming request (``e_urls[i % len(e_urls)] for i in range(len(mm_items))``), -so single-item requests -- the common case -- always hit the first encoder -instance and left the rest idle. ``encoder_rr_assignment`` threads a cursor -across calls so consecutive requests keep advancing through every URL. -""" - -import importlib.util -import sys -from collections import Counter -from pathlib import Path - -import pytest - -PROXY_REL = "examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py" - - -def _load_proxy_module(): - path = Path(__file__).parents[4] / PROXY_REL - spec = importlib.util.spec_from_file_location("disagg_epd_proxy_under_test", path) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module - - -@pytest.fixture(scope="module") -def proxy(): - return _load_proxy_module() - - -@pytest.fixture(scope="module") -def assign(proxy): - return proxy.encoder_rr_assignment - - -def _drive(assign, e_urls, counts): - """Feed a sequence of request item-counts through the cursor, as - fanout_encoder_primer does one request at a time.""" - cursor = 0 - all_urls = [] - for count in counts: - urls, cursor = assign(e_urls, cursor, count) - assert len(urls) == count - all_urls.append(urls) - return all_urls - - -@pytest.mark.parametrize("n_urls", [1, 2, 3, 5]) -def test_full_url_space_is_covered_uniformly(assign, n_urls): - e_urls = [f"E{i}" for i in range(n_urls)] - # Three full cycles of single-item requests. - urls = _drive(assign, e_urls, counts=[1] * (n_urls * 3)) - hits = Counter(u for req in urls for u in req) - assert set(hits) == set(e_urls) - assert max(hits.values()) == min(hits.values()) - - -def test_single_item_requests_rotate_through_all_encoders(assign): - # Exact scenario from the bug report: 3 encoders, one image/request. - e_urls = ["E0", "E1", "E2"] - routed = _drive(assign, e_urls, counts=[1] * 6) - assert [urls[0] for urls in routed] == ["E0", "E1", "E2", "E0", "E1", "E2"] - - -def test_cursor_stays_contiguous_across_varying_item_counts(assign): - e_urls = ["E0", "E1", "E2"] - routed = _drive(assign, e_urls, counts=[2, 1, 3, 1]) - assert routed == [ - ["E0", "E1"], - ["E2"], - ["E0", "E1", "E2"], - ["E0"], - ] - - -def test_single_encoder_always_resolves_to_it(assign): - urls, next_cursor = assign(["E0"], 0, 4) - assert urls == ["E0"] * 4 - assert next_cursor == 0 - - -def test_generic_e_p_d_routing_without_mooncake_is_supported(proxy): - proxy.validate_ec_consumer_routing(["http://prefill"], []) - - -def test_mooncake_e_pd_routing_is_supported(proxy): - proxy.validate_ec_consumer_routing([], ["tcp://decode:19019"]) - - -def test_mooncake_independent_prefill_routing_fails_fast(proxy): - with pytest.raises(ValueError, match=r"supports E\+PD only"): - proxy.validate_ec_consumer_routing(["http://prefill"], ["tcp://decode:19019"]) - - -def test_decode_rewrite_preserves_engine_reported_ec_hash(proxy): - request = { - "messages": [ - { - "role": "user", - "content": [{"type": "image_url", "image_url": {"url": "image"}}], - } - ] - } - - rewritten = proxy.rewrite_for_decode( - request, - { - 0: { - "mm_hash": "proxy-uuid", - "ec_mm_hash": "engine-derived-hash", - "transfer_id": "transfer", - "image_grid_thw": [1, 2, 3], - } - }, - ) - - assert rewritten["messages"][0]["content"][0]["uuid"] == "proxy-uuid" - assert rewritten["messages"][0]["content"][0]["image_embeds"] == { - "image_grid_thw": [1, 2, 3] - } - assert rewritten["ec_transfer_params"]["ec_items"] == [ - {"mm_hash": "engine-derived-hash", "transfer_id": "transfer"} - ] diff --git a/tests/v1/ec_connector/unit/test_metadata.py b/tests/v1/ec_connector/unit/test_metadata.py deleted file mode 100644 index 65fdc6bbb3..0000000000 --- a/tests/v1/ec_connector/unit/test_metadata.py +++ /dev/null @@ -1,96 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for ECCPUConnectorMetadata.""" - -from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata -from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( - ECCPUConnectorMetadata, -) - - -def test_metadata_is_ec_connector_metadata(): - """ECCPUConnectorMetadata must be a subclass of ECConnectorMetadata.""" - meta = ECCPUConnectorMetadata() - assert isinstance(meta, ECConnectorMetadata) - - -def test_metadata_initialization(): - """ECCPUConnectorMetadata initializes with empty saves and loads dicts.""" - meta = ECCPUConnectorMetadata() - assert hasattr(meta, "saves") - assert hasattr(meta, "loads") - assert meta.saves == {} - assert meta.loads == {} - - -def test_metadata_saves_dict_operations(): - """The saves dict supports standard dict operations.""" - meta = ECCPUConnectorMetadata() - - # Add multiple entries - meta.saves["mm_hash_1"] = [0, 1, 2] - meta.saves["mm_hash_2"] = [10, 11] - - assert "mm_hash_1" in meta.saves - assert "mm_hash_2" in meta.saves - assert len(meta.saves) == 2 - assert meta.saves["mm_hash_1"] == [0, 1, 2] - assert meta.saves["mm_hash_2"] == [10, 11] - - -def test_metadata_loads_carry_a_transfer_id_and_blocks(): - """Each load entry is a (transfer_id, block_ids) pair, so the worker can - report the transfer rather than the mm_hash.""" - meta = ECCPUConnectorMetadata() - - meta.loads["mm_hash_1"] = (0, [5, 6, 7]) - meta.loads["mm_hash_2"] = (1, [100]) - - assert len(meta.loads) == 2 - assert meta.loads["mm_hash_1"] == (0, [5, 6, 7]) - transfer_id, block_ids = meta.loads["mm_hash_2"] - assert transfer_id == 1 - assert block_ids == [100] - - -def test_metadata_saves_and_loads_are_independent(): - """Saves and loads dicts are independent.""" - meta = ECCPUConnectorMetadata() - - meta.saves["key"] = [1, 2] - meta.loads["key"] = (0, [3, 4]) - - # Same key but different dicts should not interfere - assert meta.saves["key"] == [1, 2] - assert meta.loads["key"] == (0, [3, 4]) - - # Modifying one should not affect the other - meta.saves["other"] = [99] - assert "other" not in meta.loads - - -def test_metadata_block_indices_are_integers(): - """Block indices in saves/loads lists should be integers.""" - meta = ECCPUConnectorMetadata() - - meta.saves["hash"] = [0, 100, 999, 65535] - meta.loads["hash"] = (0, [1, 2, 3]) - - for idx in meta.saves["hash"]: - assert isinstance(idx, int) - for idx in meta.loads["hash"][1]: - assert isinstance(idx, int) - - -def test_metadata_multiple_instances_are_independent(): - """Different ECCPUConnectorMetadata instances should not share state.""" - meta1 = ECCPUConnectorMetadata() - meta2 = ECCPUConnectorMetadata() - - meta1.saves["x"] = [1] - meta2.saves["y"] = [2] - - assert "x" in meta1.saves - assert "x" not in meta2.saves - assert "y" not in meta1.saves - assert "y" in meta2.saves diff --git a/tests/v1/ec_connector/unit/test_protocol.py b/tests/v1/ec_connector/unit/test_protocol.py deleted file mode 100644 index 928a1acd6c..0000000000 --- a/tests/v1/ec_connector/unit/test_protocol.py +++ /dev/null @@ -1,41 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import msgspec - -from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import ( - EC_CONNECTOR_VERSION, - XferAck, - XferReq, - XferStatus, - compute_ec_compatibility_hash, -) - - -def test_xferreq_roundtrip(): - req = XferReq(mm_hash="h1", compatibility_hash="c1", session_id="s1") - data = msgspec.msgpack.encode(req) - out = msgspec.msgpack.decode(data, type=XferReq) - assert out.mm_hash == "h1" - assert out.compatibility_hash == "c1" - assert out.connector_version == EC_CONNECTOR_VERSION - - -def test_xferack_roundtrip_ok(): - ack = XferAck( - mm_hash="h1", - status=XferStatus.OK, - src_block_indices=[1, 2], - agent_metadata=b"meta", - mem_descriptor=b"desc", - ) - out = msgspec.msgpack.decode(msgspec.msgpack.encode(ack), type=XferAck) - assert out.status == XferStatus.OK - assert out.src_block_indices == [1, 2] - - -def test_compat_hash_deterministic_and_sensitive(): - a = compute_ec_compatibility_hash("0.1", "m", "float16", 64) - b = compute_ec_compatibility_hash("0.1", "m", "float16", 64) - c = compute_ec_compatibility_hash("0.1", "m", "float16", 128) - assert a == b - assert a != c diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py deleted file mode 100644 index 8c9e14563f..0000000000 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py +++ /dev/null @@ -1,461 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import time -import uuid - -import pytest -import torch - -import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod -from tests.v1.ec_connector.unit.utils import create_ec_vllm_config -from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( - ECSharedRegion, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler -from vllm.multimodal.inputs import ( - MultiModalFieldElem, - MultiModalKwargsItem, - MultiModalSharedField, -) - -_N, _BS, _HID, _ES = 16, 64, 32, 2 - - -class _Pos: - def __init__(self, offset, length): - self.offset, self.length = offset, length - - def get_num_embeds(self): - return self.length - - -class _Feature: - def __init__(self, mm_hash, length=1, data=None): - self.mm_hash = mm_hash - self.identifier = mm_hash - self.mm_position = _Pos(0, length) - self.data = data - self.modality = "image" - - -class _Request: - def __init__(self, features, params=None, req_id="r1"): - self.mm_features = features - self.request_id = req_id - self.ec_transfer_params = params - - -class _FakeResults: - def __init__(self): - self.completed = set() - self.tombstoned = set() - self.quarantined = set() - self.retryable = set() - self.settled = [] - - -class _FakeSession: - """Records started transfers; lets the test flip them to completed.""" - - def __init__(self): - self.started = [] - self._results = _FakeResults() - - def start_xfer(self, mm_hash, block_indices, deadline): - self.started.append(mm_hash) - - def poll(self, messages, now): - pass - - def take_results(self): - r, self._results = self._results, _FakeResults() - return r - - def close(self): - pass - - -class _FakeTransport: - """Stand-in for ZmqClientTransport so _poll_step runs without ZMQ.""" - - def poll(self): - return {} - - def poll_dead(self): - return [] - - def close(self): - pass - - -_MISSING = object() - - -def _consumer_sched(monkeypatch): - def _region(cfg): - return ECSharedRegion( - engine_id="eng-" + str(uuid.uuid4()), num_blocks=_N, block_size_bytes=_BS - ) - - monkeypatch.setattr(sched_mod, "create_ec_shared_region", _region) - - # ec_enable_nixl defaults to False, so this builds gate-off; the fields - # below flip on the NIXL consumer state directly. - # No measured widths: a consumer takes each shape from the announcement. - s = ECCPUScheduler( - create_ec_vllm_config(ec_role="ec_consumer", encoder_output_widths={}) - ) - # Turn on NIXL consumer state without constructing real transports. - s._nixl_enabled = True - s._transport = _FakeTransport() - s._data = None - s._compat_hash = "c" - s._metadata_resolver._cache["image"] = {"image_grid_thw"} - return s - - -def _params(mm_hash, length): - return { - mm_hash: { - "peer_host": "h", - "peer_port": 1, - "size_bytes": length * _HID * _ES, - "shape": [length, _HID], - } - } - - -def test_new_remote_read_defers_then_completes(monkeypatch): - """The read is sized from the announced shape, not a local width.""" - width = 4 * _HID - s = _consumer_sched(monkeypatch) - fake = _FakeSession() - - # Route _start_xfer to our fake session instead of real ZMQ/NIXL, but - # still reserve a real not-ready cache entry so mark_ready works. - def _fake_start(mm_hash, info, size, shape): - assert size == width * _ES - assert shape == (1, width) - entry = s._cache.alloc(mm_hash, size // _BS, shape) - assert entry is not None - fake.started.append(mm_hash) - return True - - monkeypatch.setattr(s, "_start_xfer", _fake_start) - params = _params("h1", 1) - params["h1"]["size_bytes"] = width * _ES - params["h1"]["shape"] = [1, width] - req = _Request([_Feature("h1", 1)], params=params) - - # Step 1: unseen remote item -> read started, request deferred. - assert s.ensure_cache_available(req, 0) is False - assert "h1" in s._in_flight - assert "h1" in fake.started - s.build_connector_meta(scheduler_output=None) # re-arms _first_in_batch - - # Session reports completion this step. - fake._results.completed.add("h1") - s._sessions[("h", 1)] = fake - - # Step 2: poll drains completion -> entry marked ready, request admitted. - assert s.ensure_cache_available(req, 0) is True - assert "h1" not in s._in_flight - assert "h1" in s._step_completed - entry = s._cache.get("h1") - assert entry is not None and entry.ready - # The scheduler now allocates encoder input; the ready entry loads through - # the same local path as a natively cached encoding. - s.update_state_after_alloc(req, 0) - meta = s.build_connector_meta(scheduler_output=None) - assert "h1" in meta.loads - assert meta.loads["h1"][2] == (1, width) - assert "h1" not in s._step_completed # cleared by promote - - # Step 3: still cached -> admitted directly, no new transfer. - assert s.ensure_cache_available(req, 0) is True - s.shutdown() - - -@pytest.mark.parametrize("params", [None, {"x": 123}]) -def test_no_remote_announcement_is_ready(monkeypatch, params): - s = _consumer_sched(monkeypatch) - assert s.ensure_cache_available(_Request([_Feature("x")], params=params), 0) - assert s.take_unavailable_requests() == set() - s.shutdown() - - -@pytest.mark.parametrize("payload", [None, "metadata", "pixels", "embeds", "shm"]) -def test_failed_read_falls_back_only_with_local_input(monkeypatch, payload): - """A remote miss is fatal only when local model input is unavailable.""" - s = _consumer_sched(monkeypatch) - fake = _FakeSession() - - def _fake_start(mm_hash, info, size, shape): - entry = s._cache.alloc(mm_hash, 1, shape) - assert entry is not None - fake.started.append(mm_hash) - return True - - monkeypatch.setattr(s, "_start_xfer", _fake_start) - fields = { - "metadata": {"image_grid_thw": torch.tensor([1, 2, 2])}, - "pixels": {"pixel_values": torch.ones(1, 3)}, - "embeds": {"image_embeds": torch.ones(1, _HID)}, - "shm": {"address": 123, "monotonic_id": 1}, - } - data = ( - MultiModalKwargsItem( - { - key: MultiModalFieldElem(value, MultiModalSharedField(batch_size=1)) - for key, value in fields[payload].items() - } - ) - if payload is not None - else None - ) - req = _Request([_Feature("h1", 1, data)], params=_params("h1", 1)) - - # Step 1: start the read. - assert s.ensure_cache_available(req, 0) is False - s.build_connector_meta(scheduler_output=None) - - fake._results.tombstoned.add("h1") - s._sessions[("h", 1)] = fake - - can_fallback = payload in ("pixels", "embeds", "shm") - assert s.ensure_cache_available(req, 0) is can_fallback - assert s._cache.get("h1") is None - assert "h1" not in s._in_flight - assert s.take_unavailable_requests() == (set() if can_fallback else {"r1"}) - # Drained by the read, so the scheduler cannot abort it twice. - assert s.take_unavailable_requests() == set() - if can_fallback: - # A later scheduling step must not retry the failed remote source. - s.build_connector_meta(scheduler_output=None) - assert s.ensure_cache_available(req, 0) - assert fake.started == ["h1"] - s.shutdown() - - -@pytest.mark.parametrize("retry", [False, True], ids=["in-flight", "retryable"]) -def test_remote_wait_budget_survives_retries_and_long_steps(monkeypatch, retry): - """Neither retries nor delayed scheduling may restart the request budget.""" - s = _consumer_sched(monkeypatch) - fake = _FakeSession() - s._sessions[("h", 1)] = fake - - def start(mm_hash, info, size, shape): - assert s._cache.alloc(mm_hash, 1, shape) is not None - fake.started.append(mm_hash) - return True - - monkeypatch.setattr(s, "_start_xfer", start) - now = [1000.0] - monkeypatch.setattr(time, "monotonic", lambda: now[0]) - req = _Request([_Feature("h1")], params=_params("h1", 1)) - assert not s.ensure_cache_available(req, 0) - for elapsed in (20, 40, 121): - now[0] = 1000.0 + elapsed - if retry: - fake._results.retryable.add("h1") - s.build_connector_meta(scheduler_output=None) - assert not s.ensure_cache_available(req, 0) - assert s.take_unavailable_requests() == ({"r1"} if elapsed > 60 else set()) - assert len(fake.started) == (3 if retry else 1) - s.shutdown() - - -def test_retryable_read_re_requests_without_admitting(monkeypatch): - """A not-ready NACK defers the request and re-requests the read. - - Unlike a tombstone it never admits the request, because an admitted request - whose media was rewritten away upstream has no embedding and no way to - recompute one. - """ - s = _consumer_sched(monkeypatch) - fake = _FakeSession() - - def _fake_start(mm_hash, info, size, shape): - entry = s._cache.alloc(mm_hash, 1, shape) - assert entry is not None - fake.started.append(mm_hash) - return True - - monkeypatch.setattr(s, "_start_xfer", _fake_start) - req = _Request([_Feature("h1", 1)], params=_params("h1", 1)) - - assert s.ensure_cache_available(req, 0) is False - s.build_connector_meta(scheduler_output=None) - - # Producer's save has not landed: retryable, not tombstoned. - fake._results.retryable.add("h1") - s._sessions[("h", 1)] = fake - - # The entry is released and the read re-issued in the same step, so the - # request stays deferred rather than being admitted unfetched. - assert s.ensure_cache_available(req, 0) is False - assert "h1" in s._in_flight - assert "h1" not in s._tombstones - assert fake.started.count("h1") == 2 - s.shutdown() - - -@pytest.mark.parametrize( - "fields", - [ - {"size_bytes": 999}, - {"size_bytes": None}, - {"size_bytes": "big"}, - {"size_bytes": _MISSING}, - {"shape": _MISSING}, - {"shape": [2, _HID]}, - ], - ids=[ - "size-mismatch", - "null-size", - "text-size", - "no-size", - "no-shape", - "wrong-rows", - ], -) -def test_invalid_announcement_fails_the_request(monkeypatch, fields): - """An unusable announced size or shape fails the request, not the engine.""" - s = _consumer_sched(monkeypatch) - announced = {**_params("h1", 1)["h1"], **fields} - announced = {k: v for k, v in announced.items() if v is not _MISSING} - req = _Request([_Feature("h1", 1)], params={"h1": announced}) - assert s.ensure_cache_available(req, 0) is False - assert s.take_unavailable_requests() == {"r1"} - assert "h1" not in s._in_flight - assert s._cache.get("h1") is None - s.shutdown() - - -def test_deferral_budget_is_per_request(monkeypatch): - """Shared hashes have independent wait budgets and cancellation cleanup.""" - s = _consumer_sched(monkeypatch) - monkeypatch.setattr(s._cache, "alloc", lambda key, n, shape=None: None) - params = _params("h1", 1) - old = _Request([_Feature("h1", 1)], params=params, req_id="old") - new = _Request([_Feature("h1", 1)], params=params, req_id="new") - - now = [1000.0] - monkeypatch.setattr(time, "monotonic", lambda: now[0]) - assert s.ensure_cache_available(old, 0) is False - now[0] += sched_mod._ADMIT_DEFER_TIMEOUT_S - 1 - assert s.ensure_cache_available(new, 0) is False - assert s.take_unavailable_requests() == set() - - now[0] += 2 - # `old` is now past its budget; `new` is not. - assert s.ensure_cache_available(new, 0) is False - assert s.take_unavailable_requests() == set() - assert s.request_finished(new) == (False, None) - assert set(s._deferred_since) == {("old", "h1")} - assert s.ensure_cache_available(old, 0) is False - assert s.take_unavailable_requests() == {"old"} - assert "h1" not in s._in_flight - s.shutdown() - - -def test_metadata_only_entry_admits_without_transfer(monkeypatch): - # A producer that reported placeholder metadata but had no cache entry - # to transfer omits peer_host/peer_port/size_bytes entirely; the consumer - # must fall back to local compute rather than treating that as a size - # mismatch. - s = _consumer_sched(monkeypatch) - metadata_only = {"h1": {"image_grid_thw": [1, 2, 3]}} - req = _Request([_Feature("h1", 1)], params=metadata_only) - assert s.ensure_cache_available(req, 0) is True - assert "h1" not in s._in_flight - assert s._cache.get("h1") is None - s.shutdown() - - -def test_orphan_not_ready_entry_defers_no_realloc(monkeypatch): - # Post-quarantine, post-tombstone-consumed orphan: a not-ready cache - # entry exists for the mm_hash, but the hash is tracked in none of - # _in_flight / _step_completed / _tombstones. Its blocks are held by a - # quarantined/settling DMA and must not be re-allocated. - s = _consumer_sched(monkeypatch) - entry = s._cache.alloc("h1", 1) - assert entry is not None and not entry.ready - assert "h1" not in s._in_flight - assert "h1" not in s._step_completed - assert "h1" not in s._tombstones - - # Spy on alloc: record calls, delegate to the real allocator (which - # asserts on a duplicate key). A re-alloc for "h1" is the bug. - real_alloc = s._cache.alloc - calls: list[str] = [] - - def _spy_alloc(key, n): - calls.append(key) - return real_alloc(key, n) - - monkeypatch.setattr(s._cache, "alloc", _spy_alloc) - req = _Request([_Feature("h1", 1)], params=_params("h1", 1)) - - # Must not raise; the request is deferred while the DMA settles, the - # orphan entry is untouched, and alloc is never called again. - assert s.ensure_cache_available(req, 0) is False - assert calls == [] - assert s._cache.get("h1") is entry - assert not entry.ready - assert s.take_unavailable_requests() == set() - s.shutdown() - - -def test_already_computed_feature_admits_without_transfer(monkeypatch): - # A feature entirely within num_computed_tokens is skipped by admit; no - # cache lookup/alloc or transfer should be triggered for it. - s = _consumer_sched(monkeypatch) - calls: list[str] = [] - real_alloc = s._cache.alloc - - def _spy_alloc(key, n): - calls.append(key) - return real_alloc(key, n) - - monkeypatch.setattr(s._cache, "alloc", _spy_alloc) - # Feature spans [0, 4); with num_computed_tokens=4 it is fully computed. - req = _Request([_Feature("h1", 4)], params=_params("h1", 4)) - - assert s.ensure_cache_available(req, 4) is True - assert "h1" not in s._in_flight - assert calls == [] - s.shutdown() - - -def test_local_ready_entry_admits_without_transfer(monkeypatch): - # A READY entry already present in the cache is a local hit: admit - # returns True and no new transfer is started for that mm_hash. - s = _consumer_sched(monkeypatch) - entry = s._cache.alloc("h1", 1) - assert entry is not None - s._cache.mark_ready("h1") - req = _Request([_Feature("h1", 1)], params=_params("h1", 1)) - - assert s.ensure_cache_available(req, 0) is True - assert "h1" not in s._in_flight - s.shutdown() - - -def test_in_flight_hash_defers_without_second_transfer(monkeypatch): - # A mm_hash already tracked in _in_flight must defer the request rather - # than starting a duplicate transfer. - s = _consumer_sched(monkeypatch) - started: list[str] = [] - - def _spy_start_xfer(mm_hash, info, size, shape): - started.append(mm_hash) - return True - - monkeypatch.setattr(s, "_start_xfer", _spy_start_xfer) - s._in_flight.add("h1") - req = _Request([_Feature("h1", 1)], params=_params("h1", 1)) - - assert s.ensure_cache_available(req, 0) is False - assert started == [] - s.shutdown() diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py deleted file mode 100644 index 8cad9ec65c..0000000000 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py +++ /dev/null @@ -1,146 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import uuid -from pathlib import Path - -import pytest - -import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod -import vllm.distributed.nixl_utils as nixl_utils -from tests.v1.ec_connector.unit.utils import create_ec_vllm_config -from vllm.config.ec_transfer import ECTransferConfig -from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( - ECSharedRegion, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler -from vllm.distributed.nixl_utils import NixlWrapper - -_N = 16 -_BS = 64 - - -def _region() -> ECSharedRegion: - return ECSharedRegion( - engine_id="eng-" + str(uuid.uuid4()), num_blocks=_N, block_size_bytes=_BS - ) - - -def test_ec_enable_nixl_absent_from_extra_config_by_default(): - cfg = ECTransferConfig() - assert cfg.get_from_extra_config("ec_enable_nixl", False) is False - - -def test_ec_enable_nixl_read_from_extra_config(): - cfg = ECTransferConfig( - ec_connector="ECCPUConnector", - ec_role="ec_both", - ec_connector_extra_config={"ec_enable_nixl": True}, - ) - assert cfg.get_from_extra_config("ec_enable_nixl", False) is True - - -def test_gate_off_builds_no_nixl(monkeypatch): - monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: _region()) - # ec_enable_nixl defaults to False. - s = ECCPUScheduler(create_ec_vllm_config(ec_role="ec_both")) - assert s._nixl_enabled is False - assert getattr(s, "_data", None) is None - assert getattr(s, "_producer_session", None) is None - s.shutdown() - - -def test_string_false_in_extra_config_leaves_nixl_off(monkeypatch): - """Extra config is not type coerced, and bool("false") is True.""" - monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: _region()) - cfg = create_ec_vllm_config(ec_role="ec_both") - cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = "false" - s = ECCPUScheduler(cfg) - assert s._nixl_enabled is False - s.shutdown() - - -@pytest.mark.cpu_test -def test_missing_nixl_releases_shared_region(monkeypatch): - """A failed constructor leaves no scheduler for the caller to shut down.""" - # Import before patching to avoid caching the missing-NIXL values. - from vllm.distributed.ec_transfer.ec_connector.cpu.data import nixl # noqa: F401 - - region = _region() - monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region) - monkeypatch.setattr(nixl_utils, "NixlWrapper", None) - monkeypatch.setattr(nixl_utils, "nixl_agent_config", None) - cfg = create_ec_vllm_config(ec_role="ec_both") - cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = True - - try: - with pytest.raises(RuntimeError, match="ec_enable_nixl requires NIXL"): - ECCPUScheduler(cfg) - - assert not Path(region._mmap_path).exists() - assert region._fd is None - assert region._mmap_obj is None - finally: - region.cleanup() - - -@pytest.mark.cpu_test -def test_control_transport_failure_releases_nixl_and_shared_region(monkeypatch): - import vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq as zmq_mod - import vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl as data_mod - - class RegisteredData: - registered = True - - def deregister(self): - self.registered = False - - data = RegisteredData() - error = RuntimeError("control transport failed") - - def fail_control_transport(): - raise error - - region = _region() - monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region) - monkeypatch.setattr(nixl_utils, "NixlWrapper", object()) - monkeypatch.setattr(nixl_utils, "nixl_agent_config", object()) - monkeypatch.setattr(data_mod, "NixlDataTransport", lambda **kwargs: data) - monkeypatch.setattr(zmq_mod, "ZmqClientTransport", fail_control_transport) - cfg = create_ec_vllm_config(ec_role="ec_consumer") - cfg.model_config.hf_config = None - cfg.model_config.get_inputs_embeds_size.return_value = 8 - cfg.model_config.model = "test-model" - cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = True - - try: - with pytest.raises(RuntimeError, match="control transport failed") as exc: - ECCPUScheduler(cfg) - - assert exc.value is error - assert not data.registered - assert not Path(region._mmap_path).exists() - assert region._fd is None - assert region._mmap_obj is None - finally: - region.cleanup() - - -@pytest.mark.skipif(NixlWrapper is None, reason="Requires NIXL package") -def test_gate_on_wires_data_transport_and_producer_session(monkeypatch): - # Port 0 lets the OS pick an ephemeral port so the real ZMQ ROUTER bind - # in ProducerSession.start() cannot collide with another test/process. - monkeypatch.setenv("VLLM_EC_SIDE_CHANNEL_HOST", "127.0.0.1") - monkeypatch.setenv("VLLM_EC_SIDE_CHANNEL_PORT", "0") - monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: _region()) - cfg = create_ec_vllm_config(ec_role="ec_both") - cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = True - s = ECCPUScheduler(cfg) - try: - assert s._nixl_enabled is True - assert s._data is not None - assert s._compat_hash is not None - assert s._producer_session is not None - assert s._peer_host == "127.0.0.1" - assert s._transport is not None # is_ec_consumer is also True - finally: - s.shutdown() diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py deleted file mode 100644 index 167f333f73..0000000000 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py +++ /dev/null @@ -1,116 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import uuid - -import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod -from tests.v1.ec_connector.unit.utils import create_ec_vllm_config -from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( - ECSharedRegion, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler - -_N, _BS = 16, 64 - - -class _Pos: - def __init__(self, offset, length): - self.offset, self.length = offset, length - - def get_num_embeds(self): - return self.length - - -class _Feature: - def __init__(self, mm_hash, length=1): - self.mm_hash = mm_hash - self.identifier = mm_hash - self.mm_position = _Pos(0, length) - self.data = None - self.modality = "image" - - -class _Request: - def __init__(self, features, req_id="r1"): - self.mm_features = features - self.request_id = req_id - - -def _sched_gate_off(monkeypatch): - def _region(cfg): - return ECSharedRegion( - engine_id="eng-" + str(uuid.uuid4()), num_blocks=_N, block_size_bytes=_BS - ) - - monkeypatch.setattr(sched_mod, "create_ec_shared_region", _region) - return ECCPUScheduler(create_ec_vllm_config(ec_role="ec_both")) - - -def test_request_finished_gate_off_returns_none(monkeypatch): - s = _sched_gate_off(monkeypatch) - assert s.request_finished(_Request([_Feature("h1")])) == (False, None) - s.shutdown() - - -def test_request_finished_producer_emits_params(monkeypatch): - s = _sched_gate_off(monkeypatch) - # Simulate NIXL-enabled producer bookkeeping without building real NIXL. - s._nixl_enabled = True - s._peer_host, s._peer_port = "1.2.3.4", 5601 - # feature length=2, hidden_dim=32, element_size=2 -> 128 bytes -> 2 blocks. - entry = s._cache.alloc("h1", 2, (2, 32)) - assert entry is not None - s._cache.mark_ready("h1") - - delay, params = s.request_finished(_Request([_Feature("h1", length=2)])) - assert delay is False - assert entry.evictable - assert params == { - "h1": { - "metadata": {}, - "item_indices": [0], - "peer_host": "1.2.3.4", - "peer_port": 5601, - "size_bytes": 2 * 32 * 2, - "shape": [2, 32], - } - } - s.shutdown() - - -def test_request_finished_announces_not_ready_entry(monkeypatch): - """A save whose GPU->mmap copy hasn't been confirmed complete yet is - still announced: the entry can't be evicted before it's ready, and a read - arriving too early is NACKed NACK_NOT_READY for the consumer to retry — - protecting it is not this method's job.""" - s = _sched_gate_off(monkeypatch) - s._nixl_enabled = True - s._peer_host, s._peer_port = "1.2.3.4", 5601 - s._cache.alloc("h1", 2, (2, 32)) # allocated but not marked ready - - delay, params = s.request_finished(_Request([_Feature("h1", length=2)])) - assert delay is False - assert params == { - "h1": { - "metadata": {}, - "item_indices": [0], - "peer_host": "1.2.3.4", - "peer_port": 5601, - "size_bytes": 2 * 32 * 2, - "shape": [2, 32], - } - } - s.shutdown() - - -def test_request_finished_skips_unallocated_entry(monkeypatch): - s = _sched_gate_off(monkeypatch) - s._nixl_enabled = True - s._peer_host, s._peer_port = "1.2.3.4", 5601 - # No alloc() for "h1" at all — e.g. the cache was full at save time. - - delay, params = s.request_finished(_Request([_Feature("h1", length=2)])) - assert delay is False - # The item's placeholder metadata is still reported even though there's - # no cache entry to transfer (empty "metadata": no fields, no transfer). - assert params == {"h1": {"metadata": {}, "item_indices": [0]}} - s.shutdown() diff --git a/tests/v1/ec_connector/unit/test_session.py b/tests/v1/ec_connector/unit/test_session.py deleted file mode 100644 index 5c25fdf228..0000000000 --- a/tests/v1/ec_connector/unit/test_session.py +++ /dev/null @@ -1,471 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for ProducerXfer, ProducerSession, ConsumerXfer, ConsumerSession.""" - -import time -from unittest.mock import MagicMock - -import pytest - -from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import ( - XferAck, - XferReq, - XferStatus, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import ( - EmbeddingCache, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.session import ( - _CONSUMER_QUARANTINE_TIMEOUT_S, - ConsumerSession, - ConsumerXfer, - ProducerSession, - ProducerXfer, - XferState, -) - -# ── helpers ─────────────────────────────────────────────────────────────────── - - -def _make_data(xfer_state: str = "PROC") -> MagicMock: - data = MagicMock() - data.check_xfer_state.return_value = xfer_state - data.post_read.return_value = 99 - return data - - -def _ok_ack( - mm_hash: str = "h1", - session_id: str = "sess-1", - src_indices: list | None = None, -) -> XferAck: - return XferAck( - mm_hash=mm_hash, - status=XferStatus.OK, - session_id=session_id, - src_block_indices=src_indices or [0, 1], - agent_metadata=b"meta", - mem_descriptor=b"desc", - ) - - -def _xfer( - data: MagicMock | None = None, - deadline: float | None = None, - consumer_session_id: str = "sess-1", -) -> ConsumerXfer: - return ConsumerXfer( - mm_hash="h1", - block_indices=[10, 11], - addr=("host", 1234), - deadline=deadline if deadline is not None else time.monotonic() + 60, - data=data or _make_data(), - consumer_session_id=consumer_session_id, - ) - - -def _started( - data: MagicMock | None = None, deadline: float | None = None -) -> ConsumerXfer: - """ConsumerXfer that has already received an OK XferAck.""" - x = _xfer(data, deadline) - x.handle_ack(_ok_ack(session_id="sess-1", src_indices=[5, 6]), agent_name="agent-1") - return x - - -# ── ProducerXfer ────────────────────────────────────────────────────────────── - - -def test_producer_xfer_not_expired_before_deadline(): - x = ProducerXfer("h1", [0, 1], deadline=time.monotonic() + 60) - assert not x.is_expired() - - -def test_producer_xfer_expired_after_deadline(): - x = ProducerXfer("h1", [0, 1], deadline=time.monotonic() - 1) - assert x.is_expired() - - -# ── ProducerSession ─────────────────────────────────────────────────────────── - - -def _xfer_req(mm_hash: str = "h1", session_id: str = "sess-1") -> XferReq: - return XferReq( - mm_hash=mm_hash, compatibility_hash="hash-abc", session_id=session_id - ) - - -def _make_producer_session(cache: EmbeddingCache | None = None) -> ProducerSession: - data = _make_data() - data.get_agent_metadata.return_value = b"meta" - data.get_mem_descriptor.return_value = b"desc" - return ProducerSession( - transport=MagicMock(), - data=data, - cache=cache or EmbeddingCache(num_blocks=8), - compat_hash="hash-abc", - ) - - -def test_producer_session_grant_pins_ready_entry_and_returns_blocks(): - cache = EmbeddingCache(num_blocks=8) - cache.alloc("h1", 2) - cache.mark_ready("h1") - s = _make_producer_session(cache) - - ack = s._grant_or_nack(_xfer_req(mm_hash="h1")) - - assert ack.status == XferStatus.OK - entry = cache.get("h1") - assert not entry.evictable # pinned by the grant - assert list(ack.src_block_indices) == list(entry.block_ids) - - -def test_producer_session_grant_nacks_missing_entry(): - s = _make_producer_session() - ack = s._grant_or_nack(_xfer_req(mm_hash="nope")) - assert ack.status == XferStatus.NACK_MISSING - - -def test_producer_session_grant_nacks_not_ready_entry_as_retryable(): - """An announced entry whose save is still in flight is distinguished from a - miss, so the consumer retries instead of falling back to local compute.""" - cache = EmbeddingCache(num_blocks=8) - cache.alloc("h1", 2) # not ready - s = _make_producer_session(cache) - ack = s._grant_or_nack(_xfer_req(mm_hash="h1")) - assert ack.status == XferStatus.NACK_NOT_READY - - -def test_producer_session_grant_nacks_version_mismatch(): - cache = EmbeddingCache(num_blocks=8) - cache.alloc("h1", 2) - cache.mark_ready("h1") - s = _make_producer_session(cache) - req = XferReq( - mm_hash="h1", - compatibility_hash="hash-abc", - session_id="sess-1", - connector_version=-1, - ) - - ack = s._grant_or_nack(req) - - assert ack.status == XferStatus.NACK_VERSION - entry = cache.get("h1") - assert entry.evictable # nothing was pinned - - -def test_producer_session_grant_nacks_compat_hash_mismatch(): - cache = EmbeddingCache(num_blocks=8) - cache.alloc("h1", 2) - cache.mark_ready("h1") - s = _make_producer_session(cache) - req = XferReq(mm_hash="h1", compatibility_hash="wrong-hash", session_id="sess-1") - - ack = s._grant_or_nack(req) - - assert ack.status == XferStatus.NACK_INCOMPAT - entry = cache.get("h1") - assert entry.evictable # nothing was pinned - - -def test_producer_session_notif_unpins_via_cache(): - cache = EmbeddingCache(num_blocks=8) - cache.alloc("h1", 2) - cache.mark_ready("h1") - s = _make_producer_session(cache) - s._grant_or_nack(_xfer_req(mm_hash="h1", session_id="sess-1")) - entry = cache.get("h1") - assert not entry.evictable # pinned - - s._data.get_new_notifs.return_value = {"agent": [b"sess-1:h1"]} - s._drain_notifs() - - assert entry.evictable # unpinned via cache - assert s._active_xfers == {} - - -def test_producer_session_sweep_timeouts_unpins_via_cache(): - cache = EmbeddingCache(num_blocks=8) - cache.alloc("h1", 2) - cache.mark_ready("h1") - s = _make_producer_session(cache) - s._grant_or_nack(_xfer_req(mm_hash="h1", session_id="sess-1")) - key = next(iter(s._active_xfers)) - s._active_xfers[key].deadline = time.monotonic() - 1 - - s._sweep_timeouts() - - assert cache.get("h1").evictable - assert s._active_xfers == {} - - -# ── ConsumerXfer.handle_ack ─────────────────────────────────────────────────── - - -def test_consumer_xfer_handle_ack_ok_starts_read(): - data = _make_data() - x = _xfer(data) - assert x.handle_ack(_ok_ack(), agent_name="agent-1") is True - assert x.transfer_handle == 99 - - -def test_consumer_xfer_handle_ack_notif_msg_is_session_id_colon_mm_hash(): - data = _make_data() - x = _xfer(data, consumer_session_id="my-sess") - x.handle_ack(_ok_ack(mm_hash="h1"), agent_name="agent-1") - _, kwargs = data.post_read.call_args - assert kwargs["notif_msg"] == b"my-sess:h1" - - -def test_consumer_xfer_handle_ack_forwards_correct_args(): - data = _make_data() - x = _xfer(data) - x.handle_ack(_ok_ack(src_indices=[5, 6]), agent_name="agent-1") - args, _ = data.post_read.call_args - assert args[0] == [10, 11] # local block_indices - assert args[1] == "agent-1" # agent_name (not a raw handle) - assert args[2] == [5, 6] # remote src_block_indices - - -def test_consumer_xfer_handle_ack_resets_deadline(): - original = time.monotonic() + 2 - x = _xfer(deadline=original) - x.handle_ack(_ok_ack(), agent_name="agent-1") - assert x.deadline > original - - -def test_consumer_xfer_handle_ack_nack_returns_false(): - data = _make_data() - x = _xfer(data) - assert ( - x.handle_ack(XferAck(mm_hash="h1", status=XferStatus.NACK_MISSING), "a") - is False - ) - assert x.transfer_handle is None - data.post_read.assert_not_called() - - -# ── ConsumerXfer.poll — WAITING_ACK ────────────────────────────────────────── - - -def test_consumer_xfer_poll_waiting_ack_before_deadline(): - assert _xfer().poll(time.monotonic()) == XferState.WAITING_ACK - - -def test_consumer_xfer_poll_waiting_ack_timeout(): - assert ( - _xfer(deadline=time.monotonic() - 1).poll(time.monotonic()) - == XferState.ACK_TIMEOUT - ) - - -# ── ConsumerXfer.poll — READING ─────────────────────────────────────────────── - - -def test_consumer_xfer_poll_reading_proc(): - assert _started().poll(time.monotonic()) == XferState.READING - - -def test_consumer_xfer_poll_reading_done_releases_handle(): - data = _make_data(xfer_state="DONE") - x = _started(data) - assert x.poll(time.monotonic()) == XferState.DONE - data.release_xfer_handle.assert_called_once_with(99) - assert x.transfer_handle is None - - -def test_consumer_xfer_poll_reading_unexpected_state_fails(): - data = _make_data(xfer_state="ERR") - x = _started(data) - assert x.poll(time.monotonic()) == XferState.READ_FAILED - data.release_xfer_handle.assert_called_once_with(99) - - -def test_consumer_xfer_poll_reading_exception_fails_gracefully(): - data = _make_data() - data.check_xfer_state.side_effect = RuntimeError("nixl crash") - x = _started(data) - assert x.poll(time.monotonic()) == XferState.READ_FAILED - - -def test_consumer_xfer_poll_timeout_quarantines_without_releasing(): - data = _make_data(xfer_state="PROC") - x = _started(data) - assert x.poll(time.monotonic() + 9999) == XferState.QUARANTINED - data.release_xfer_handle.assert_not_called() - assert x.transfer_handle == 99 - - -# ── ConsumerXfer.poll — QUARANTINED ─────────────────────────────────────────── - - -def test_consumer_xfer_quarantined_proc_stays(): - data = _make_data(xfer_state="PROC") - x = _started(data) - x.poll(time.monotonic() + 9999) # → QUARANTINED - assert x.poll(time.monotonic()) == XferState.QUARANTINED - data.release_xfer_handle.assert_not_called() - - -def test_consumer_xfer_quarantined_terminal_settles(): - data = _make_data(xfer_state="PROC") - x = _started(data) - x.poll(time.monotonic() + 9999) # → QUARANTINED - data.check_xfer_state.return_value = "DONE" - assert x.poll(time.monotonic()) == XferState.SETTLED - data.release_xfer_handle.assert_called_once_with(99) - assert x.transfer_handle is None - - -def test_consumer_xfer_quarantined_exception_settles(): - data = _make_data(xfer_state="PROC") - x = _started(data) - x.poll(time.monotonic() + 9999) - data.check_xfer_state.side_effect = RuntimeError("gone") - assert x.poll(time.monotonic()) == XferState.SETTLED - - -def test_consumer_xfer_quarantined_timeout_evicts_blocks(): - """Quarantine timeout should release handle and return READ_FAILED.""" - data = _make_data(xfer_state="PROC") - x = _started(data) - # First poll transitions to QUARANTINED (deadline extended by 60s) - x.poll(time.monotonic() + 9999) - # Second poll with expired quarantine deadline should evict - expired_time = time.monotonic() + 9999 + _CONSUMER_QUARANTINE_TIMEOUT_S + 1 - assert x.poll(expired_time) == XferState.READ_FAILED - data.release_xfer_handle.assert_called_once_with(99) - assert x.transfer_handle is None - - -# ── ConsumerXfer.cancel / release ──────────────────────────────────────────── - - -def test_consumer_xfer_cancel_valid_in_waiting_ack(): - _xfer().cancel() # must not raise - - -def test_consumer_xfer_cancel_raises_when_read_started(): - with pytest.raises(AssertionError): - _started().cancel() - - -def test_consumer_xfer_release_releases_handle(): - data = _make_data() - x = _started(data) - x.release() - data.release_xfer_handle.assert_called_once_with(99) - - -def test_consumer_xfer_release_noop_when_no_handle(): - data = _make_data() - _xfer(data).release() - data.release_xfer_handle.assert_not_called() - - -# ── ConsumerSession ─────────────────────────────────────────────────────────── - - -def _make_consumer_session(data=None): - zmq_conn = MagicMock() - zmq_conn.recv.return_value = [] - return ConsumerSession( - addr=("host", 1234), - zmq_conn=zmq_conn, - transport=MagicMock(), - data=data or _make_data(), - compat_hash="hash-abc", - ) - - -def test_consumer_session_has_stable_session_id(): - s = _make_consumer_session() - assert isinstance(s._session_id, str) and len(s._session_id) > 0 - assert s._session_id == s._session_id # stable - - -def test_consumer_session_start_xfer_sends_req_with_session_id(): - import msgspec - - s = _make_consumer_session() - s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10) - raw = s._zmq.send.call_args[0][0] - req = msgspec.msgpack.decode(raw, type=XferReq) - assert req.mm_hash == "h1" - assert req.session_id == s._session_id - - -def test_consumer_session_poll_done_goes_to_completed(): - import msgspec - - data = _make_data(xfer_state="PROC") - s = _make_consumer_session(data) - s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10) - - ack = XferAck( - mm_hash="h1", - status=XferStatus.OK, - session_id=s._session_id, - src_block_indices=[5, 6], - agent_metadata=b"meta", - mem_descriptor=b"desc", - ) - data.add_remote_peer.return_value = "agent-1" - raw_ack = msgspec.msgpack.encode(ack) - - data.check_xfer_state.return_value = "DONE" - s.poll([raw_ack], time.monotonic()) - - results = s.take_results() - assert "h1" in results.completed - assert not results.tombstoned - - -def test_consumer_session_poll_nack_goes_to_tombstoned(): - import msgspec - - s = _make_consumer_session() - s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10) - nack = XferAck(mm_hash="h1", status=XferStatus.NACK_MISSING) - s.poll([msgspec.msgpack.encode(nack)], time.monotonic()) - results = s.take_results() - assert "h1" in results.tombstoned - assert not results.completed - - -def test_consumer_session_poll_not_ready_nack_goes_to_retryable(): - """A not-ready NACK must not tombstone: tombstoning admits the request with - no embedding, which is fatal once the media has been rewritten away.""" - import msgspec - - s = _make_consumer_session() - s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10) - nack = XferAck(mm_hash="h1", status=XferStatus.NACK_NOT_READY) - s.poll([msgspec.msgpack.encode(nack)], time.monotonic()) - results = s.take_results() - assert "h1" in results.retryable - assert not results.tombstoned - assert not results.completed - - -def test_consumer_session_ack_timeout_reports_retry_once(): - """A delayed producer ACK is retryable, and its result is drained once.""" - s = _make_consumer_session() - s.start_xfer("h1", [0, 1], deadline=time.monotonic() - 1) - s.poll([], time.monotonic()) - r1 = s.take_results() - assert "h1" in r1.retryable - assert not r1.tombstoned - r2 = s.take_results() - assert not r2.retryable # cleared after first take - - -def test_consumer_session_on_peer_down_cancels_waiting_ack(): - s = _make_consumer_session() - s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 60) - s.on_peer_down() - results = s.take_results() - assert "h1" in results.retryable - assert not results.tombstoned diff --git a/tests/v1/ec_connector/unit/test_utils.py b/tests/v1/ec_connector/unit/test_utils.py deleted file mode 100644 index 92925030bb..0000000000 --- a/tests/v1/ec_connector/unit/test_utils.py +++ /dev/null @@ -1,54 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for ECCPUConnector scheduler utilities.""" - -import msgspec -import pytest - -from vllm.distributed.ec_transfer.ec_connector.cpu.utils import ( - build_block_descs, - deserialize_mem_descriptor, - serialize_mem_descriptor, -) - -# ── build_block_descs ──────────────────────────────────────────────────────── - - -@pytest.mark.parametrize("device_id,expected_dev", [(7, 7), (None, 0)]) -def test_build_block_descs(device_id, expected_dev): - kwargs = dict(base_ptr=1000, num_blocks=4, block_size_bytes=256) - if device_id is not None: - kwargs["device_id"] = device_id - descs = build_block_descs(**kwargs) - assert len(descs) == 4 - for i, (addr, size, dev) in enumerate(descs): - assert addr == 1000 + i * 256 - assert size == 256 - assert dev == expected_dev - - -def test_build_block_descs_zero_blocks_returns_empty(): - assert build_block_descs(base_ptr=100, num_blocks=0, block_size_bytes=64) == [] - - -# ── serialize / deserialize_mem_descriptor ─────────────────────────────────── - - -def test_mem_descriptor_roundtrip(): - descs = [(100, 64, 0), (164, 64, 0), (228, 64, 1)] - assert deserialize_mem_descriptor(serialize_mem_descriptor(descs)) == descs - - -@pytest.mark.parametrize( - "bad_value", - [ - [(1, 2)], # 2-tuple instead of 3 - [("a", "b", "c")], # strings, not ints - ], -) -def test_mem_descriptor_rejects_malformed_payload(bad_value): - """Malformed descriptor lists must fail to decode.""" - encoder = msgspec.msgpack.Encoder() - bad_payload = encoder.encode(bad_value) - with pytest.raises((msgspec.DecodeError, msgspec.ValidationError)): - deserialize_mem_descriptor(bad_payload) diff --git a/tests/v1/ec_connector/unit/test_worker_ec_connector.py b/tests/v1/ec_connector/unit/test_worker_ec_connector.py deleted file mode 100644 index 911810b53a..0000000000 --- a/tests/v1/ec_connector/unit/test_worker_ec_connector.py +++ /dev/null @@ -1,115 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for the V2 GPU model runner's EC connector wrapper.""" - -from types import SimpleNamespace -from unittest.mock import MagicMock, patch - -import pytest - -from vllm.distributed.ec_transfer.ec_connector.base import ( - ECConnectorBase, - ECConnectorMetadata, -) -from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT -from vllm.v1.worker.ec_connector_model_runner_mixin import ( - ECConnectorModelRunnerMixin, -) -from vllm.v1.worker.gpu.ec_connector import NO_OP_EC_CONNECTOR, ActiveECConnector - -pytestmark = pytest.mark.cpu_test - -WORKER_META = object() -CONNECTOR_STATS = object() - - -def _scheduler_output() -> SimpleNamespace: - return SimpleNamespace( - ec_connector_metadata=ECConnectorMetadata(), finished_req_ids=frozenset() - ) - - -def _connector( - encoder_cache: dict | None = None, - is_producer: bool = True, - is_consumer: bool = False, -) -> tuple[ActiveECConnector, MagicMock]: - fake = MagicMock(spec=ECConnectorBase) - fake.is_producer = is_producer - fake.is_consumer = is_consumer - fake.get_finished.return_value = (None, None) - fake.build_connector_worker_meta.return_value = WORKER_META - fake.get_ec_connector_stats.return_value = CONNECTOR_STATS - with patch("vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=fake): - return ActiveECConnector(SimpleNamespace(), encoder_cache or {}), fake - - -@pytest.mark.parametrize( - ("is_producer", "is_consumer"), [(True, False), (True, True), (False, True)] -) -def test_saves_newly_added_caches_for_every_producer(is_producer, is_consumer): - """An ec_both node is also a producer: it must offload what it just computed.""" - encoder_cache = {"mm_old": None} - connector, fake = _connector(encoder_cache, is_producer, is_consumer) - - with connector.maybe_get_output(_scheduler_output()): - encoder_cache["mm_new"] = None - - saved = [call.kwargs["mm_hash"] for call in fake.save_caches.call_args_list] - assert saved == (["mm_new"] if is_producer else []) - assert fake.start_save_caches.called == is_producer - if is_producer: - assert fake.start_save_caches.call_args.kwargs["encoder_cache"] is encoder_cache - assert fake.start_load_caches.called == is_consumer - - -def test_worker_meta_is_reported_on_context_exit(): - """Reported in the finally block, so is_empty() sees it only after the exit.""" - connector, fake = _connector() - - with connector.maybe_get_output(_scheduler_output()) as output: - assert output.ec_connector_worker_meta is None - - assert output.ec_connector_worker_meta is WORKER_META - assert fake.clear_connector_metadata.called - - -def test_connector_stats_are_reported_on_context_exit(): - """Reported in the finally block, alongside the worker meta.""" - connector, fake = _connector() - - with connector.maybe_get_output(_scheduler_output()) as output: - assert output.ec_connector_stats is None - - assert output.ec_connector_stats is CONNECTOR_STATS - - -def test_v1_producer_receives_existing_encoder_cache(): - encoder_cache = {"mm_cached": None} - fake = MagicMock(spec=ECConnectorBase) - fake.is_producer = True - fake.is_consumer = False - fake.get_finished.return_value = (None, None) - - module = "vllm.v1.worker.ec_connector_model_runner_mixin" - with ( - patch(f"{module}.has_ec_transfer", return_value=True), - patch(f"{module}.get_ec_transfer", return_value=fake), - ECConnectorModelRunnerMixin.maybe_get_ec_connector_output( - _scheduler_output(), encoder_cache - ), - ): - pass - - assert fake.start_save_caches.call_args.kwargs["encoder_cache"] is encoder_cache - - -def test_no_forward_reports_without_running_the_model(): - connector, _ = _connector() - - output = connector.no_forward(_scheduler_output()) - - assert output.ec_connector_output.ec_connector_worker_meta is WORKER_META - - empty = NO_OP_EC_CONNECTOR.no_forward(_scheduler_output()) - assert empty is EMPTY_MODEL_RUNNER_OUTPUT diff --git a/tests/v1/ec_connector/unit/utils.py b/tests/v1/ec_connector/unit/utils.py deleted file mode 100644 index 705a8ff7ad..0000000000 --- a/tests/v1/ec_connector/unit/utils.py +++ /dev/null @@ -1,81 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Shared config builders for EC connector unit tests.""" - -import uuid -from unittest.mock import Mock - -import torch - -from vllm.config import ModelConfig, VllmConfig -from vllm.config.ec_transfer import ECRole, ECTransferConfig -from vllm.config.parallel import ParallelConfig - -# The scheduler derives its per-load ack count from the PCP size only under the -# `mp` backend -- which also keeps the world size off the local GPU count, so -# TP > 1 stays expressible on a single-GPU host. -_EXECUTOR_BACKEND = "mp" -_PCP_SIZE = 1 - - -_EMBED_WIDTH = 32 - - -def create_ec_vllm_config( - *, - ec_role: ECRole = "ec_both", - tensor_parallel_size: int = 1, - rank: int = 0, - dtype: torch.dtype = torch.float16, - encoder_output_widths: dict[str, int] | None = None, -) -> Mock: - """Build a `VllmConfig` stand-in for EC connector unit tests. - - `ec_transfer_config` and `parallel_config` are the real config objects, so - role derivation and the rank arithmetic the connector depends on behave as - they do in production. - - `model_config` supplies a dtype and embedding width without resolving an - HF model or inspecting its architecture. - - Args: - ec_role: EC role this instance plays. `is_ec_producer` and - `is_ec_consumer` are derived from it, as in production. - tensor_parallel_size: TP degree. - rank: Global rank of this worker. - dtype: Encoder cache dtype. - encoder_output_widths: Per-modality encoder output width, as a - producer measures it at startup. Defaults to the embedding width. - - Returns: - A `VllmConfig`-specced mock carrying the configs above. - - """ - parallel_config = ParallelConfig( - tensor_parallel_size=tensor_parallel_size, - prefill_context_parallel_size=_PCP_SIZE, - distributed_executor_backend=_EXECUTOR_BACKEND, - rank=rank, - ) - ec_transfer_config = ECTransferConfig( - # The role properties report False unless a connector is named. - ec_connector="ECCPUConnector", - ec_role=ec_role, - engine_id=str(uuid.uuid4()), - ) - ec_transfer_config.mm_encoder_output_widths = ( - dict.fromkeys(("image", "video", "audio"), _EMBED_WIDTH) - if encoder_output_widths is None - else encoder_output_widths - ) - - model_config = Mock(spec=ModelConfig) - model_config.dtype = dtype - model_config.get_inputs_embeds_size.return_value = _EMBED_WIDTH - - vllm_config = Mock(spec=VllmConfig) - vllm_config.ec_transfer_config = ec_transfer_config - vllm_config.parallel_config = parallel_config - vllm_config.model_config = model_config - vllm_config.instance_id = f"ec-test-{uuid.uuid4()}" - return vllm_config diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index a1b141237f..413fc80c4f 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -124,7 +124,6 @@ MoEBackend = Literal[ "triton", "batched_triton", "deep_gemm", - "deep_gemm_mega_moe", "cutlass", "flashinfer_trtllm", "flashinfer_cutlass", @@ -163,7 +162,7 @@ PASSTHROUGH_ALL2ALL_BACKEND: Final = "passthrough" # "Native" here meaning that the backend does not use the fused MoE path. # Currently only relevant for the vLLM provided DeepGEMM MegaMoE., # TODO: migrate to the fused MoE path ultimately. -NATIVE_MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"}) +NATIVE_MEGA_MOE_BACKENDS: frozenset[str] = frozenset() # Every megakernel backend. The FlashInfer variants run through the standard # fused-MoE model path; the quantization method binds the kernel. diff --git a/vllm/distributed/device_communicators/cpu_communicator.py b/vllm/distributed/device_communicators/cpu_communicator.py deleted file mode 100644 index c190c9ded6..0000000000 --- a/vllm/distributed/device_communicators/cpu_communicator.py +++ /dev/null @@ -1,339 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import os -from typing import Any - -import torch -from torch.distributed import ProcessGroup - -from vllm.distributed.utils import pickle -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.platforms.interface import CpuArchEnum - -from .base_device_communicator import DeviceCommunicatorBase - -logger = init_logger(__name__) - - -class CpuCommunicator(DeviceCommunicatorBase): - def __init__( - self, - cpu_group: ProcessGroup, - device: torch.device | None = None, - device_group: ProcessGroup | None = None, - unique_name: str = "", - use_all2all: bool = False, - ): - super().__init__( - cpu_group, device, device_group, unique_name, use_all2all=use_all2all - ) - self.dist_module = torch.distributed - - if ( - ( - current_platform.get_cpu_architecture() == CpuArchEnum.X86 - or current_platform.get_cpu_architecture() == CpuArchEnum.ARM - or current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC - or current_platform.get_cpu_architecture() == CpuArchEnum.S390X - ) - and hasattr(torch.ops._C, "init_shm_manager") - and (unique_name.startswith("tp") or unique_name.startswith("pp")) - and self._all_group_ranks_share_shm_group_name() - ): - self.dist_module = _CPUSHMDistributed(self) - elif unique_name.startswith("tp") or unique_name.startswith("pp"): - logger.info( - "CPU SHM communicator disabled for group %s: ranks do not share " - "the same SHM group name, falling back to torch.distributed.", - unique_name, - ) - - # send/recv tensor_dict is only supported through the SHM communicator backend - self.supports_tensor_dict = isinstance(self.dist_module, _CPUSHMDistributed) - - if self.use_all2all: - if self.all2all_backend not in ( - "naive", - "allgather_reducescatter", - ): # type: ignore[has-type] - logger.warning( - "`%s` all2all manager is not supported on CPU. " - "Falling back to `allgather_reducescatter` manager.", - self.all2all_backend, # type: ignore[has-type] - ) - from .all2all import AgRsAll2AllManager - - self.all2all_manager = AgRsAll2AllManager(self.cpu_group) - logger.info("Using allgather_reducescatter all2all manager.") - - def _all_group_ranks_share_shm_group_name(self) -> bool: - """CPUSHM requires all ranks in this group to agree on one SHM group name. - This is a lightweight consistency check for VLLM_DIST_IDENT/name inputs. - """ - local_name = _CPUSHMDistributed.make_group_name(self) - names: list[str] = [""] * self.world_size - torch.distributed.all_gather_object( - names, - local_name, - group=self.device_group, - ) - return len(set(names)) == 1 - - def all_reduce(self, input_): - self.dist_module.all_reduce(input_, group=self.device_group) - return input_ - - def gather( - self, input_: torch.Tensor, dst: int = 0, dim: int = -1 - ) -> torch.Tensor | None: - """NOTE: We assume that the input tensor is on the same device across - all the ranks. - NOTE: `dst` is the local rank of the destination rank. - """ - world_size = self.world_size - assert -input_.dim() <= dim < input_.dim(), ( - f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" - ) - if dim < 0: - # Convert negative dim to positive. - dim += input_.dim() - - # Allocate output tensor. - if self.rank_in_group == dst: - gather_list = [torch.empty_like(input_) for _ in range(world_size)] - else: - gather_list = None - - # Gather. - self.dist_module.gather( - input_, gather_list, dst=self.ranks[dst], group=self.device_group - ) - - if self.rank_in_group == dst: - output_tensor = torch.cat(gather_list, dim=dim) - else: - output_tensor = None - return output_tensor - - def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: - if dim < 0: - # Convert negative dim to positive. - dim += input_.dim() - input_size = input_.size() - # NOTE: we have to use concat-style all-gather here, - # stack-style all-gather has compatibility issues with - # torch.compile . see https://github.com/pytorch/pytorch/issues/138795 - output_size = (input_size[0] * self.world_size,) + input_size[1:] - # Allocate output tensor. - output_tensor = torch.empty( - output_size, dtype=input_.dtype, device=input_.device - ) - # All-gather. - self.dist_module.all_gather_into_tensor( - output_tensor, input_, group=self.device_group - ) - - # Reshape - output_tensor = output_tensor.reshape((self.world_size,) + input_size) - output_tensor = output_tensor.movedim(0, dim) - output_tensor = output_tensor.reshape( - input_size[:dim] - + (self.world_size * input_size[dim],) - + input_size[dim + 1 :] - ) - return output_tensor - - def send_tensor_dict( - self, - tensor_dict: dict[str, torch.Tensor | Any], - dst: int, - ) -> None: - if not self.supports_tensor_dict: - raise NotImplementedError( - "CpuCommunicator does not support tensor dict fastpath with " - "torch.distributed backend." - ) - return self.dist_module.send_tensor_dict(tensor_dict, dst) - - def recv_tensor_dict( - self, - src: int, - ) -> dict[str, torch.Tensor | Any]: - if not self.supports_tensor_dict: - raise NotImplementedError( - "CpuCommunicator does not support tensor dict fastpath with " - "torch.distributed backend." - ) - return self.dist_module.recv_tensor_dict(src) - - def dispatch_router_logits( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - is_sequence_parallel: bool = False, - extra_tensors: list[torch.Tensor] | None = None, - ) -> ( - tuple[torch.Tensor, torch.Tensor] - | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] - ): - """Dispatch the hidden states and router logits to the appropriate device. - This is a no-op in the base class. - """ - assert self.all2all_manager is not None - return self.all2all_manager.dispatch_router_logits( - hidden_states, - router_logits, - is_sequence_parallel, - extra_tensors, - ) - - def dispatch( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - is_sequence_parallel: bool = False, - extra_tensors: list[torch.Tensor] | None = None, - ) -> ( - tuple[torch.Tensor, torch.Tensor, torch.Tensor] - | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] - ): - """Dispatch the hidden states and topk weights/ids to the appropriate device. - This is a no-op in the base class. - """ - assert self.all2all_manager is not None - return self.all2all_manager.dispatch( - hidden_states, - topk_weights, - topk_ids, - is_sequence_parallel, - extra_tensors=extra_tensors, - ) - - def combine( - self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False - ) -> torch.Tensor: - """Combine the hidden states and router logits from the appropriate device. - This is a no-op in the base class. - """ - assert self.all2all_manager is not None - return self.all2all_manager.combine( - hidden_states, - is_sequence_parallel, - ) - - -class _CPUSHMDistributed: - def __init__(self, communicator: CpuCommunicator): - self.communicator = communicator - - self.group_name = self.make_group_name(communicator) - - self.handle = self._init_cpu_shm() - - @staticmethod - def make_group_name(communicator: CpuCommunicator) -> str: - instance_identifier = os.environ["VLLM_DIST_IDENT"] - unique_name = communicator.unique_name - instance_identifier = f"{instance_identifier}-{unique_name}" - group_ranks = [str(rank) for rank in communicator.ranks] - shm_group_identifier = f"[{'-'.join(group_ranks)}]" - return f"{instance_identifier}-{shm_group_identifier}-cpushm" - - def _init_cpu_shm(self) -> int: - thread_num_tensor = torch.tensor( - [torch.get_num_threads()], - dtype=torch.int64, - ) - torch.distributed.all_reduce( - thread_num_tensor, - op=torch.distributed.ReduceOp.MIN, - group=self.communicator.device_group, - ) - thread_num = thread_num_tensor.item() - - handle = torch.ops._C.init_shm_manager( - self.group_name, - self.communicator.world_size, - self.communicator.rank, - thread_num, - ) - torch.distributed.barrier(self.communicator.device_group) - torch.ops._C.join_shm_manager( - handle, - self.group_name, - ) - torch.distributed.barrier(self.communicator.device_group) - - return handle - - def all_reduce( - self, input: torch.Tensor, group: ProcessGroup | None = None - ) -> None: - torch.ops._C.shm_allreduce(self.handle, input) - - def gather( - self, - input: torch.Tensor, - gather_list: list[torch.Tensor] | None, - dst: int = -1, - group: ProcessGroup | None = None, - ) -> None: - # Note: different from the torch gather, here we use local dst rank. - torch.ops._C.shm_gather( - self.handle, - input, - gather_list, - torch.distributed.get_group_rank(group, dst), - ) - - def all_gather_into_tensor( - self, - output: torch.Tensor, - input: torch.Tensor, - group: ProcessGroup | None = None, - ) -> None: - torch.ops._C.shm_all_gather(self.handle, input, output) - - def send_tensor_dict( - self, - tensor_dict: dict[str, torch.Tensor | Any], - dst: int, - ) -> None: - key_list = list(tensor_dict.keys()) - value_list = list(tensor_dict.values()) - size_list = [] - for v in value_list: - if not isinstance(v, torch.Tensor): - raise RuntimeError("CpuCommunicator only supports sending tensors.") - size_list.append(v.size()) - key_size_tensor = torch.frombuffer( - pickle.dumps([key_list, size_list]), dtype=torch.uint8 - ) - value_list.append(key_size_tensor) - - torch.ops._C.shm_send_tensor_list(self.handle, value_list, dst) - - return None - - def recv_tensor_dict( - self, - src: int, - ) -> dict[str, torch.Tensor | Any]: - tensor_list = torch.ops._C.shm_recv_tensor_list(self.handle, src) - - value_list: list[torch.Tensor] = tensor_list[:-1] - key_size_tensor = tensor_list[-1] - - key_size = pickle.loads(key_size_tensor.numpy().tobytes()) - key_list = key_size[0] - size_list = key_size[1] - assert len(key_list) == len(size_list) - assert len(key_list) == len(value_list) - - tensor_dict: dict[str, torch.Tensor] = {} - for key, size, t in zip(key_list, size_list, value_list): - tensor_dict[key] = t.view(size) - return tensor_dict diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py deleted file mode 100644 index 461cd6a55a..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py +++ /dev/null @@ -1,98 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Shared types for the ECCPUConnector scheduler and worker delegates.""" - -from dataclasses import dataclass, field -from typing import TYPE_CHECKING - -import torch - -from vllm.distributed.ec_transfer.ec_connector.base import ( - ECConnectorMetadata, - ECConnectorWorkerMetadata, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( - ECSharedRegion, -) - -if TYPE_CHECKING: - from vllm.config import VllmConfig - - -@dataclass -class ECCPUConnectorMetadata(ECConnectorMetadata): - """Per-step scheduler → worker payload for the ECCPUConnector. - - Populated by `ECCPUScheduler.build_connector_meta`; consumed by - `ECCPUWorker` via the mixin's `bind_connector_metadata`. - """ - - # Producer role: mm_hashes the scheduler has just allocated CPU - # blocks for this step; the worker's save_caches copies - # encoder_cache[mm_hash] → mmap at these block IDs. - saves: dict[str, list[int]] = field(default_factory=dict) - - # Consumer role: mm_hashes whose bytes are available in the local mmap, - # mapped to (transfer_id, block_ids, shape); the worker's start_load_caches - # copies mmap[block_ids] → GPU encoder_cache and reports the transfer_id - # once its copy lands. The id is minted per dispatch, so two loads of the - # same mm_hash in different steps are distinct and a late report from the - # earlier one cannot release the later one's pin. - loads: dict[str, tuple[int, list[int], tuple[int, int]]] = field( - default_factory=dict - ) - - -@dataclass -class ECCPUWorkerMetadata(ECConnectorWorkerMetadata): - """Per-step worker → scheduler payload for the ECCPUConnector. - - Reports the GPU copies that completed this step: saved mm_hashes become - safe to mark ready, and loaded transfers become safe to unpin once every - participating rank has reported them. Built by - `ECCPUWorker.build_connector_worker_meta`; consumed by - `ECCPUScheduler.update_connector_output`. - - Loads are reported by transfer id rather than mm_hash because every rank - copies the same blocks: `aggregate` concatenates, so an id appears once - per reporting rank and the scheduler can count participants off the list. - """ - - completed_saves: list[str] = field(default_factory=list) - completed_loads: list[int] = field(default_factory=list) - - def aggregate(self, other: ECConnectorWorkerMetadata) -> ECConnectorWorkerMetadata: - assert isinstance(other, ECCPUWorkerMetadata) - self.completed_saves.extend(other.completed_saves) - self.completed_loads.extend(other.completed_loads) - return self - - -def create_ec_shared_region(vllm_config: "VllmConfig") -> ECSharedRegion: - """Build the EC mmap region from `vllm_config`. - - Both `ECCPUScheduler` and `ECCPUWorker` call this to get the same - shared region (same engine_id, same block_size_bytes). - """ - ec_config = vllm_config.ec_transfer_config - assert ec_config is not None, "ec_transfer_config required to build region" - - dp_index = vllm_config.parallel_config.data_parallel_index - engine_id = f"{vllm_config.instance_id}_dp{dp_index}" - - dtype = vllm_config.model_config.dtype - hidden_dim = vllm_config.model_config.get_inputs_embeds_size() - element_size = torch.empty(0, dtype=dtype).element_size() - block_size_bytes = hidden_dim * element_size - - cpu_bytes = ec_config.ec_connector_extra_config.get("ec_cpu_bytes") - if not cpu_bytes: - raise ValueError("ec_cpu_bytes must be specified in ec_connector_extra_config") - cpu_bytes = int(cpu_bytes) - num_blocks = cpu_bytes // block_size_bytes - - return ECSharedRegion( - engine_id=engine_id, - num_blocks=num_blocks, - block_size_bytes=block_size_bytes, - ) diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/connector.py b/vllm/distributed/ec_transfer/ec_connector/cpu/connector.py deleted file mode 100644 index a47688f64b..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/connector.py +++ /dev/null @@ -1,153 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""ECCPUConnector — CPU encoder-cache offloading. - -A role-routed shell: one instance per process. The scheduler delegate owns the -mmap region and offload bookkeeping; the worker delegate owns the mmap view and -GPU copy plumbing. An ec_both instance reuses encoder outputs it has already -offloaded to CPU instead of recomputing them. -""" - -from typing import TYPE_CHECKING, Any - -import torch - -from vllm.distributed.ec_transfer.ec_connector.base import ( - ECConnectorBase, - ECConnectorRole, - ECConnectorWorkerMetadata, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( - ECCPUConnectorMetadata, -) -from vllm.logger import init_logger - -if TYPE_CHECKING: - from vllm.config import VllmConfig - from vllm.v1.core.sched.output import SchedulerOutput - from vllm.v1.outputs import ECConnectorOutput - from vllm.v1.request import Request - -logger = init_logger(__name__) - - -class ECCPUConnector(ECConnectorBase): - """EC connector that offloads encoder cache to a shared CPU mmap region.""" - - def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole) -> None: - super().__init__(vllm_config=vllm_config, role=role) - - if not vllm_config.use_v2_model_runner: - raise ValueError( - "ECCPUConnector requires the V2 model runner. " - "Set VLLM_USE_V2_MODEL_RUNNER=1." - ) - - self.connector_worker = None - self.connector_scheduler = None - - if role == ECConnectorRole.WORKER: - self.connector_worker = self._make_worker(vllm_config) - elif role == ECConnectorRole.SCHEDULER: - self.connector_scheduler = self._make_scheduler(vllm_config) - else: - raise ValueError(f"Unknown ECConnectorRole: {role}") - - # Construction seams. - def _make_worker(self, vllm_config: "VllmConfig"): - # Deferred import: the worker module touches torch/CUDA at import time - # via the region, so keep that cost off the scheduler path. - from vllm.distributed.ec_transfer.ec_connector.cpu.worker import ( - ECCPUWorker, - ) - - return ECCPUWorker(vllm_config) - - def _make_scheduler(self, vllm_config: "VllmConfig"): - from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ( - ECCPUScheduler, - ) - - return ECCPUScheduler(vllm_config) - - # Worker-side forwarders. - def start_load_caches( - self, encoder_cache: dict[str, torch.Tensor], **kwargs - ) -> None: - assert self.connector_worker is not None - metadata = self._get_connector_metadata() - assert isinstance(metadata, ECCPUConnectorMetadata) - self.connector_worker.start_load_caches( - encoder_cache, connector_metadata=metadata - ) - - def save_caches( - self, encoder_cache: dict[str, torch.Tensor], mm_hash: str, **kwargs - ) -> None: - assert self.connector_worker is not None - metadata = self._get_connector_metadata() - assert isinstance(metadata, ECCPUConnectorMetadata) - self.connector_worker.save_caches( - encoder_cache, mm_hash, connector_metadata=metadata - ) - - # Scheduler-side forwarders. - def has_cache_item(self, identifier: str) -> bool: - assert self.connector_scheduler is not None - return self.connector_scheduler.has_cache_item(identifier) - - def ensure_cache_available( - self, request: "Request", num_computed_tokens: int - ) -> bool: - assert self.connector_scheduler is not None - return self.connector_scheduler.ensure_cache_available( - request, num_computed_tokens - ) - - def take_unavailable_requests(self) -> set[str]: - assert self.connector_scheduler is not None - return self.connector_scheduler.take_unavailable_requests() - - def update_state_after_alloc(self, request: "Request", index: int) -> None: - assert self.connector_scheduler is not None - self.connector_scheduler.update_state_after_alloc(request, index) - - def build_connector_meta( - self, scheduler_output: "SchedulerOutput" - ) -> ECCPUConnectorMetadata: - assert self.connector_scheduler is not None - return self.connector_scheduler.build_connector_meta(scheduler_output) - - def request_finished( - self, request: "Request" - ) -> tuple[bool, "dict[str, Any] | None"]: - assert self.connector_scheduler is not None - return self.connector_scheduler.request_finished(request) - - def get_finished( - self, finished_req_ids: set[str] - ) -> tuple[set[str] | None, set[str] | None]: - if self.connector_worker is not None: - self.connector_worker.flush_saves() - return None, None - - def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None: - if self.connector_worker is not None: - return self.connector_worker.build_connector_worker_meta() - return None - - def update_connector_output(self, connector_output: "ECConnectorOutput") -> None: - if self.connector_scheduler is not None: - self.connector_scheduler.update_connector_output(connector_output) - - def has_pending_push_work(self) -> bool: - if self.connector_scheduler is not None: - return self.connector_scheduler.has_pending_push_work() - return False - - # Shared. - def shutdown(self) -> None: - if self.connector_scheduler is not None: - self.connector_scheduler.shutdown() - if self.connector_worker is not None: - self.connector_worker.shutdown() diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/control/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/control/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/control/base.py b/vllm/distributed/ec_transfer/ec_connector/cpu/control/base.py deleted file mode 100644 index 07114f095d..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/control/base.py +++ /dev/null @@ -1,37 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Abstract base classes for the EC control plane. - -ControlConnection wraps a single peer channel (one DEALER socket). -""" - -from abc import ABC, abstractmethod - - -class ControlConnection(ABC): - """Bidirectional message channel to a single remote peer. - - All I/O is non-blocking: send enqueues without blocking; recv drains - all buffered inbound messages since the last call. - """ - - @property - @abstractmethod - def alive(self) -> bool: - """True if the connection is still usable.""" - - @abstractmethod - def send(self, msg: bytes) -> None: - """Enqueue msg for delivery. Must not block.""" - - @abstractmethod - def recv(self) -> list[bytes]: - """Drain and return all buffered inbound messages.""" - - @abstractmethod - def mark_dead(self) -> None: - """Signal that the peer disconnected. Sets alive to False.""" - - @abstractmethod - def close(self) -> None: - """Release all resources. Idempotent.""" diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/control/zmq.py b/vllm/distributed/ec_transfer/ec_connector/cpu/control/zmq.py deleted file mode 100644 index d87bdc7920..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/control/zmq.py +++ /dev/null @@ -1,233 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""ZMQ control-plane implementations for ECCPUConnector. - -ZmqClientConnection — one DEALER socket to a single producer peer. -ZmqClientTransport — pool of ZmqClientConnection objects (consumer side). -ZmqServerTransport — ROUTER socket (producer side). -""" - -import contextlib -from dataclasses import dataclass, field - -import zmq -from zmq.utils.monitor import recv_monitor_message - -from vllm.distributed.ec_transfer.ec_connector.cpu.control.base import ( - ControlConnection, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.utils import PeerAddr -from vllm.logger import init_logger -from vllm.utils.network_utils import make_zmq_path, make_zmq_socket - -logger = init_logger(__name__) - -_HEARTBEAT_IVL_MS = 2000 -_HEARTBEAT_TIMEOUT_MS = 4000 -_HEARTBEAT_TTL_MS = 8000 - - -# ── Client side ─────────────────────────────────────────────────────────────── - - -@dataclass -class ZmqClientConnection(ControlConnection): - """One DEALER socket to one producer peer. - - Owns the socket and its optional monitor socket. NIXL state is NOT - stored here — it belongs in ECSession above this layer. - """ - - dealer: zmq.Socket - monitor: zmq.Socket | None - _alive: bool = field(default=True, init=False) - - @property - def alive(self) -> bool: - return self._alive - - def send(self, msg: bytes) -> None: - self.dealer.send_multipart([b"", msg]) - - def recv(self) -> list[bytes]: - msgs: list[bytes] = [] - while True: - try: - frames = self.dealer.recv_multipart(flags=zmq.NOBLOCK) - except zmq.Again: - break - except Exception: - logger.exception("ec: DEALER recv failed") - break - if len(frames) != 2 or frames[0] != b"": - logger.warning( - "ec: dropped malformed frame envelope (%d frames)", len(frames) - ) - continue - msgs.append(frames[1]) - return msgs - - def mark_dead(self) -> None: - self._alive = False - - def close(self) -> None: - if self.monitor is not None: - with contextlib.suppress(Exception): - self.dealer.disable_monitor() - with contextlib.suppress(Exception): - self.monitor.close(linger=0) - with contextlib.suppress(Exception): - self.dealer.close(linger=0) - self._alive = False - - -class ZmqClientTransport: - """Pool of ZmqClientConnection objects — one per producer peer. - - Creates DEALER sockets lazily on first contact. Does not touch NIXL; - NIXL registration is handled by ECCPUConsumer once an XferAck arrives. - """ - - def __init__(self) -> None: - self._ctx = zmq.Context() - self._connections: dict[PeerAddr, ZmqClientConnection] = {} - - def connect(self, addr: PeerAddr) -> ZmqClientConnection: - """Return the existing connection for addr, or create a new one.""" - existing = self._connections.get(addr) - if existing is not None: - return existing - - host, port = addr - dealer = make_zmq_socket( - ctx=self._ctx, - path=make_zmq_path(scheme="tcp", host=host, port=port), - socket_type=zmq.DEALER, - bind=False, - ) - dealer.setsockopt(zmq.HEARTBEAT_IVL, _HEARTBEAT_IVL_MS) - dealer.setsockopt(zmq.HEARTBEAT_TIMEOUT, _HEARTBEAT_TIMEOUT_MS) - dealer.setsockopt(zmq.HEARTBEAT_TTL, _HEARTBEAT_TTL_MS) - - monitor_addr = f"inproc://ec-peer-mon-{host}-{port}" - dealer.monitor(monitor_addr, zmq.EVENT_DISCONNECTED) - monitor = self._ctx.socket(zmq.PAIR) - monitor.connect(monitor_addr) - - conn = ZmqClientConnection(dealer=dealer, monitor=monitor) - self._connections[addr] = conn - return conn - - def poll(self) -> dict[PeerAddr, list[bytes]]: - """Non-blocking drain of messages from every peer's DEALER.""" - msgs: dict[PeerAddr, list[bytes]] = {} - for addr, conn in list(self._connections.items()): - conn_msgs = conn.recv() - msgs[addr] = conn_msgs - return msgs - - def poll_dead(self) -> list[PeerAddr]: - """Non-blocking sweep of peer monitors; return addresses that died.""" - dead: list[PeerAddr] = [] - for addr, conn in list(self._connections.items()): - if conn.monitor is None: - continue - try: - while True: - evt = recv_monitor_message(conn.monitor, flags=zmq.NOBLOCK) - if evt["event"] == zmq.EVENT_DISCONNECTED: - conn.mark_dead() - dead.append(addr) - break - except zmq.Again: - pass - except Exception: - logger.warning("ec: monitor poll failed for %s", addr, exc_info=True) - return dead - - def remove(self, addr: PeerAddr) -> ZmqClientConnection | None: - """Remove and close the connection for addr, returning it (or None).""" - conn = self._connections.pop(addr, None) - if conn is not None: - conn.close() - return conn - - def close(self) -> None: - for conn in list(self._connections.values()): - conn.close() - self._connections.clear() - try: - self._ctx.destroy(linger=0) - except Exception: - logger.debug("ec: zmq context destroy failed", exc_info=True) - - -# ── Server side ─────────────────────────────────────────────────────────────── - - -class ZmqServerTransport: - """ROUTER socket for the producer side. Byte-level only, driven by - ProducerSession. - - poll(timeout_ms) waits up to timeout_ms for messages, then drains all that - are available; timeout_ms=0 drains whatever has already arrived. send() - routes a reply back to a specific peer by ZMQ identity. - """ - - def __init__(self, host: str, port: int) -> None: - self._ctx = zmq.Context() - self._router = make_zmq_socket( - ctx=self._ctx, - path=make_zmq_path(scheme="tcp", host=host, port=port), - socket_type=zmq.ROUTER, - bind=True, - ) - self._poller = zmq.Poller() - self._poller.register(self._router, zmq.POLLIN) - - def poll(self, timeout_ms: int = 0) -> list[tuple[bytes, bytes]]: - """Wait up to timeout_ms for incoming messages, then drain all pending. - - Returns a list of (identity, payload) pairs. - """ - if timeout_ms > 0: - try: - events = dict(self._poller.poll(timeout=timeout_ms)) - except zmq.ContextTerminated: - return [] - except Exception: - logger.exception("ec: router poller failed") - return [] - if self._router not in events: - return [] - - msgs: list[tuple[bytes, bytes]] = [] - while True: - try: - identity, _, payload = self._router.recv_multipart(flags=zmq.NOBLOCK) - msgs.append((identity, payload)) - except zmq.Again: - break - except zmq.ContextTerminated: - break - except Exception: - logger.exception("ec: router recv failed") - break - return msgs - - def send(self, identity: bytes, payload: bytes) -> None: - """Send a reply payload to the peer identified by ZMQ identity.""" - try: - self._router.send_multipart([identity, b"", payload]) - except Exception: - logger.exception("ec: router send failed") - - def close(self) -> None: - try: - self._router.close(linger=0) - except Exception: - logger.debug("ec: router close failed", exc_info=True) - try: - self._ctx.destroy(linger=0) - except Exception: - logger.debug("ec: zmq context destroy failed", exc_info=True) diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/data/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/data/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/data/base.py b/vllm/distributed/ec_transfer/ec_connector/cpu/data/base.py deleted file mode 100644 index c4250c4ce8..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/data/base.py +++ /dev/null @@ -1,67 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Abstract base class for the EC data-plane transport (NIXL/UCX).""" - -from abc import ABC, abstractmethod -from typing import Any - - -class DataTransport(ABC): - """Interface wrapping the NIXL data-plane for EC block transfers. - - Peer lifecycle: - add_remote_peer() — register a peer and store its dlist handle internally - remove_remote_peer() — deregister and release the stored handle - post_read() — issue a READ using the stored handle for agent_name - - The dlist handle (prepped xfer descriptor list) is an implementation detail - of the transport and is never exposed to callers. - """ - - @abstractmethod - def get_agent_metadata(self) -> bytes: - """Return this agent's serialized NIXL metadata blob.""" - - @abstractmethod - def get_mem_descriptor(self) -> bytes: - """Return the msgpack-encoded block descriptor list for this agent.""" - - @abstractmethod - def add_remote_peer(self, metadata: bytes, mem_descriptor: bytes) -> str: - """Register a remote peer, prep its READ-source dlist, and return agent_name. - - The dlist handle is stored internally and used by post_read(). - """ - - @abstractmethod - def remove_remote_peer(self, agent_name: str) -> None: - """Deregister a previously registered remote peer and release its handle.""" - - @abstractmethod - def post_read( - self, - local_indices: list[int], - agent_name: str, - remote_indices: list[int], - notif_msg: bytes, - ) -> Any: - """Issue a consumer-initiated READ using the stored dlist for agent_name. - - Returns a transfer handle for use with check_xfer_state / release_xfer_handle. - """ - - @abstractmethod - def get_new_notifs(self) -> dict[str, list[bytes]]: - """Drain NIXL completion notifications addressed to this agent.""" - - @abstractmethod - def check_xfer_state(self, handle: Any) -> str: - """Return the NIXL state string: 'DONE', 'PROC', or an error string.""" - - @abstractmethod - def release_xfer_handle(self, handle: Any) -> None: - """Release a transfer handle returned by post_read.""" - - @abstractmethod - def deregister(self) -> None: - """Deregister all memory and tear down the NIXL agent.""" diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/data/nixl.py b/vllm/distributed/ec_transfer/ec_connector/cpu/data/nixl.py deleted file mode 100644 index 48b3514b27..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/data/nixl.py +++ /dev/null @@ -1,162 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""NIXL-backed data-plane transport for ECCPUConnector.""" - -from typing import Any - -from vllm.distributed.ec_transfer.ec_connector.cpu.data.base import ( - DataTransport, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.utils import ( - build_block_descs, - deserialize_mem_descriptor, - serialize_mem_descriptor, -) -from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config -from vllm.logger import init_logger - -logger = init_logger(__name__) -_NIXL_DRAM = "DRAM" - - -class NixlDataTransport(DataTransport): - """Adapts NixlWrapper to the DataTransport interface. - - Dlist handles (prepped xfer descriptor lists for remote peers) are stored - internally keyed by agent_name. Callers never touch raw NIXL handles — - post_read() resolves the handle from agent_name internally. - - NixlWrapper is not thread-safe, and neither is this class. Every call site - runs on the scheduler thread, which serializes the transport by - construction; transfers are async inside NIXL and are drained from that - same thread via get_new_notifs() / check_xfer_state(). - """ - - def __init__( - self, - agent_name: str, - base_ptr: int, - num_blocks: int, - block_size_bytes: int, - total_size_bytes: int, - num_threads: int = 1, - ) -> None: - if NixlWrapper is None or nixl_agent_config is None: - raise RuntimeError( - "ECCPUConnector requires NIXL; " - "install the `nixl` package or set a different ec_connector." - ) - # agent_name → remote dlist handle (from add_remote_peer, used by post_read) - self._peer_handles: dict[str, int] = {} - # None until the local region is registered, so deregister() is safe on - # a transport whose registration failed part-way. - self._reg_descs: Any = None - - self._nixl = NixlWrapper( - agent_name, - nixl_agent_config(num_threads=num_threads, capture_telemetry=True), - ) - self._register_local_memory( - base_ptr, num_blocks, block_size_bytes, total_size_bytes - ) - - def _register_local_memory( - self, - base_ptr: int, - num_blocks: int, - block_size_bytes: int, - total_size_bytes: int, - ) -> None: - """Register the local region with NIXL and prep its transfer dlist. - - Unwinds the registration if a later step fails: __init__ propagates the - error, so the caller never receives an object it could deregister with. - """ - block_descs = build_block_descs(base_ptr, num_blocks, block_size_bytes) - self._reg_descs = self._nixl.get_reg_descs( - [(base_ptr, total_size_bytes, 0, "")], _NIXL_DRAM - ) - self._nixl.register_memory(self._reg_descs, backends=["UCX"]) - try: - xfer_descs = self._nixl.get_xfer_descs(block_descs, _NIXL_DRAM) - self._local_xfer_handle = self._nixl.prep_xfer_dlist( - "NIXL_INIT_AGENT", xfer_descs - ) - self._agent_metadata = self._nixl.get_agent_metadata() - self._mem_descriptor_bytes = serialize_mem_descriptor(block_descs) - except Exception: - self.deregister() - raise - - def get_agent_metadata(self) -> bytes: - return self._agent_metadata - - def get_mem_descriptor(self) -> bytes: - return self._mem_descriptor_bytes - - def add_remote_peer(self, metadata: bytes, mem_descriptor: bytes) -> str: - """Register remote peer, prep its dlist, store handle internally. - - Returns agent_name for use in post_read() and remove_remote_peer(). - """ - agent_name = self._nixl.add_remote_agent(metadata) - remote_blocks = deserialize_mem_descriptor(mem_descriptor) - remote_xfer_descs = self._nixl.get_xfer_descs(remote_blocks, _NIXL_DRAM) - self._peer_handles[agent_name] = self._nixl.prep_xfer_dlist( - agent_name, remote_xfer_descs - ) - return agent_name - - def remove_remote_peer(self, agent_name: str) -> None: - self._peer_handles.pop(agent_name, None) - try: - self._nixl.remove_remote_agent(agent_name) - except Exception: - logger.warning( - "EC: remove_remote_peer failed for %s", agent_name, exc_info=True - ) - - def post_read( - self, - local_indices: list[int], - agent_name: str, - remote_indices: list[int], - notif_msg: bytes, - ) -> Any: - if len(local_indices) != len(remote_indices): - raise ValueError( - f"EC: local/remote block count mismatch " - f"({len(local_indices)} vs {len(remote_indices)})" - ) - remote_handle = self._peer_handles[agent_name] - handle = self._nixl.make_prepped_xfer( - "READ", - self._local_xfer_handle, - local_indices, - remote_handle, - remote_indices, - notif_msg=notif_msg, - ) - self._nixl.transfer(handle) - return handle - - def get_new_notifs(self) -> dict[str, list[bytes]]: - return self._nixl.get_new_notifs() - - def check_xfer_state(self, handle: Any) -> str: - return self._nixl.check_xfer_state(handle) - - def release_xfer_handle(self, handle: Any) -> None: - try: - self._nixl.release_xfer_handle(handle) - except Exception: - logger.warning("EC: release_xfer_handle failed", exc_info=True) - - def deregister(self) -> None: - if self._reg_descs is None: - return - try: - self._nixl.deregister_memory(self._reg_descs) - except Exception: - logger.warning("EC: deregister failed", exc_info=True) - self._reg_descs = None diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py b/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py deleted file mode 100644 index bfc8aa009e..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py +++ /dev/null @@ -1,193 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Lightweight mmap-backed shared memory region for encoder cache (EC) data. - -Modeled after SharedOffloadRegion (vllm/v1/kv_offload/cpu/) but simplified -for EC: flat shared layout, no multi-tensor cursor, no block_size_factor. -""" - -import errno -import mmap -import os -import time -from collections.abc import Callable - -import numpy as np -import torch - -from vllm.logger import init_logger - -logger = init_logger(__name__) - -# MADV_POPULATE_WRITE was added in Linux 5.14 (value 23). -_MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23) - - -def _madvise_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None: - mmap_obj.madvise(_MADV_POPULATE_WRITE, offset, length) - - -def _fallback_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None: - # Touch one byte per page via a read-modify-write so existing bytes are - # preserved — a peer worker may have already written EC data into this - # shared mmap by the time we run on a kernel without MADV_POPULATE_WRITE. - arr = np.frombuffer(mmap_obj, dtype=np.uint8) - arr[offset : offset + length : mmap.PAGESIZE] |= 0 - - -def _get_populate_write_fn( - mmap_obj: mmap.mmap, -) -> Callable[[mmap.mmap, int, int], None]: - """Select the pre-faulting method once for this mmap.""" - try: - _madvise_populate_write(mmap_obj, 0, mmap.PAGESIZE) - except OSError as e: - if e.errno != errno.EINVAL: - raise - logger.warning( - "MADV_POPULATE_WRITE is not supported; falling back to per-page " - "writes for EC mmap pre-population. Startup may be slower." - ) - return _fallback_populate_write - return _madvise_populate_write - - -def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0): - """Spin-wait until the file reaches expected_size (creator truncated it).""" - deadline = time.monotonic() + timeout - while True: - if os.fstat(fd).st_size >= expected_size: - return - if time.monotonic() > deadline: - raise TimeoutError( - f"Timed out waiting for EC mmap file to reach {expected_size} bytes" - ) - time.sleep(0.005) - - -class ECSharedRegion: - """Flat mmap-backed memory region shared across TP workers for - encoder cache blocks. - - Layout: (num_blocks, block_size_bytes) — contiguous, no per-worker - interleaving. All workers map the same file and see identical data. - - File path: /dev/shm/vllm_ec_{engine_id}.mmap - - This class owns only the shared memory substrate (mmap lifecycle, the - `blocks` view, CUDA host registration). Block allocation and eviction - are tracked by `EmbeddingCache` in the scheduler process. - """ - - def __init__( - self, - engine_id: str, - num_blocks: int, - block_size_bytes: int, - ) -> None: - self.num_blocks = num_blocks - self.block_size_bytes = block_size_bytes - - total_size_bytes = num_blocks * block_size_bytes - # Path in /dev/shm (tmpfs); unique per engine instance. - self._mmap_path = f"/dev/shm/vllm_ec_{engine_id}.mmap" - # True for the process that created the file (responsible for unlink). - self._is_creator = False - # True after successful cudaHostRegister (cleanup must unregister). - self._is_pinned = False - - # File descriptor for the shared memory backing file. - try: - self._fd: int | None = os.open( - self._mmap_path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o600 - ) - os.ftruncate(self._fd, total_size_bytes) - self._is_creator = True - logger.info( - "Created EC mmap file %s (%.2f MiB)", - self._mmap_path, - total_size_bytes / (1 << 20), - ) - except FileExistsError: - self._fd = os.open(self._mmap_path, os.O_RDWR) - try: - _wait_for_file_size(self._fd, total_size_bytes) - except Exception: - os.close(self._fd) - self._fd = None - raise - logger.info("Opened existing EC mmap file %s", self._mmap_path) - - # MAP_SHARED mmap over _fd; all processes see the same pages. - self._mmap_obj: mmap.mmap | None = mmap.mmap( - self._fd, - total_size_bytes, - flags=mmap.MAP_SHARED, - prot=mmap.PROT_READ | mmap.PROT_WRITE, - ) - - if self._is_creator: - populate_write_fn = _get_populate_write_fn(self._mmap_obj) - populate_write_fn(self._mmap_obj, 0, total_size_bytes) - - # (num_blocks, block_size_bytes) int8 tensor over the mmap buffer. - self.blocks: torch.Tensor = torch.frombuffer( - memoryview(self._mmap_obj), dtype=torch.int8 - ).view(num_blocks, block_size_bytes) - # Cached for cudaHostRegister/Unregister and pointer math. - self._blocks_ptr: int = self.blocks.data_ptr() - self._blocks_nbytes: int = self.blocks.nbytes - - def pin_memory(self) -> None: - """Register the entire mmap as CUDA pinned memory for fast DMA. - - Each TP worker process owns its own virtual address for the same - shared physical pages, so every process must register independently. - No-op when CUDA is not available or already pinned. - """ - if self._is_pinned or not torch.cuda.is_available(): - return - result = torch.cuda.cudart().cudaHostRegister( - self._blocks_ptr, self._blocks_nbytes, 0 - ) - if result.value != 0: - logger.warning( - "cudaHostRegister failed (code=%d) — " - "transfers will still work but may be slower (unpinned DMA)", - result.value, - ) - else: - logger.debug("cudaHostRegister %.2f MB", self._blocks_nbytes / 1e6) - self._is_pinned = True - - def cleanup(self) -> None: - """Tear down the region. Lifecycle method; no concurrent access.""" - logger.info("Starting ECSharedRegion cleanup...") - if self._is_creator: - try: - os.unlink(self._mmap_path) - logger.info("Removed EC mmap file %s", self._mmap_path) - except Exception: - logger.warning( - "Failed to unlink path %s", self._mmap_path, exc_info=True - ) - self._is_creator = False - if self._is_pinned: - result = torch.cuda.cudart().cudaHostUnregister(self._blocks_ptr) - if result.value != 0: - logger.warning("cudaHostUnregister failed (code=%d)", result) - self._is_pinned = False - if hasattr(self, "blocks"): - del self.blocks - if self._mmap_obj: - try: - self._mmap_obj.close() - except Exception: - logger.warning("Failed to close mmap_obj", exc_info=True) - self._mmap_obj = None - if self._fd is not None: - try: - os.close(self._fd) - except Exception: - logger.warning("Failed to close fd %s", self._fd, exc_info=True) - self._fd = None diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py b/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py deleted file mode 100644 index e20bf48f0b..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py +++ /dev/null @@ -1,91 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Scheduler-side wire types for the ECCPUConnector. - -Two peer-to-peer message types exchanged between consumer and producer over -ZMQ (XferReq, XferAck) plus a peer-compatibility fingerprint -(compute_ec_compatibility_hash) sent on every XferReq so a producer refuses -to serve a mismatched peer. -""" - -import enum -import hashlib -import json - -import msgspec - -# Bump when the wire format or tensor/block layout changes in a -# backward-incompatible way. Peers reject mismatched versions. -EC_CONNECTOR_VERSION: int = 2 - - -class XferStatus(enum.IntEnum): - """Outcome the producer reports for an XferReq.""" - - OK = 0 - NACK_MISSING = 1 # producer no longer holds the encoding - NACK_INCOMPAT = 2 # peer-compatibility hash mismatch - NACK_VERSION = 3 # connector wire-version mismatch - NACK_INTERNAL = 4 # producer hit an unexpected error - NACK_NOT_READY = 5 # announced, but its GPU->mmap save has not landed yet - - -# The producer expects a retryable condition to clear on its own, so the -# consumer re-requests the read rather than giving up on the encoding. -RETRYABLE_NACKS = frozenset((XferStatus.NACK_NOT_READY,)) -# Ordinary outcomes rather than signs of a broken peer. -EXPECTED_NACKS = RETRYABLE_NACKS | {XferStatus.NACK_MISSING} - - -class XferReq( # type: ignore[call-arg] - msgspec.Struct, - tag="req", - omit_defaults=True, -): - """Consumer → Producer: I want to READ mm_hash; pin it and tell me where.""" - - mm_hash: str - compatibility_hash: str - # Consumer's session identity. Together with mm_hash it uniquely identifies - # this transfer on both sides: the producer keys its active xfers by - # (session_id, mm_hash), and the NIXL notif_msg carries the same pair. - session_id: str = "" - connector_version: int = EC_CONNECTOR_VERSION - - -class XferAck( # type: ignore[call-arg] - msgspec.Struct, - tag="ack", - omit_defaults=True, -): - """Producer → Consumer: response to an XferReq. - - On status == OK: producer has pinned source blocks and returns their - indices plus fresh NIXL metadata for the consumer to register and READ. - On any NACK_*: optional fields are empty; consumer falls back to local encode. - """ - - mm_hash: str - status: XferStatus - # Consumer's session_id echoed back from XferReq. Together with mm_hash it - # identifies the exact transfer; the consumer uses both to build notif_msg. - session_id: str = "" - src_block_indices: list[int] = [] - agent_metadata: bytes = b"" - mem_descriptor: bytes = b"" - - -def compute_ec_compatibility_hash( - vllm_version: str, - model: str, - dtype: str, - block_size_bytes: int, -) -> str: - """Peer-compatibility fingerprint over the four factors that determine - byte-layout compatibility. Producer NACKs any XferReq whose hash differs. - """ - factors = json.dumps( - [vllm_version, model, dtype, block_size_bytes], - separators=(",", ":"), - ) - return hashlib.sha256(factors.encode("utf-8")).hexdigest() diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py deleted file mode 100644 index 567f2d59c9..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ /dev/null @@ -1,700 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""ECCPUScheduler — CPU offload scheduler delegate. - -Owns the mmap region and the embedding cache, and handles the producer -(GPU->CPU offload) and consumer (CPU->GPU reload) scheduler-side logic -for the ECCPUConnector. -""" - -import time -from collections.abc import Collection -from typing import TYPE_CHECKING, Any - -import torch - -from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( - ECCPUConnectorMetadata, - ECCPUWorkerMetadata, - create_ec_shared_region, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import ( - EmbeddingCache, -) -from vllm.distributed.ec_transfer.ec_connector.encoder_output_width import ( - get_encoder_output_width, -) -from vllm.distributed.ec_transfer.ec_connector.utils import ( - PlaceholderMetadataResolver, - collect_ec_item_metadata, -) -from vllm.logger import init_logger - -if TYPE_CHECKING: - from vllm.config import VllmConfig - from vllm.multimodal.inputs import MultiModalFeatureSpec - from vllm.v1.core.sched.output import SchedulerOutput - from vllm.v1.outputs import ECConnectorOutput - from vllm.v1.request import Request - -logger = init_logger(__name__) - -# Total remote wait budget per request item, including all retries. -_ADMIT_DEFER_TIMEOUT_S = 60.0 - - -class _RemoteUnavailable(Exception): - """An item cannot be fetched within its request's remote wait budget.""" - - -class ECCPUScheduler: - """Scheduler delegate for the ECCPUConnector.""" - - def __init__(self, vllm_config: "VllmConfig") -> None: - self._vllm_config = vllm_config - ec_config = vllm_config.ec_transfer_config - assert ec_config is not None - self._is_producer: bool = ec_config.is_ec_producer - self._is_consumer: bool = ec_config.is_ec_consumer - - self._region = create_ec_shared_region(vllm_config) - # Block allocator + LRU eviction policy for the shared region. - self._cache = EmbeddingCache(self._region.num_blocks) - self._metadata_resolver = PlaceholderMetadataResolver(vllm_config.model_config) - - # mm_hash → block IDs allocated this step for GPU→mmap saves. - self._pending_saves: dict[str, list[int]] = {} - # mm_hash → (transfer_id, block IDs, shape) to load from mmap→GPU. - self._pending_loads: dict[str, tuple[int, list[int], tuple[int, int]]] = {} - - # Dispatched loads awaiting completion reports, keyed by transfer id: - # transfer_id → (mm_hash, reports still outstanding). The pin taken at - # dispatch is released only once the count reaches zero. - self._load_acks: dict[int, tuple[str, int]] = {} - self._next_transfer_id = 0 - - pc = vllm_config.parallel_config - # Reports to expect per load. Only pipeline stage 0 runs the EC - # connector, so the participants are the tp × pcp ranks of that stage. - # Other executors deliver a single rank's output, and expecting reports - # that cannot arrive would hold the pin forever. - self._expected_load_reports = ( - pc.tensor_parallel_size * pc.prefill_context_parallel_size - if pc.distributed_executor_backend == "mp" - else 1 - ) - - self._ec_config = ec_config - # NIXL p2p is an option of this connector, not engine-wide behavior, so - # it lives in extra config. Extra config is not type coerced, so a value - # that arrived as a JSON/CLI string is parsed, not truth-tested. - raw_nixl = ec_config.get_from_extra_config("ec_enable_nixl", False) - self._nixl_enabled = ( - raw_nixl - if isinstance(raw_nixl, bool) - else str(raw_nixl).strip().lower() in ("true", "1", "yes") - ) - # NIXL fields default to None/empty so the gate-off path is untouched. - self._data: Any = None - self._compat_hash: str | None = None - self._first_in_batch = True - self._transport: Any = None - self._producer_session: Any = None - self._sessions: dict = {} - self._in_flight: set[str] = set() - self._tombstones: set[str] = set() - self._step_completed: set[str] = set() - # Per-request clocks survive retries and are cleared on completion. - self._deferred_since: dict[tuple[str, str], float] = {} - # Requests needing a remote encoding that will not arrive. Drained by - # take_unavailable_requests(); the scheduler aborts them. - self._unrecoverable: set[str] = set() - self._peer_host: str | None = None - self._peer_port: int | None = None - dtype = vllm_config.model_config.dtype - assert isinstance(dtype, torch.dtype) - self._dtype = dtype - self._ack_timeout_s: float = 0.0 - if self._nixl_enabled: - try: - self._setup_nixl(vllm_config) - except Exception: - self.shutdown() - raise - - def _setup_nixl(self, vllm_config: "VllmConfig") -> None: - # Lazy imports keep nixl/zmq off the gate-off path. - from vllm import envs - from vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq import ( - ZmqClientTransport, - ZmqServerTransport, - ) - from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import ( - NixlDataTransport, - ) - from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import ( - compute_ec_compatibility_hash, - ) - from vllm.distributed.ec_transfer.ec_connector.cpu.session import ( - PRODUCER_PIN_LEASE_S, - ProducerSession, - ) - from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config - from vllm.version import __version__ as VLLM_VERSION - - if NixlWrapper is None or nixl_agent_config is None: - raise RuntimeError( - "ec_enable_nixl requires NIXL; install the `nixl` package or " - "remove ec_enable_nixl from ec_connector_extra_config." - ) - engine_id = self._ec_config.engine_id - assert engine_id is not None - # How long a consumer waits for an XferAck. The producer answers from - # its scheduler step, so its reply latency scales with the encoder's - # batch size: a deployment whose steps run longer must raise this. - self._ack_timeout_s = float( - self._ec_config.get_from_extra_config( - "consumer_ack_timeout_s", PRODUCER_PIN_LEASE_S - ) - ) - self._compat_hash = compute_ec_compatibility_hash( - vllm_version=VLLM_VERSION, - model=str(vllm_config.model_config.model), - dtype=str(self._dtype), - block_size_bytes=self._region.block_size_bytes, - ) - if self._is_producer: - self._peer_host = envs.VLLM_EC_SIDE_CHANNEL_HOST - self._peer_port = envs.VLLM_EC_SIDE_CHANNEL_PORT - - self._data = NixlDataTransport( - agent_name=engine_id, - base_ptr=self._region.blocks.data_ptr(), - num_blocks=self._region.num_blocks, - block_size_bytes=self._region.block_size_bytes, - total_size_bytes=self._region.num_blocks * self._region.block_size_bytes, - ) - if self._is_producer: - assert self._peer_host is not None - assert self._peer_port is not None - self._producer_session = ProducerSession( - transport=ZmqServerTransport( - host=self._peer_host, port=self._peer_port - ), - data=self._data, - cache=self._cache, - compat_hash=self._compat_hash, - ) - if self._is_consumer: - self._transport = ZmqClientTransport() - - def has_cache_item(self, identifier: str) -> bool: - if not self._is_consumer: - return False - entry = self._cache.get(identifier) - return entry is not None and entry.ready - - def ensure_cache_available( - self, - request: "Request", - num_computed_tokens: int, - local_cache_hashes: Collection[str] | None = None, - ) -> bool: - if not self._nixl_enabled: - return True # CPU offload never blocks. - first = self._first_in_batch - self._first_in_batch = False - if not self._is_consumer: - return True - if first: - self._poll_step() - return self._nixl_consumer_admit(request, num_computed_tokens) - - def _nixl_consumer_admit( - self, request: "Request", num_computed_tokens: int - ) -> bool: - """Admit a request once its required remote encodings are ready. - - Each item has a per-request wait budget that survives transfer retries. - Unavailable items fall back to local input; without it, the request is - reported by take_unavailable_requests() for the scheduler to abort. - - Returns: - True if no item needs to wait for a remote encoding. False if any - item is pending or the request must be aborted. - - """ - if not request.ec_transfer_params: - return True - now = time.monotonic() - pending = False - for feature in request.mm_features: - try: - if not self._admit_item(request, feature, num_computed_tokens, now): - pending = True - except _RemoteUnavailable as error: - if not self._fail_or_fallback(request, feature, str(error)): - return False - return not pending - - def _admit_item( - self, - request: "Request", - feature: "MultiModalFeatureSpec", - num_computed_tokens: int, - now: float, - ) -> bool: - """Check one encoding, starting a remote read if needed. - - Returns: - True if already computed, cached, or not remotely sourced. False - while waiting for a read or cache space, including across retries. - - Raises: - _RemoteUnavailable: The remote source is invalid, the read failed, - or the item's total wait budget expired. - - """ - pos = feature.mm_position - mm_hash = feature.identifier - if pos.offset + pos.length <= num_computed_tokens: - self._deferred_since.pop((request.request_id, mm_hash), None) - return True - announced = (request.ec_transfer_params or {}).get(mm_hash) - # Without a producer address there is nothing to fetch: the request - # carries what the model needs and the encoder runs locally. - # `ec_transfer_params` reaches us from the request, so its shape is - # checked rather than assumed. - remote: dict[str, Any] | None = ( - announced - if isinstance(announced, dict) and "peer_host" in announced - else None - ) - - entry = self._cache.get(mm_hash) - if entry is not None and entry.ready: - # Local hit: upstream's update_state_after_alloc pins and - # loads it through the same path as a natively cached entry. - self._deferred_since.pop((request.request_id, mm_hash), None) - return True - if remote is None: - return True - since = self._deferred_since.setdefault((request.request_id, mm_hash), now) - if now - since > _ADMIT_DEFER_TIMEOUT_S: - raise _RemoteUnavailable(f"the remote wait exceeded {now - since:.0f}s") - if mm_hash in self._in_flight or mm_hash in self._step_completed: - return False - - if mm_hash in self._tombstones: - self._tombstones.discard(mm_hash) - raise _RemoteUnavailable("the remote read failed") - - if entry is not None: - # Present but not ready and not being fetched: its blocks are - # held by a quarantined/settling DMA and cannot be reused. - return False - - try: - size = int(remote["size_bytes"]) - shape = tuple(int(dim) for dim in remote["shape"]) - except (KeyError, TypeError, ValueError) as error: - raise _RemoteUnavailable("the announced size was unusable") from error - # The width is the producer's to report; the row count is not. - num_embeds = pos.get_num_embeds() - if ( - len(shape) != 2 - or shape[0] != num_embeds - or size != self._num_blocks(shape) * self._region.block_size_bytes - ): - logger.warning( - "EC consumer: size mismatch mm_hash=%s announced=%d shape=%s " - "expected %d rows", - mm_hash, - size, - shape, - num_embeds, - ) - raise _RemoteUnavailable("the announced size was wrong") - - try: - started = self._start_xfer(mm_hash, remote, size, (shape[0], shape[1])) - except Exception as error: - logger.exception( - "EC consumer: failed to start NIXL xfer mm_hash=%s", mm_hash - ) - raise _RemoteUnavailable("the read could not be started") from error - if started: - self._in_flight.add(mm_hash) - return False - - def _fail_or_fallback( - self, request: "Request", feature: "MultiModalFeatureSpec", why: str - ) -> bool: - """Resolve an unavailable remote encoding and clear its wait budget. - - Returns: - True if local media or embeddings are available, removing this - item's remote announcement so subsequent steps do not retry it. - False if only placeholder metadata (or no input) is available, - recording the request for take_unavailable_requests() to drain. - - """ - mm_hash = feature.identifier - self._deferred_since.pop((request.request_id, mm_hash), None) - data = feature.data - metadata = ( - self._metadata_resolver.fields_for(feature.modality) if data else set() - ) - if data and any( - value.data is not None for key, value in data.items() if key not in metadata - ): - # Passthrough embeddings bypass the processor cache. SHM address - # items therefore refer to media that the worker can encode. - request.ec_transfer_params = dict(request.ec_transfer_params or {}) - request.ec_transfer_params.pop(mm_hash, None) - logger.warning( - "EC consumer: request %s mm_hash=%s: %s; using local input", - request.request_id, - mm_hash, - why, - ) - return True - self._unrecoverable.add(request.request_id) - logger.error( - "EC consumer: request %s needs remote encoding mm_hash=%s but %s. " - "No local model input is available; failing the request.", - request.request_id, - mm_hash, - why, - ) - return False - - def take_unavailable_requests(self) -> set[str]: - """Return and clear IDs of requests that cannot obtain their encodings. - - The scheduler must abort these requests rather than leave them waiting. - Requests that can fall back to local input are not included, and each - recorded failure is returned only once. - """ - if not self._unrecoverable: - return set() - failed = self._unrecoverable - self._unrecoverable = set() - return failed - - def _start_xfer( - self, - mm_hash: str, - info: "dict[str, Any]", - size_bytes: int, - shape: tuple[int, int], - ) -> bool: - """Allocate a not-ready cache entry and start a NIXL READ into it. - - Returns True if the transfer was started. Returns False when the - cache cannot accommodate the encoding; admission will retry it. - """ - from math import ceil - - from vllm.distributed.ec_transfer.ec_connector.cpu.session import ( - ConsumerSession, - ) - - n_blocks = max(1, ceil(size_bytes / self._region.block_size_bytes)) - entry = self._cache.alloc(mm_hash, n_blocks, shape) - if entry is None: - logger.debug( - "EC consumer: cache full for mm_hash=%s (%d blocks); deferring", - mm_hash, - n_blocks, - ) - return False - indices = list(entry.block_ids) - addr = (info["peer_host"], int(info["peer_port"])) - if addr not in self._sessions: - zmq_conn = self._transport.connect(addr) - assert self._compat_hash is not None - self._sessions[addr] = ConsumerSession( - addr=addr, - zmq_conn=zmq_conn, - transport=self._transport, - data=self._data, - compat_hash=self._compat_hash, - ) - deadline = time.monotonic() + self._ack_timeout_s - try: - self._sessions[addr].start_xfer(mm_hash, indices, deadline) - except Exception: - self._cache.discard(mm_hash) - raise - logger.debug( - "EC consumer: starting NIXL xfer mm_hash=%s from %s:%d blocks=%d", - mm_hash, - addr[0], - addr[1], - n_blocks, - ) - return True - - def _poll_step(self) -> None: - now = time.monotonic() - all_messages = self._transport.poll() - for addr, session in list(self._sessions.items()): - session.poll(all_messages.get(addr, []), now) - for addr in self._transport.poll_dead(): - self._on_peer_down(addr) - for session in self._sessions.values(): - self._process_session_results(session) - - def _process_session_results(self, session: Any) -> None: - r = session.take_results() - for mm_hash in r.completed: - self._in_flight.discard(mm_hash) - self._cache.mark_ready(mm_hash) - self._step_completed.add(mm_hash) - logger.debug("EC consumer: NIXL xfer complete mm_hash=%s", mm_hash) - for mm_hash in r.tombstoned: - self._in_flight.discard(mm_hash) - self._cache.discard(mm_hash) - self._tombstones.add(mm_hash) - logger.debug("EC consumer: NIXL xfer failed mm_hash=%s", mm_hash) - for mm_hash in r.quarantined: - # DMA still running: keep the blocks reserved (entry stays - # not-ready, hence non-evictable) until the xfer settles. - self._in_flight.discard(mm_hash) - self._tombstones.add(mm_hash) - logger.debug( - "EC consumer: NIXL xfer mm_hash=%s quarantined; DMA still running", - mm_hash, - ) - for mm_hash in r.retryable: - # Release this attempt's blocks, retaining the request wait budget. - self._in_flight.discard(mm_hash) - self._cache.discard(mm_hash) - logger.debug( - "EC consumer: NIXL xfer mm_hash=%s retryable; will re-request", - mm_hash, - ) - for mm_hash, _block_indices in r.settled: - self._cache.discard(mm_hash) - logger.debug( - "EC consumer: quarantined NIXL xfer mm_hash=%s settled", mm_hash - ) - - def _on_peer_down(self, addr: Any) -> None: - session = self._sessions.pop(addr, None) - if session is None: - return - session.on_peer_down() - self._process_session_results(session) - session.close() - logger.info("EC consumer: producer peer down addr=%s", addr) - - def _promote_completed_reads(self) -> None: - """Drain the just-completed set built during ``_poll_step``. - - Reads are marked ready in ``_process_session_results``, so a deferred - request is re-admitted by ``ensure_cache_available`` and loaded through - the same local path as a natively cached entry - (``update_state_after_alloc`` -> ``_pending_loads``). No explicit - pin/load bookkeeping is needed here beyond clearing the set. - """ - self._step_completed.clear() - - def _num_blocks(self, shape: tuple[int, ...]) -> int: - size_bytes = shape[0] * shape[1] * self._dtype.itemsize - block_size = self._region.block_size_bytes - return (size_bytes + block_size - 1) // block_size - - def update_state_after_alloc(self, request: "Request", index: int) -> None: - feature = request.mm_features[index] - mm_hash = feature.identifier - - if self._is_producer and self._cache.get(mm_hash) is None: - shape = ( - feature.mm_position.get_num_embeds(), - get_encoder_output_width(self._vllm_config, feature.modality), - ) - entry = self._cache.alloc(mm_hash, self._num_blocks(shape), shape) - if entry is not None: - self._pending_saves[mm_hash] = list(entry.block_ids) - - if self._is_consumer and mm_hash not in self._pending_loads: - entry = self._cache.get(mm_hash) - if entry is not None and entry.ready: - assert entry.shape is not None - self._cache.pin(mm_hash) - transfer_id = self._next_transfer_id - self._next_transfer_id += 1 - self._pending_loads[mm_hash] = ( - transfer_id, - list(entry.block_ids), - entry.shape, - ) - self._load_acks[transfer_id] = (mm_hash, self._expected_load_reports) - - def build_connector_meta( - self, scheduler_output: "SchedulerOutput" - ) -> ECCPUConnectorMetadata: - meta = ECCPUConnectorMetadata() - if self._is_producer: - if self._nixl_enabled and self._producer_session is not None: - self._producer_session.poll_step() - meta.saves = self._pending_saves - self._pending_saves = {} - if self._is_consumer: - if self._nixl_enabled: - self._promote_completed_reads() - meta.loads = self._pending_loads - self._pending_loads = {} - if self._nixl_enabled: - self._first_in_batch = True - return meta - - def update_connector_output(self, connector_output: "ECConnectorOutput") -> None: - """Apply the worker's memcpy-completion report to the cache. - - Completed saves become safe to mark ready. A load is unpinned once - every participating rank has reported its transfer id; reports for a - transfer that has already been released, or for one this scheduler - never dispatched, are ignored. - """ - meta = connector_output.ec_connector_worker_meta - if not isinstance(meta, ECCPUWorkerMetadata): - return - for mm_hash in meta.completed_saves: - entry = self._cache.get(mm_hash) - if entry is None: - logger.debug( - "EC producer: worker reported completed save for unknown " - "mm_hash=%s (already discarded/evicted?)", - mm_hash, - ) - elif not entry.ready: - self._cache.mark_ready(mm_hash) - logger.debug("EC producer: mm_hash=%s marked ready", mm_hash) - for transfer_id in meta.completed_loads: - pending = self._load_acks.get(transfer_id) - if pending is None: - continue - mm_hash, outstanding = pending - if outstanding > 1: - self._load_acks[transfer_id] = (mm_hash, outstanding - 1) - continue - # Drop the entry before unpinning so a replayed report is treated - # as stale rather than releasing the pin a second time. - del self._load_acks[transfer_id] - self._cache.unpin(mm_hash) - logger.debug("EC consumer: mm_hash=%s unpinned after load", mm_hash) - - def has_pending_push_work(self) -> bool: - """Keep the engine stepping so this connector's polls keep running. - - With NIXL enabled the engine tick is the only driver of - ``ProducerSession.poll_step()``, and a producer's work arrives from a - remote consumer, so nothing local would otherwise wake it. Without NIXL - this only has to outlive dispatched saves and loads, so the engine can - quiesce once the worker has confirmed them. - """ - return self._nixl_enabled or self._cache.has_held_entries() - - def request_finished( - self, request: "Request" - ) -> tuple[bool, "dict[str, Any] | None"]: - for feature in request.mm_features: - self._deferred_since.pop((request.request_id, feature.identifier), None) - if not (self._nixl_enabled and self._is_producer): - return False, None - if not request.mm_features: - return False, None - - items = collect_ec_item_metadata(request.mm_features, self._metadata_resolver) - - for feature in request.mm_features: - mm_hash = feature.identifier - entry = self._cache.get(mm_hash) - if entry is None: - # Never saved, or evicted since. Publishing placeholder - # metadata without an address invites an orchestrator to - # rewrite the media into a reference to an encoding no - # consumer can fetch, leaving the decoder nothing to embed. - # Publish neither, so the media stays on the request. - items[mm_hash]["metadata"] = {} - logger.debug( - "EC producer: mm_hash=%s absent at request_finished; " - "announcing no metadata so the media is not rewritten away", - mm_hash, - ) - continue - # Announce even if the save's GPU->mmap copy hasn't been - # confirmed complete yet: a not-ready entry can't be evicted, so - # it will still be here by the time a consumer's XferReq arrives. - # A read arriving before the save lands is NACKed NACK_NOT_READY, - # which the consumer retries on a later step rather than treating - # as a miss. - assert entry.shape is not None - size_bytes = len(entry.block_ids) * self._region.block_size_bytes - items[mm_hash].update( - peer_host=self._peer_host, - peer_port=self._peer_port, - size_bytes=size_bytes, - shape=list(entry.shape), - ) - logger.debug( - "EC producer: announcing NIXL-readable encodings req_id=%s items=%s", - request.request_id, - items, - ) - return False, items - - def shutdown(self) -> None: - self._pending_saves.clear() - self._pending_loads.clear() - self._load_acks.clear() - self._deferred_since.clear() - - self._is_producer = False - self._is_consumer = False - - if self._nixl_enabled: - self._teardown_nixl() - - try: - self._region.cleanup() - except Exception: - logger.debug("ec: region cleanup failed", exc_info=True) - - def _teardown_nixl(self) -> None: - """Close the control sockets and release the NIXL agent. - - Guards each step and clears what it releases, so it is safe on a - scheduler whose NIXL setup failed part-way and safe to call twice. - The shared region is not touched — shutdown() owns that. - """ - if self._producer_session is not None: - try: - self._producer_session.close() - except Exception: - logger.debug("ec: producer session close failed", exc_info=True) - self._producer_session = None - - for session in list(self._sessions.values()): - try: - session.close() - except Exception: - logger.debug("ec: consumer session close failed", exc_info=True) - self._sessions.clear() - - if self._transport is not None: - try: - self._transport.close() - except Exception: - logger.debug("ec: client transport close failed", exc_info=True) - self._transport = None - - if self._data is not None: - try: - self._data.deregister() - except Exception: - logger.debug("ec: deregister failed", exc_info=True) - self._data = None diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py deleted file mode 100644 index 046fba0574..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py +++ /dev/null @@ -1,184 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""EmbeddingCache — named-entry block cache with FIFO eviction. - -Manages a fixed pool of block IDs keyed by content identity (mm_hash). -Entries transition through: not-ready → ready (evictable) → pinned. -Eviction targets ready + unpinned entries in FIFO order. - -All public methods are thread-safe. -""" - -import threading -from collections import OrderedDict - - -class CacheEntry: - """A single cache entry. Read `.block_ids`, `.shape` and `.ready` freely; - mutations only through EmbeddingCache methods. - - `shape` is the encoder output the blocks hold, so a load can restore it - without knowing how wide this model's encoder output is. - """ - - __slots__ = ("block_ids", "shape", "_pin_count") - - def __init__( - self, block_ids: tuple[int, ...], shape: tuple[int, int] | None = None - ) -> None: - self.block_ids = block_ids - self.shape = shape - self._pin_count = -1 # not ready - - @property - def ready(self) -> bool: - return self._pin_count >= 0 - - @property - def evictable(self) -> bool: - return self._pin_count == 0 - - def mark_ready(self): - self._pin_count = 0 - - def pin(self): - self._pin_count += 1 - - -class EmbeddingCache: - """Fixed-size block cache with FIFO eviction. - - Entries are keyed by content identity (e.g. mm_hash). Blocks are - allocated from a free-set; when space is needed, ready + unpinned - entries are evicted oldest-first. - - The caller is responsible for deciding *when* to call ``mark_ready`` - (e.g. after enough engine steps have elapsed for the worker to have - completed the write). - """ - - def __init__(self, num_blocks: int) -> None: - self._num_blocks = num_blocks - self._lock = threading.Lock() - # Available block IDs, LIFO stack. - self._free_blocks: list[int] = list(range(num_blocks)) - # Key → live entry (all states: not-ready, ready, pinned). - self._entries: dict[str, CacheEntry] = {} - # Ready + unpinned entries in insertion order (FIFO eviction). - self._entries_free_list: OrderedDict[str, None] = OrderedDict() - # Total blocks reclaimable via eviction (fail-fast check). - self._evictable_block_count: int = 0 - - def get(self, key: str) -> CacheEntry | None: - """Return the entry for *key*, or None if not present.""" - with self._lock: - return self._entries.get(key) - - def has_held_entries(self) -> bool: - """True if any entry is not-ready or pinned (an in-flight transfer). - - Held entries are exactly those absent from the eviction free list, so - they represent saves awaiting mark_ready or loads awaiting unpin. - """ - with self._lock: - return len(self._entries) != len(self._entries_free_list) - - def alloc( - self, key: str, n_blocks: int, shape: tuple[int, int] | None = None - ) -> CacheEntry | None: - """Allocate *n_blocks* for *key*, evicting as needed. - - The entry starts not-ready. Returns None if there is not enough - space even after evicting all evictable entries. - """ - with self._lock: - assert key not in self._entries, ( - f"EmbeddingCache: duplicate alloc for {key!r}" - ) - assert n_blocks <= self._num_blocks, ( - f"EmbeddingCache: {n_blocks} blocks requested but capacity " - f"is {self._num_blocks}" - ) - if len(self._free_blocks) + self._evictable_block_count < n_blocks: - return None - self._evict_until(n_blocks) - # Canonical ascending order: keeps a freed entry's blocks contiguous - # in the free list when it is evicted, and lets consumers detect - # which of an entry's blocks are adjacent. - block_ids = tuple(sorted(self._free_blocks.pop() for _ in range(n_blocks))) - entry = CacheEntry(block_ids, shape) - self._entries[key] = entry - return entry - - def mark_ready(self, key: str) -> None: - """Mark an entry as ready (data is CPU-visible).""" - with self._lock: - entry = self._entries[key] - assert entry._pin_count == -1, ( - f"EmbeddingCache: mark_ready on already-ready entry {key!r}" - ) - entry.mark_ready() - self._entries_free_list[key] = None - self._evictable_block_count += len(entry.block_ids) - - def pin(self, key: str) -> None: - """Pin an entry (prevent eviction).""" - with self._lock: - entry = self._entries[key] - assert entry.ready, f"EmbeddingCache: pin of not-ready entry {key!r}" - if entry.evictable: - del self._entries_free_list[key] - self._evictable_block_count -= len(entry.block_ids) - entry.pin() - - def pin_if_ready(self, key: str) -> CacheEntry | None: - """Atomically pin *key* if present and ready; return its entry. - - Returns None when the key is absent. An entry returned with `ready` - False was not pinned: its save is still in flight, which the producer - reports differently from a miss because such a read can be retried. - """ - with self._lock: - entry = self._entries.get(key) - if entry is None or not entry.ready: - return entry - if entry.evictable: - del self._entries_free_list[key] - self._evictable_block_count -= len(entry.block_ids) - entry.pin() - return entry - - def unpin(self, key: str) -> None: - """Unpin an entry. Asserts currently pinned.""" - with self._lock: - entry = self._entries[key] - assert entry._pin_count > 0, ( - f"EmbeddingCache: unpin of unpinned entry {key!r}" - ) - entry._pin_count -= 1 - if entry._pin_count == 0: - self._entries_free_list[key] = None - self._evictable_block_count += len(entry.block_ids) - - def discard(self, key: str) -> None: - """Remove a not-ready in-flight entry, returning its blocks to the pool. - - Used when an in-flight fill (e.g. a NIXL READ) fails before the entry - is marked ready. Asserts the entry is present and not ready — ready or - pinned entries are reclaimed through eviction/unpin, not discard. - """ - with self._lock: - entry = self._entries[key] - assert entry._pin_count == -1, ( - f"EmbeddingCache: discard of ready/pinned entry {key!r}" - ) - del self._entries[key] - self._free_blocks.extend(entry.block_ids) - - def _evict_until(self, n_blocks: int) -> None: - """Evict ready+unpinned entries FIFO until enough space. Lock held.""" - while len(self._free_blocks) < n_blocks: - key, _ = self._entries_free_list.popitem(last=False) - entry = self._entries.pop(key) - self._evictable_block_count -= len(entry.block_ids) - self._free_blocks.extend(entry.block_ids) diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/session.py b/vllm/distributed/ec_transfer/ec_connector/cpu/session.py deleted file mode 100644 index 80bbd1fd40..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/session.py +++ /dev/null @@ -1,644 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Session objects for the ECCPUConnector. - -Layer hierarchy: connector → session → transport (connection) - -Two long-lived sessions: - - ProducerSession — owns ZmqServerTransport. The connector calls poll_step() - once per engine step, which drains the transport without blocking and - feeds the raw bytes to poll(). poll() decodes XferReqs, grants or NACKs - them, pins source blocks, sends XferAck replies, drains NIXL completion - notifications, and sweeps expired xfers. - - ConsumerSession — owns one ZmqClientConnection and the NIXL agent - registration for that peer. The connector calls poll(messages, now) - once per engine step with the raw bytes collected for this peer. - poll() decodes XferAcks, dispatches them to ConsumerXfer objects, - and advances all active xfer state machines. - -Internal data classes: - - ProducerXfer — one pinned read grant, keyed by session_id. - ConsumerXfer — state machine for one NIXL READ, keyed by mm_hash. -""" - -import enum -import logging -import time -import uuid -from dataclasses import dataclass, field -from typing import Any - -import msgspec - -from vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq import ( - ZmqClientConnection, - ZmqClientTransport, - ZmqServerTransport, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.data.base import ( - DataTransport, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import ( - EC_CONNECTOR_VERSION, - EXPECTED_NACKS, - RETRYABLE_NACKS, - XferAck, - XferReq, - XferStatus, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import ( - EmbeddingCache, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.utils import PeerAddr -from vllm.logger import init_logger - -logger = init_logger(__name__) - -PRODUCER_PIN_LEASE_S = 30.0 -_CONSUMER_READ_TIMEOUT_S = 20.0 -_CONSUMER_QUARANTINE_TIMEOUT_S = 60.0 - - -# ── XferState ───────────────────────────────────────────────────────────────── - - -class XferState(enum.Enum): - """Lifecycle states of a single ConsumerXfer.""" - - WAITING_ACK = "waiting_ack" # XferReq sent, awaiting XferAck - READING = "reading" # NIXL READ in flight - DONE = "done" # READ completed; promote to local cache - ACK_TIMEOUT = "ack_timeout" # XferAck never arrived; free blocks + retry - READ_FAILED = "read_failed" # Unexpected NIXL state; free blocks + tombstone - QUARANTINED = "quarantined" # Read timed out, NIXL unabortable; keep blocks - SETTLED = "settled" # Quarantined DMA done; safe to free blocks - - -# ── ProducerXfer ────────────────────────────────────────────────────────────── - - -@dataclass -class ProducerXfer: - """One pinned read grant on the producer side. - - Pure data — lifecycle operations (pin/unpin) are performed by ProducerSession. - Keyed by (consumer_session_id, mm_hash) in ProducerSession._active_xfers. - """ - - mm_hash: str - block_indices: list[int] - deadline: float - - def is_expired(self) -> bool: - return time.monotonic() > self.deadline - - -# ── ConsumerXfer ────────────────────────────────────────────────────────────── - - -@dataclass -class ConsumerXfer: - """State machine for one NIXL READ on the consumer side. - - transfer_handle is the in-flight NIXL handle returned by post_read() — - used for check_xfer_state() / release_xfer_handle(). The dlist handle - for the remote peer lives inside DataTransport, not here. - """ - - mm_hash: str - block_indices: list[int] - addr: PeerAddr - deadline: float - data: DataTransport - consumer_session_id: str # session identity; combined with mm_hash for notif_msg - - transfer_handle: Any | None = field(default=None) - _quarantined: bool = field(default=False, init=False) - - def handle_ack(self, ack: XferAck, agent_name: str) -> bool: - """Issue the NIXL READ on an OK XferAck. Returns False on any NACK.""" - if ack.status != XferStatus.OK: - return False - # notif_msg encodes both session_id and mm_hash so the producer can - # do a direct (session_id, mm_hash) → ProducerXfer lookup on completion. - notif_msg = f"{self.consumer_session_id}:{self.mm_hash}".encode() - self.transfer_handle = self.data.post_read( - self.block_indices, - agent_name, - ack.src_block_indices, - notif_msg=notif_msg, - ) - self.deadline = time.monotonic() + _CONSUMER_READ_TIMEOUT_S - logger.debug( - "EC consumer: NIXL READ posted mm_hash=%s agent=%s local_blocks=%d " - "remote_blocks=%d", - self.mm_hash, - agent_name, - len(self.block_indices), - len(ack.src_block_indices), - ) - return True - - def poll(self, now: float) -> XferState: - if self._quarantined: - return self._poll_quarantined(now) - if self.transfer_handle is None: - return ( - XferState.ACK_TIMEOUT if now > self.deadline else XferState.WAITING_ACK - ) - return self._poll_reading(now) - - def _poll_reading(self, now: float) -> XferState: - assert self.transfer_handle is not None - try: - state = self.data.check_xfer_state(self.transfer_handle) - except Exception: - logger.exception( - "EC Connector: check_xfer_state failed for mm_hash=%s", self.mm_hash - ) - self.data.release_xfer_handle(self.transfer_handle) - self.transfer_handle = None - return XferState.READ_FAILED - if state == "DONE": - self.data.release_xfer_handle(self.transfer_handle) - self.transfer_handle = None - logger.debug("EC consumer: NIXL READ completed mm_hash=%s", self.mm_hash) - return XferState.DONE - if state == "PROC": - if now > self.deadline: - self._quarantined = True - self.deadline = now + _CONSUMER_QUARANTINE_TIMEOUT_S - logger.warning( - "EC Connector: READ for mm_hash=%s from %s:%d timed out; " - "quarantining", - self.mm_hash, - self.addr[0], - self.addr[1], - ) - return XferState.QUARANTINED - return XferState.READING - logger.warning( - "EC Connector: READ for mm_hash=%s unexpected NIXL state %r", - self.mm_hash, - state, - ) - self.data.release_xfer_handle(self.transfer_handle) - self.transfer_handle = None - return XferState.READ_FAILED - - def _poll_quarantined(self, now: float) -> XferState: - if self.transfer_handle is None: - return XferState.SETTLED - # If quarantine has been too long, assume NIXL is stuck and evict blocks. - if now > self.deadline: - self.data.release_xfer_handle(self.transfer_handle) - self.transfer_handle = None - logger.warning( - "EC Connector: quarantined mm_hash=%s from %s:%d timed out; " - "evicting blocks", - self.mm_hash, - self.addr[0], - self.addr[1], - ) - return XferState.READ_FAILED - try: - state = self.data.check_xfer_state(self.transfer_handle) - except Exception: - logger.exception( - "EC Connector: check_xfer_state failed for quarantined mm_hash=%s", - self.mm_hash, - ) - state = None - if state == "PROC": - return XferState.QUARANTINED - self.data.release_xfer_handle(self.transfer_handle) - self.transfer_handle = None - logger.debug( - "EC Connector: quarantined mm_hash=%s settled (state=%s)", - self.mm_hash, - state, - ) - return XferState.SETTLED - - def cancel(self) -> None: - """Cancel a WAITING_ACK xfer (peer down). No tombstone — caller allows retry.""" - assert self.transfer_handle is None, "cancel() requires WAITING_ACK state" - - def release(self) -> None: - """Force-release the transfer handle. Shutdown only.""" - if self.transfer_handle is not None: - self.data.release_xfer_handle(self.transfer_handle) - self.transfer_handle = None - - -# ── ProducerSession ─────────────────────────────────────────────────────────── - - -class ProducerSession: - """Long-lived producer session. Owns ZmqServerTransport. - - poll_step() runs once per engine step: drains the transport without - blocking, decodes XferReqs, grants or NACKs each one, drains NIXL completion - notifications, and sweeps expired xfers. - """ - - def __init__( - self, - transport: ZmqServerTransport, - data: DataTransport, - cache: EmbeddingCache, - compat_hash: str, - ) -> None: - self._transport = transport - self._data = data - self._cache = cache # EmbeddingCache — single block owner - self._compat_hash = compat_hash - - # Key: "{consumer_session_id}:{mm_hash}" — matches the NIXL notif_msg exactly, - # so completion notification lookup is a direct dict pop with no parsing. - self._active_xfers: dict[str, ProducerXfer] = {} - - self._req_decoder = msgspec.msgpack.Decoder(XferReq) - self._req_encoder = msgspec.msgpack.Encoder() - - def close(self) -> None: - for xfer in self._active_xfers.values(): - self._cache.unpin(xfer.mm_hash) # release any still-held pins - self._active_xfers.clear() - self._transport.close() - - def poll_step(self) -> None: - """Drain the transport without blocking and process what arrived.""" - self.poll(self._transport.poll(timeout_ms=0)) - - def poll(self, messages: list[tuple[bytes, bytes]]) -> None: - """Decode and process inbound XferReqs, drain NIXL notifs, sweep timeouts. - - messages: (identity, payload) pairs from ZmqServerTransport.poll(). - """ - for identity, payload in messages: - try: - req = self._req_decoder.decode(payload) - except (msgspec.DecodeError, msgspec.ValidationError): - logger.warning("EC Connector: dropped malformed XferReq") - continue - ack = self._grant_or_nack(req) - self._transport.send(identity, self._req_encoder.encode(ack)) - - self._drain_notifs() - self._sweep_timeouts() - - def _grant_or_nack(self, req: XferReq) -> XferAck: - if req.connector_version != EC_CONNECTOR_VERSION: - logger.warning( - "EC Connector: incompatible version for mm_hash=%s", req.mm_hash - ) - return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_VERSION) - if req.compatibility_hash != self._compat_hash: - logger.warning( - "EC Connector: incompatible compat hash for mm_hash=%s", req.mm_hash - ) - return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_INCOMPAT) - if not req.session_id: - # Grants are keyed by (session_id, mm_hash), so an omitted id from - # one peer aliases another's: the first completion would unpin - # blocks the other peer is still reading. - logger.warning( - "EC Connector: XferReq without session_id for mm_hash=%s", req.mm_hash - ) - return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_INTERNAL) - entry = self._cache.pin_if_ready(req.mm_hash) - if entry is None or not entry.ready: - status = ( - XferStatus.NACK_MISSING if entry is None else XferStatus.NACK_NOT_READY - ) - logger.log( - logging.INFO if status in EXPECTED_NACKS else logging.WARNING, - "EC Connector: mm_hash=%s not serveable in local cache (%s); NACKing", - req.mm_hash, - status.name, - ) - return XferAck(mm_hash=req.mm_hash, status=status) - block_indices = list(entry.block_ids) - key = f"{req.session_id}:{req.mm_hash}" - if key in self._active_xfers: - # This peer is re-requesting a grant it still holds, having timed - # out waiting. Overwriting it drops the pin taken for it, which - # neither a completion nor the sweep would ever release. - self._cache.unpin(req.mm_hash) - self._active_xfers[key] = ProducerXfer( - mm_hash=req.mm_hash, - block_indices=block_indices, - deadline=time.monotonic() + PRODUCER_PIN_LEASE_S, - ) - logger.debug( - "EC producer: serving NIXL READ request mm_hash=%s key=%s blocks=%d", - req.mm_hash, - key, - len(block_indices), - ) - return XferAck( - mm_hash=req.mm_hash, - status=XferStatus.OK, - session_id=req.session_id, # echo consumer's session_id back - src_block_indices=block_indices, - agent_metadata=self._data.get_agent_metadata(), - mem_descriptor=self._data.get_mem_descriptor(), - ) - - def _drain_notifs(self) -> None: - try: - notifs = self._data.get_new_notifs() - except Exception: - logger.exception("EC Connector: get_new_notifs failed") - return - for msgs in notifs.values(): - for msg in msgs: - try: - key = msg.decode("utf-8") # "{session_id}:{mm_hash}" - except Exception: - logger.exception("EC Connector: Dropped malformed NIXL notif") - continue - xfer = self._active_xfers.pop(key, None) - if xfer is not None: - self._cache.unpin(xfer.mm_hash) - logger.debug( - "EC producer: NIXL READ completed mm_hash=%s key=%s", - xfer.mm_hash, - key, - ) - - def _sweep_timeouts(self) -> None: - for key, xfer in list(self._active_xfers.items()): - if xfer.is_expired(): - logger.warning( - "EC producer: consumer did not complete NIXL READ " - "mm_hash=%s key=%s within the pin lease; releasing " - "pinned blocks", - xfer.mm_hash, - key, - ) - self._cache.unpin(xfer.mm_hash) - del self._active_xfers[key] - - -# ── ConsumerSession ─────────────────────────────────────────────────────────── - - -@dataclass -class ConsumerSessionResults: - """Results from ConsumerSession.take_results().""" - - completed: set[str] # promote to local cache - tombstoned: set[str] # free blocks + prevent retry this step - quarantined: set[str] # timed-out, DMA still running: prevent retry - # but DO NOT free blocks — session holds them - retryable: set[str] # free blocks, allow another attempt (not ready, peer down) - settled: list[tuple[str, list[int]]] # (mm_hash, block_indices) quarantine cleared - - -class ConsumerSession: - """Long-lived per-producer session on the consumer side. - - Owns the ZMQ DEALER connection and the NIXL agent registration for one - producer peer. The connector calls poll(messages, now) once per engine - step with the raw bytes collected from this peer's DEALER socket. - - poll() decodes XferAcks, dispatches them to the right ConsumerXfer, - advances all active xfer state machines, and drains quarantined xfers. - Results accumulate in internal sets; call take_results() to collect them. - """ - - def __init__( - self, - addr: PeerAddr, - zmq_conn: ZmqClientConnection, - transport: ZmqClientTransport, - data: DataTransport, - compat_hash: str, - ) -> None: - self._addr = addr - self._zmq = zmq_conn - self._transport = transport - self._data = data - self._compat_hash = compat_hash - self._session_id = str( - uuid.uuid4() - ) # stable identity for (session_id:mm_hash) keys - - self._nixl_agent_name: str | None = None - self._nixl_metadata_bytes: bytes | None = None - - self._xfers: dict[str, ConsumerXfer] = {} # mm_hash → xfer - self._quarantined: list[ConsumerXfer] = [] # DMA pending, to poll - self._newly_quarantined: set[str] = set() # quarantined since last take_results - - self._completed: set[str] = set() - self._tombstoned: set[str] = set() - self._retryable: set[str] = set() - self._settled: list[tuple[str, list[int]]] = [] - - self._encoder = msgspec.msgpack.Encoder() - self._decoder = msgspec.msgpack.Decoder(XferAck) - - @property - def addr(self) -> PeerAddr: - return self._addr - - def start_xfer( - self, mm_hash: str, block_indices: list[int], deadline: float - ) -> None: - """Create a ConsumerXfer and send the XferReq over the connection.""" - xfer = ConsumerXfer( - mm_hash=mm_hash, - block_indices=block_indices, - addr=self._addr, - deadline=deadline, - data=self._data, - consumer_session_id=self._session_id, - ) - self._zmq.send( - self._encoder.encode( - XferReq( - mm_hash=mm_hash, - compatibility_hash=self._compat_hash, - session_id=self._session_id, - ) - ) - ) - self._xfers[mm_hash] = xfer - logger.debug( - "EC Connector: XferReq sent mm_hash=%s to %s:%d", - mm_hash, - self._addr[0], - self._addr[1], - ) - - def poll(self, messages: list[bytes], now: float) -> None: - """Decode inbound XferAcks, dispatch to xfers, and advance all state machines. - - messages: raw bytes from ZmqClientTransport.poll() for this peer. - Results accumulate until take_results() is called. - """ - for raw in messages: - try: - ack = self._decoder.decode(raw) - except (msgspec.DecodeError, msgspec.ValidationError): - logger.warning( - "EC Connector: dropped malformed XferAck from %s:%d", - self._addr[0], - self._addr[1], - ) - continue - self._dispatch_ack(ack) - - for mm_hash, xfer in list(self._xfers.items()): - state = xfer.poll(now) - if state == XferState.DONE: - self._completed.add(mm_hash) - del self._xfers[mm_hash] - elif state == XferState.ACK_TIMEOUT: - # The scheduler bounds retries with a per-request deadline. - logger.warning( - "EC consumer: no XferAck for mm_hash=%s from %s:%d in time; " - "will re-request", - mm_hash, - self._addr[0], - self._addr[1], - ) - self._retryable.add(mm_hash) - del self._xfers[mm_hash] - elif state == XferState.READ_FAILED: - self._tombstoned.add(mm_hash) - del self._xfers[mm_hash] - elif state == XferState.QUARANTINED: - # DMA may still be writing — caller must NOT free blocks. - self._quarantined.append(xfer) - self._newly_quarantined.add(mm_hash) - del self._xfers[mm_hash] - - self._drain_quarantine(now) - - def _dispatch_ack(self, ack: XferAck) -> None: - xfer = self._xfers.get(ack.mm_hash) - if xfer is None or xfer.transfer_handle is not None: - return # unknown mm_hash or duplicate ack - - if ack.status != XferStatus.OK: - logger.log( - logging.INFO if ack.status in EXPECTED_NACKS else logging.WARNING, - "EC Connector: NACK %s from %s:%d for mm_hash=%s", - ack.status.name, - self._addr[0], - self._addr[1], - ack.mm_hash, - ) - del self._xfers[ack.mm_hash] - # The producer announced this encoding and its save is still in - # flight, so another attempt on a later step can be granted. Every - # other NACK is terminal for this mm_hash. - if ack.status in RETRYABLE_NACKS: - self._retryable.add(ack.mm_hash) - else: - self._tombstoned.add(ack.mm_hash) - return - - try: - agent_name = self._ensure_registered(ack.agent_metadata, ack.mem_descriptor) - except Exception: - logger.exception( - "EC Connector: failed to register peer %s:%d for mm_hash=%s", - self._addr[0], - self._addr[1], - ack.mm_hash, - ) - del self._xfers[ack.mm_hash] - self._tombstoned.add(ack.mm_hash) - return - - if not xfer.handle_ack(ack, agent_name): - del self._xfers[ack.mm_hash] - self._tombstoned.add(ack.mm_hash) - - def _ensure_registered(self, metadata: bytes, mem_descriptor: bytes) -> str: - """Register or re-register the peer in DataTransport. Returns agent_name.""" - if self._nixl_agent_name is not None and self._nixl_metadata_bytes == metadata: - return self._nixl_agent_name - if self._nixl_agent_name is not None: - logger.info( - "EC Connector: producer %s:%d restarted; re-registering", - self._addr[0], - self._addr[1], - ) - self._data.remove_remote_peer(self._nixl_agent_name) - agent_name = self._data.add_remote_peer(metadata, mem_descriptor) - self._nixl_agent_name = agent_name - self._nixl_metadata_bytes = metadata - logger.debug( - "EC Connector: registered peer %s:%d agent=%s", - self._addr[0], - self._addr[1], - agent_name, - ) - return agent_name - - def _drain_quarantine(self, now: float) -> None: - still_pending: list[ConsumerXfer] = [] - for xfer in list(self._quarantined): - state = xfer.poll(now) - if state == XferState.SETTLED: - self._settled.append((xfer.mm_hash, xfer.block_indices)) - else: - still_pending.append(xfer) - self._quarantined = still_pending - - def on_peer_down(self) -> None: - """WAITING_ACK xfers → cancel (retry allowed); READING xfers → quarantine.""" - for mm_hash, xfer in list(self._xfers.items()): - if xfer.transfer_handle is None: - xfer.cancel() - self._retryable.add(mm_hash) - else: - self._quarantined.append(xfer) - self._newly_quarantined.add(mm_hash) - self._xfers.clear() - - def take_results(self) -> ConsumerSessionResults: - """Return and clear all accumulated results since the last call.""" - results = ConsumerSessionResults( - completed=self._completed, - tombstoned=self._tombstoned, - quarantined=self._newly_quarantined, - retryable=self._retryable, - settled=self._settled, - ) - self._completed = set() - self._tombstoned = set() - self._newly_quarantined = set() - self._retryable = set() - self._settled = [] - return results - - def has_in_flight(self) -> bool: - return bool(self._xfers) or bool(self._quarantined) - - def close(self) -> None: - """Release all transfer handles, deregister NIXL peer, close connection. - - Connection teardown is routed through the transport so the pooled - connection is dropped in lock-step with the session — a session and its - DEALER connection have a 1:1 lifecycle, so a later connect() to the same - peer builds a fresh socket instead of handing back this dead one. - """ - for xfer in self._xfers.values(): - xfer.release() - self._xfers.clear() - for xfer in self._quarantined: - xfer.release() - self._quarantined.clear() - if self._nixl_agent_name is not None: - self._data.remove_remote_peer(self._nixl_agent_name) - self._transport.remove(self._addr) diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/utils.py b/vllm/distributed/ec_transfer/ec_connector/cpu/utils.py deleted file mode 100644 index 6a8d8e07e0..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/utils.py +++ /dev/null @@ -1,47 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Supporting utilities for the ECCPUConnector scheduler.""" - -import msgspec - -from vllm.logger import init_logger - -# (host, port) pair identifying a remote peer's ZMQ side-channel endpoint. -PeerAddr = tuple[str, int] - -logger = init_logger(__name__) - - -# Msgpack (de)serialization for the list-of-(addr, size, device_id) -# descriptor form that NIXL's `get_xfer_descs` expects. -_MemDescList = list[tuple[int, int, int]] - -_mem_desc_encoder = msgspec.msgpack.Encoder() -_mem_desc_decoder = msgspec.msgpack.Decoder(_MemDescList) - - -def serialize_mem_descriptor(descs: _MemDescList) -> bytes: - return _mem_desc_encoder.encode(descs) - - -def deserialize_mem_descriptor(payload: bytes) -> _MemDescList: - return _mem_desc_decoder.decode(payload) - - -def build_block_descs( - base_ptr: int, - num_blocks: int, - block_size_bytes: int, - device_id: int = 0, -) -> _MemDescList: - """Per-block `(addr, size, device_id)` tuples for an `ECSharedRegion`. - - Handed to `nixl.get_xfer_descs` to build a dlist where descriptor `i` - addresses block `i` in the mmap. This makes a WRITE to block indices - `[i1, i2, ...]` a simple `make_prepped_xfer(..., [i1, i2])` on both - ends — no per-request address math. - """ - return [ - (base_ptr + i * block_size_bytes, block_size_bytes, device_id) - for i in range(num_blocks) - ] diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py deleted file mode 100644 index 95c6c96128..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py +++ /dev/null @@ -1,455 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Worker-side of the ECCPUConnector. - -Thin, stateless across steps: opens the shared mmap region and uses the -per-step connector metadata (`ECCPUConnectorMetadata`) to decide which -blocks to copy in each direction. -""" - -from collections import deque -from collections.abc import Sequence -from dataclasses import dataclass -from enum import Enum, auto -from itertools import chain -from typing import TYPE_CHECKING - -import numpy as np -import torch - -from vllm._custom_ops import swap_blocks_batch -from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( - ECCPUConnectorMetadata, - ECCPUWorkerMetadata, - create_ec_shared_region, -) -from vllm.distributed.ec_transfer.ec_connector.cpu.worker.descriptor_buffers import ( - DescriptorBufferPool, - DescriptorBuffers, -) -from vllm.distributed.parallel_state import ( - get_pcp_group, - get_tensor_model_parallel_rank, -) -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.utils.platform_utils import is_pin_memory_available - -if TYPE_CHECKING: - from vllm.config import VllmConfig - -logger = init_logger(__name__) - - -class ECCPUTransferDirection(Enum): - HOST_TO_DEVICE = auto() - DEVICE_TO_HOST = auto() - - -def _coalesce_runs( - block_ids: Sequence[int], -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Split block ids into runs that are consecutive in the region. - - Returns `(slots, first_blocks, num_blocks)`, one element per run, where - `slots` is the run's offset in the flat block sequence. Per-descriptor - overhead dominates these copies, so describing a run with one descriptor - rather than one per block is worth an order of magnitude. - - Both callers derive the non-region side of a run from its slot, so that - side advances in lockstep with the region and only the region has to be - split. `EmbeddingCache` hands out ascending ids, which is what makes runs - findable at all. - """ - ids = np.fromiter(block_ids, dtype=np.int64, count=len(block_ids)) - breaks = np.flatnonzero(np.diff(ids) != 1) + 1 - slots = np.concatenate(([0], breaks)) - ends = np.concatenate((breaks, [ids.size])) - return slots, ids[slots], ends - slots - - -@dataclass -class Transfer: - """A batched copy in flight on a dedicated GPU stream. - - Its descriptor buffers and stream are held until `end_event` fires; only - then are the buffers and stream safe to recycle and the `completions` safe - to report. `start_event`/`end_event` bracket the copy so its elapsed time - can be measured on completion. - - `completions` is what the scheduler is told about this copy: mm_hashes for - saves, transfer ids for loads. - """ - - start_event: torch.Event - end_event: torch.Event - completions: list[str] | list[int] - bufs: DescriptorBuffers - stream: torch.Stream - num_bytes: int - - -class ECCPUWorker: - """Worker-side delegate for the ECCPUConnector. - - - Producer role: copies `encoder_cache[mm_hash]` → `mmap[block_ids]` - for each entry in `metadata.saves`. Descriptor buffers are filled - directly in `save_caches`; the actual DMA is issued as a single - batched call in `flush_saves`. - - Consumer role: copies `mmap[block_ids]` → `encoder_cache[mm_hash]` - for all entries in `metadata.loads` via a single `swap_blocks_batch` - call on a dedicated stream. - - On `ec_both` nodes both paths run back-to-back in a single step. - - Every batched copy runs on a stream drawn from a pool so it overlaps the - compute stream, bracketed by a start/end event pair. The stream, events, - and descriptor buffers are recycled once the end event fires. - - Both paths move memory between the two streams, so each buffer is registered - against the stream that is not its own: save sources against the save stream, - load destinations against the compute stream. Without that, the caching - allocator would be free to reissue memory a copy or a queued read still - needs, since it only tracks the stream a buffer was allocated on. - """ - - def __init__(self, vllm_config: "VllmConfig") -> None: - self._region = create_ec_shared_region(vllm_config) - # Model dtype; used to reinterpret raw int8 blocks on load. - self._dtype = vllm_config.model_config.dtype - - self._pin_shared_region() - - # All TP/PCP ranks hold identical encoder output. Only one rank - # per mmap needs to write — saves host memory bandwidth. - # DCP is a subdivision of TP, so tp_rank==0 covers it. - self._is_save_rank = ( - get_tensor_model_parallel_rank() == 0 and get_pcp_group().rank_in_group == 0 - ) - - # Descriptor buffer pool (recycled across steps, shared by both paths). - self._buf_pool = DescriptorBufferPool() - - # Active save buffer being filled during save_caches calls this step. - self._save_bufs: DescriptorBuffers | None = None - self._save_count: int = 0 - # Bytes accumulated into the active save buffer this step, for logging. - self._save_bytes: int = 0 - # mm_hashes accumulated into the active save buffer this step. - self._save_mm_hashes: list[str] = [] - # Stream the active save buffer's copy will run on, held from the moment - # the buffer opens so `save_caches` can register its sources against it. - self._save_stream: torch.Stream | None = None - - # Batched copies whose end event has not yet fired. Buffers and stream - # stay held until the event fires; the completions are reported to the - # scheduler only then. Separate per direction because saves and loads - # can be in flight concurrently. - self._inflight_saves: deque[Transfer] = deque() - self._inflight_loads: deque[Transfer] = deque() - - # Streams and timing events recycled across steps to avoid per-step - # create/destroy churn. Shared by both directions. - self._stream_pool: list[torch.Stream] = [] - self._event_pool: list[torch.Event] = [] - - def _acquire_stream(self) -> torch.Stream: - if self._stream_pool: - return self._stream_pool.pop() - return current_platform.Stream() - - def _acquire_event(self) -> torch.Event: - return ( - self._event_pool.pop() - if self._event_pool - else torch.Event(enable_timing=True) - ) - - def _collect_finished(self, inflight: deque[Transfer], direction: str) -> list: - """Pop transfers whose end event has fired, recycle their stream, - events, and buffers, and return their completions. - - The front check is conservative: a transfer is reported only once its - own end event fires, so completions are always genuinely done, and each - is reported exactly once because the transfer is popped as it is - reported. A later transfer that happens to finish first simply waits - behind the front and is reported on a subsequent poll. - """ - done: list = [] - while inflight and inflight[0].end_event.query(): - transfer = inflight.popleft() - done.extend(transfer.completions) - elapsed_ms = transfer.start_event.elapsed_time(transfer.end_event) - logger.debug( - "EC %s: %d entr%s (%d bytes) took %.3f ms", - direction, - len(transfer.completions), - "y" if len(transfer.completions) == 1 else "ies", - transfer.num_bytes, - elapsed_ms, - ) - self._buf_pool.release(transfer.bufs) - self._stream_pool.append(transfer.stream) - self._event_pool.append(transfer.start_event) - self._event_pool.append(transfer.end_event) - return done - - def _pin_shared_region(self) -> None: - """Prepare the shared region for device transfers.""" - if is_pin_memory_available(): - self._region.pin_memory() - - def _submit_transfer( - self, - bufs: DescriptorBuffers, - count: int, - direction: ECCPUTransferDirection, - ) -> None: - """Submit a transfer on the current backend stream.""" - src_ptrs = bufs.src_ptrs[:count] - dst_ptrs = bufs.dst_ptrs[:count] - sizes = bufs.sizes[:count] - if direction == ECCPUTransferDirection.HOST_TO_DEVICE: - swap_blocks_batch( - src_ptrs, - dst_ptrs, - sizes, - is_src_access_order_any=True, - ) - else: - swap_blocks_batch(src_ptrs, dst_ptrs, sizes) - - def _shutdown_transfer_backend(self) -> None: - """Release backend-specific transfer resources.""" - - def save_caches( - self, - encoder_cache: dict[str, torch.Tensor], - mm_hash: str, - connector_metadata: ECCPUConnectorMetadata, - ) -> None: - """Fill descriptor buffers directly for batched flush. - - Registers `encoder_cache[mm_hash]` against the stream the copy will run - on, so the caching allocator will not hand that memory to another - allocation before the copy has read it. The descriptor buffers hold raw - addresses, and the caller is free to evict the entry as soon as this - step's saves are dispatched. - """ - if not self._is_save_rank: - return - block_ids = connector_metadata.saves.get(mm_hash) - if block_ids is None: - return - - src = encoder_cache[mm_hash] - total_bytes = src.numel() * src.element_size() - block_size = self._region.block_size_bytes - allocated_bytes = len(block_ids) * block_size - assert total_bytes <= allocated_bytes, ( - f"EC: encoder output exceeds allocated blocks for mm_hash={mm_hash}: " - f"{total_bytes} bytes but only {allocated_bytes} allocated " - f"({len(block_ids)} blocks × {block_size}). " - f"shape={list(src.shape)}" - ) - - if self._save_bufs is None: - total = sum(len(v) for v in connector_metadata.saves.values()) - self._save_bufs = self._buf_pool.acquire(total) - self._save_stream = self._acquire_stream() - self._save_mm_hashes = [] - - assert self._save_count + len(block_ids) <= self._save_bufs.src_ptrs.numel() - - bufs = self._save_bufs - assert self._save_stream is not None - src.record_stream(self._save_stream) - # Device addresses are unsigned and can have the high bit set (XPU USM), - # so they are added up in uint64; signed arithmetic rejects them. - src_base = np.uint64(src.view(-1).view(torch.uint8).data_ptr()) - dst_base = self._region.blocks.data_ptr() - # Copy only the tensor bytes, excluding any block padding. - n_used = -(-total_bytes // block_size) - if n_used: - slots, first_blocks, num_blocks = _coalesce_runs(block_ids[:n_used]) - run_bytes = num_blocks * block_size - # The output can stop partway into the final run. - run_bytes[-1] = total_bytes - slots[-1] * block_size - bufs.add_copies( - self._save_count, - src_base + slots.astype(np.uint64) * np.uint64(block_size), - dst_base + first_blocks * block_size, - run_bytes, - ) - self._save_count += slots.size - - self._save_bytes += total_bytes - self._save_mm_hashes.append(mm_hash) - - def flush_saves(self) -> None: - """Flush all accumulated saves in a single swap_blocks_batch call. - - Runs the copy on the stream the batch acquired when it opened, gated - behind the compute stream that produced the encoder outputs, brackets it - with start/end events, and enqueues the batch as in-flight. The stream, - events, and descriptor buffers are recycled once the end event fires (see - `_collect_finished`), which is also when the saved mm_hashes become - safe to mark ready. - """ - if self._save_count == 0: - return - - bufs = self._save_bufs - stream = self._save_stream - assert bufs is not None and stream is not None - n = self._save_count - num_bytes = self._save_bytes - - # Gate the GPU→CPU copy behind the compute stream: it reads the encoder - # outputs the model just produced there. - stream.wait_stream(current_platform.current_stream()) - start_event = self._acquire_event() - end_event = self._acquire_event() - with current_platform.stream(stream): - start_event.record(stream) - self._submit_transfer( - bufs, - n, - ECCPUTransferDirection.DEVICE_TO_HOST, - ) - end_event.record(stream) - - self._inflight_saves.append( - Transfer( - start_event=start_event, - end_event=end_event, - completions=self._save_mm_hashes, - bufs=bufs, - stream=stream, - num_bytes=num_bytes, - ) - ) - - self._save_bufs = None - self._save_stream = None - self._save_count = 0 - self._save_bytes = 0 - self._save_mm_hashes = [] - - def start_load_caches( - self, - encoder_cache: dict[str, torch.Tensor], - connector_metadata: ECCPUConnectorMetadata, - ) -> None: - """Consumer path: single batched copy of all loads from mmap→GPU. - - The destination is registered against the compute stream that consumes - it, so the caching allocator will not hand that memory to a later load - while reads of it are still queued there. - """ - if not connector_metadata.loads: - return - - block_size = self._region.block_size_bytes - blocks = self._region.blocks - dtype = self._dtype - device_type = current_platform.device_type - src_base = blocks.data_ptr() - - # Copy every dispatched load. The scheduler routes an input here only - # after missing the GPU encoder cache, so a resident hit is not expected; - # copying unconditionally means every participating rank reports each - # transfer exactly once, which is what the scheduler counts to release - # the pin. - # Ordered by first block so entries that neighbour each other in the - # region land next to each other in the destination and coalesce into - # one descriptor. Every loop below consumes this same order. - load_items = sorted( - connector_metadata.loads.items(), key=lambda kv: kv[1][1][0] - ) - total_blocks = sum(len(block_ids) for _, (_, block_ids, _) in load_items) - - stream = self._acquire_stream() - compute_stream = current_platform.current_stream() - start_event = self._acquire_event() - end_event = self._acquire_event() - with current_platform.stream(stream): - # Single contiguous destination buffer for all loads. - dst_buf = torch.empty( - total_blocks, block_size, dtype=torch.int8, device=device_type - ) - dst_buf.record_stream(compute_stream) - # See save_caches: device addresses are added up in uint64. - dst_buf_base = np.uint64(dst_buf.data_ptr()) - - bufs = self._buf_pool.acquire(total_blocks) - slots, first_blocks, num_blocks = _coalesce_runs( - list(chain.from_iterable(ids for _, (_, ids, _) in load_items)) - ) - bufs.add_copies( - 0, - src_base + first_blocks * block_size, - dst_buf_base + slots.astype(np.uint64) * np.uint64(block_size), - num_blocks * block_size, - ) - op_idx = slots.size - - start_event.record(stream) - self._submit_transfer( - bufs, - op_idx, - ECCPUTransferDirection.HOST_TO_DEVICE, - ) - end_event.record(stream) - self._inflight_loads.append( - Transfer( - start_event=start_event, - end_event=end_event, - completions=[tid for _, (tid, _, _) in load_items], - bufs=bufs, - stream=stream, - num_bytes=total_blocks * block_size, - ) - ) - - # Slice contiguous buffer into per-hash views. - offset = 0 - for mm_hash, (_, block_ids, shape) in load_items: - n = len(block_ids) - flat = dst_buf[offset : offset + n].view(dtype).reshape(-1) - encoder_cache[mm_hash] = flat[: shape[0] * shape[1]].view(shape) - offset += n - - current_platform.current_stream().wait_stream(stream) - - def build_connector_worker_meta(self) -> ECCPUWorkerMetadata | None: - """Report the GPU copies that completed on this rank this step: saved - mm_hashes and loaded transfer ids. - - Returns None when nothing finished, so the scheduler sees no payload. - """ - completed_saves = self._collect_finished(self._inflight_saves, "save") - completed_loads = self._collect_finished(self._inflight_loads, "load") - if not completed_saves and not completed_loads: - return None - return ECCPUWorkerMetadata( - completed_saves=completed_saves, - completed_loads=completed_loads, - ) - - def shutdown(self) -> None: - for transfer in (*self._inflight_saves, *self._inflight_loads): - transfer.end_event.synchronize() - self._shutdown_transfer_backend() - self._save_bufs = None - self._save_stream = None - self._save_count = 0 - self._save_bytes = 0 - self._save_mm_hashes = [] - self._inflight_saves.clear() - self._inflight_loads.clear() - self._stream_pool.clear() - self._event_pool.clear() - try: - self._region.cleanup() - except Exception: - logger.debug("EC: worker region cleanup failed", exc_info=True) diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py deleted file mode 100644 index 5b2109187f..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py +++ /dev/null @@ -1,90 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Reusable pool of (src_ptrs, dst_ptrs, sizes) tensor triples. - -Used by ECCPUWorker to batch swap_blocks_batch descriptors without -per-step allocation overhead. -""" - -from typing import NamedTuple - -import numpy as np -import torch - -from vllm.platforms import current_platform - -# CUDA/ROCm cache_kernels.cu requires int64 pointers; the XPU DMA engine -# requires uint64 (see vllm._custom_ops.swap_blocks_batch). -_PTR_DTYPE = torch.uint64 if current_platform.is_xpu() else torch.int64 - - -class DescriptorBuffers(NamedTuple): - src_ptrs: torch.Tensor - dst_ptrs: torch.Tensor - sizes: torch.Tensor - # Numpy aliases of the three tensors, written via add_copies(). - src_np: np.ndarray - dst_np: np.ndarray - sizes_np: np.ndarray - - def add_copies( - self, - first: int, - sources: np.ndarray, - destinations: np.ndarray, - byte_counts: np.ndarray, - ) -> None: - """Store a batch of copies in these buffers, the first at *first*. - - One copy is one source address, one destination address, and a number - of bytes. `swap_blocks_batch` performs every copy held in the buffers - in a single call, so a caller fills in all of a step's copies before - asking for any of them to run, and passes *first* to say where its own - copies begin. - - Whole batches go in at once because storing a single number into a - torch tensor costs microseconds. Adding copies one at a time can take - longer than performing them. - - TODO(torch>=2.14): write the tensors directly and drop these numpy - aliases once the minimum supported torch is 2.14. They exist only - because torch's setitem unpacks the value as a signed long long before - 2.14 (pytorch#191458), rejecting XPU USM addresses >= 2**63, while the - two's-complement rewrite those versions accept is in turn rejected by - uint64 from 2.14 on. Numpy casts against the array's own type and so - works on either. Byte counts are small enough to need none of this - care, and use their alias only to match the addresses. - """ - last = first + len(sources) - self.src_np[first:last] = sources - self.dst_np[first:last] = destinations - self.sizes_np[first:last] = byte_counts - - -class DescriptorBufferPool: - """Pool of descriptor buffer triples for swap_blocks_batch. - - Each buffer is a `DescriptorBuffers` namedtuple of three 1-D tensors - (dtype `_PTR_DTYPE`, platform-dependent) of equal length, paired with - numpy aliases used to fill them. Buffers are recycled across steps; if a - returned buffer is too small it is discarded and a fresh one allocated. - """ - - def __init__(self) -> None: - # LIFO stack of idle buffer triples. - self._pool: list[DescriptorBuffers] = [] - - def acquire(self, n: int) -> DescriptorBuffers: - """Get a buffer triple with capacity >= *n*.""" - if self._pool: - bufs = self._pool.pop() - if bufs.src_ptrs.numel() >= n: - return bufs - src, dst, sizes = (torch.empty(n, dtype=_PTR_DTYPE) for _ in range(3)) - return DescriptorBuffers( - src, dst, sizes, src.numpy(), dst.numpy(), sizes.numpy() - ) - - def release(self, bufs: DescriptorBuffers) -> None: - """Return a buffer triple to the pool for reuse.""" - self._pool.append(bufs) diff --git a/vllm/distributed/ec_transfer/ec_connector/example_connector.py b/vllm/distributed/ec_transfer/ec_connector/example_connector.py deleted file mode 100644 index 97075cdf38..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/example_connector.py +++ /dev/null @@ -1,231 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import os -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any - -import safetensors - -from vllm.config import VllmConfig -from vllm.distributed.ec_transfer.ec_connector.base import ( - ECConnectorBase, - ECConnectorMetadata, - ECConnectorRole, -) -from vllm.distributed.ec_transfer.ec_connector.utils import ( - PlaceholderMetadataResolver, - collect_ec_item_metadata, -) -from vllm.logger import init_logger -from vllm.utils.gpu_sync_debug import gpu_sync_allowed -from vllm.v1.core.sched.output import SchedulerOutput - -if TYPE_CHECKING: - from vllm.v1.request import Request - -logger = init_logger(__name__) - - -@dataclass -class MMMeta: - mm_hash: str - num_token: int - - @staticmethod - def make_meta(mm_hash, num_token) -> "MMMeta": - return MMMeta(mm_hash=mm_hash, num_token=num_token) - - -@dataclass -class ECExampleConnectorMetadata(ECConnectorMetadata): - mm_datas: list[MMMeta] - - def __init__(self): - self.mm_datas = [] - - def add_mm_data(self, mm_data: MMMeta): - self.mm_datas.append(mm_data) - - -class ECExampleConnector(ECConnectorBase): - # NOTE: This is Simple debug implementation of the EC connector. - # It save / load the EC cache to / from the disk. - - def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole): - super().__init__(vllm_config=vllm_config, role=role) - # req_id -> index - self._mm_datas_need_loads: dict[str, int] = {} - self._metadata_resolver = PlaceholderMetadataResolver(vllm_config.model_config) - transfer_config = vllm_config.ec_transfer_config - if transfer_config is not None: - self._storage_path = transfer_config.get_from_extra_config( - "shared_storage_path", "/tmp" - ) - logger.debug(transfer_config) - logger.debug("Shared storage path is %s", self._storage_path) - else: - raise ValueError("ec_transfer_config must be set for ECConnectorBase") - - def start_load_caches(self, encoder_cache, **kwargs) -> None: - """Start loading the cache from the connector into vLLM's encoder cache. - - This method loads the encoder cache based on metadata provided by the scheduler. - It is called before `_gather_mm_embeddings` for the EC Connector. For EC, - the `encoder_cache` and `mm_hash` are stored in `kwargs`. - - Args: - encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal - data hashes (`mm_hash`) to encoder cache tensors. - kwargs (dict): Additional keyword arguments for the connector. - - """ - from vllm.platforms import current_platform - - # Get the metadata - metadata: ECConnectorMetadata = self._get_connector_metadata() - assert isinstance(metadata, ECExampleConnectorMetadata) - assert encoder_cache is not None - if metadata is None: - logger.warning( - "In connector.start_load_caches, but the connector metadata is None" - ) - return - # A bare "cuda" makes safetensors load onto cuda:0. Pin the load to the - # selected device so tensor-parallel workers load the cache locally. - device = current_platform.device_type - if current_platform.is_cuda_alike(): - device = f"{device}:{current_platform.current_device()}" - # Load the EC for each mm data - for mm_data in metadata.mm_datas: - if mm_data.mm_hash in encoder_cache: - continue - filename = self._generate_filename_debug(mm_data.mm_hash) - with gpu_sync_allowed(): - ec_cache = safetensors.torch.load_file(filename, device=device)[ - "ec_cache" - ] - encoder_cache[mm_data.mm_hash] = ec_cache - logger.debug("Success load encoder cache for hash %s", mm_data.mm_hash) - - def save_caches(self, encoder_cache, mm_hash, **kwargs) -> None: - """Save the encoder cache to the connector. - - This method saves the encoder cache from the worker's local storage - to shared storage or another external connector. - - Args: - encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal - data hashes (`mm_hash`) to encoder cache tensors. - mm_hash (str): The hash of the multimodal data whose cache is being saved. - kwargs (dict): Additional keyword arguments for the connector. - - """ - # Return if it is PD Instance - if not self.is_producer: - return - filename = self._generate_filename_debug(mm_hash) - ec_cache = encoder_cache[mm_hash] - with gpu_sync_allowed(): - tensors = {"ec_cache": ec_cache.detach().cpu()} - safetensors.torch.save_file(tensors, filename) - logger.debug("Save cache successful for mm_hash %s", mm_hash) - - def has_cache_item( - self, - identifier: str, - ) -> bool: - """Check if cache exist externally for the media. - - Args: - identifier (str): the identifier of the media. - - Returns: - Bool indicate that media exists in cache or not - - """ - return self._found_match_for_mm_data(identifier) - - def update_state_after_alloc( - self, - request: "Request", - index: int, - ) -> None: - """Update ECConnector state after encoder cache allocation.""" - mm_hash = request.mm_features[index].identifier - # Only load cache if it is consumer and cache exists - if not self.is_consumer or not self.has_cache_item(mm_hash): - return - num_encoder_token = request.get_num_encoder_embeds(index) - self._mm_datas_need_loads[mm_hash] = num_encoder_token - - def build_connector_meta( - self, - scheduler_output: SchedulerOutput, - ) -> ECConnectorMetadata: - """Build the connector metadata for this step. - - This function should NOT modify any fields in the scheduler_output. - Also, calling this function will reset the state of the connector. - This only build for load mm_data only - Args: - scheduler_output (SchedulerOutput): the scheduler output object. - """ - meta = ECExampleConnectorMetadata() - for mm_hash, num_encoder_token in self._mm_datas_need_loads.items(): - meta.add_mm_data(MMMeta.make_meta(mm_hash, num_encoder_token)) - self._mm_datas_need_loads.clear() - return meta - - def request_finished( - self, - request: "Request", - ) -> tuple[bool, dict[str, Any] | None]: - """Report each item's cache key and grid so a consumer can skip the - image transform. - - A consumer only needs the grid to size the prompt's placeholder range; - the embedding itself arrives through this connector. Reporting the grid - the producer actually computed keeps the two sides in agreement without - the caller re-deriving it from the raw media. - """ - if not self.is_producer: - return False, None - - items = collect_ec_item_metadata(request.mm_features, self._metadata_resolver) - if not items: - return False, None - return False, items - - # ============================== - # Helper functions - # ============================== - - def _found_match_for_mm_data(self, mm_hash) -> bool: - """Check if the cache is hit for the request.""" - filename = self._generate_filename_debug(mm_hash, create_folder=False) - return os.path.exists(filename) - - def _generate_foldername_debug( - self, - mm_hash: str, - create_folder: bool = True, # <- now defaults to True - ) -> str: - """Return the folder in which the cache for this mm_hash lives. - If `create_folder` is True (default) the directory is created - recursively the first time it is needed. - """ - foldername = os.path.join(self._storage_path, mm_hash) - if create_folder: - os.makedirs(foldername, exist_ok=True) - return foldername - - def _generate_filename_debug( - self, - mm_hash: str, - create_folder: bool = True, - ) -> str: - """Return the full path of the safetensors file for this mm_hash. - If `create_folder` is True (default) the parent directory is created. - """ - foldername = self._generate_foldername_debug(mm_hash, create_folder) - return os.path.join(foldername, "encoder_cache.safetensors") diff --git a/vllm/entrypoints/cli/main.py b/vllm/entrypoints/cli/main.py index 76ecce9d81..713ff1261d 100644 --- a/vllm/entrypoints/cli/main.py +++ b/vllm/entrypoints/cli/main.py @@ -47,7 +47,6 @@ def main(): import vllm.entrypoints.cli.preload import vllm.entrypoints.cli.run_batch import vllm.entrypoints.cli.serve - import vllm.entrypoints.cli.snapshot from vllm.entrypoints.serve.utils.api_utils import ( VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup, @@ -62,11 +61,9 @@ def main(): vllm.entrypoints.cli.download_kernels, vllm.entrypoints.cli.preload, vllm.entrypoints.cli.run_batch, - vllm.entrypoints.cli.snapshot, ] - if sys.argv[1:2] != ["snapshot"]: - cli_env_setup() + cli_env_setup() parser = FlexibleArgumentParser( description="vLLM CLI", @@ -81,12 +78,7 @@ def main(): subparsers = parser.add_subparsers(required=False, dest="subparser") cmds = {} for cmd_module in CMD_MODULES: - if cmd_module is vllm.entrypoints.cli.snapshot: - new_cmds = cmd_module.cmd_init( - create_requested=sys.argv[1:3] == ["snapshot", "create"] - ) - else: - new_cmds = cmd_module.cmd_init() + new_cmds = cmd_module.cmd_init() for cmd in new_cmds: cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd) cmds[cmd.name] = cmd diff --git a/vllm/entrypoints/cli/snapshot.py b/vllm/entrypoints/cli/snapshot.py deleted file mode 100644 index f5201675d4..0000000000 --- a/vllm/entrypoints/cli/snapshot.py +++ /dev/null @@ -1,143 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import argparse -import json -import os -import platform -import typing -from pathlib import Path - -from vllm.entrypoints.cli.types import CLISubcommand - -if typing.TYPE_CHECKING: - from vllm.utils.argparse_utils import FlexibleArgumentParser -else: - FlexibleArgumentParser = argparse.ArgumentParser - - -def validate_create_args(args: argparse.Namespace) -> None: - model = str(getattr(args, "model_tag", None) or args.model) - revision = str(getattr(args, "revision", None) or "") - if ( - Path(model).exists() - or len(revision) != 40 - or any(character not in "0123456789abcdefABCDEF" for character in revision) - ): - raise ValueError( - "snapshot create requires an immutable remote model and " - "40-character --revision" - ) - parallel_sizes = ( - ("tensor parallel", args.tensor_parallel_size), - ("pipeline parallel", args.pipeline_parallel_size), - ("data parallel", args.data_parallel_size), - ("prefill context parallel", args.prefill_context_parallel_size), - ) - for name, size in parallel_sizes: - if size != 1: - raise ValueError(f"snapshot create currently requires {name} size 1") - unsupported_frontend = ( - args.api_server_count not in (None, 1) - or args.grpc - or args.headless - or args.api_key - or os.environ.get("VLLM_API_KEY") - or args.uds - or args.middleware - ) - if unsupported_frontend or any( - ( - args.ssl_keyfile, - args.ssl_certfile, - args.ssl_ca_certs, - args.enable_ssl_refresh, - args.ssl_ciphers, - ) - ): - raise ValueError( - "snapshot create requires one unauthenticated plaintext HTTP server" - ) - if args.logprobs_mode in {"raw_logits", "processed_logits"}: - raise ValueError("snapshot canary requires a log-probability mode") - if any( - getattr(args, name, None) is not None - for name in ("speculative_config", "spec_method", "spec_model", "spec_tokens") - ): - raise ValueError("snapshot create does not support speculative decoding") - if platform.system() != "Linux" or platform.machine() != "x86_64": - raise ValueError("snapshot create requires Linux x86_64") - - -def run_create(args: argparse.Namespace) -> None: - validate_create_args(args) - from vllm.snapshot.controller import create_snapshot - - create_snapshot(args) - - -def run_inspect(args: argparse.Namespace) -> None: - from vllm.snapshot.manifest import inspect_snapshot - - print(json.dumps(inspect_snapshot(Path(args.snapshot_dir)), indent=2)) - - -def run_restore(args: argparse.Namespace) -> None: - from vllm.snapshot.controller import restore_snapshot - - restore_snapshot(args) - - -class SnapshotSubcommand(CLISubcommand): - """The `snapshot` subcommand for the vLLM CLI.""" - - name = "snapshot" - - def __init__(self, *, create_requested: bool = False) -> None: - self.create_requested = create_requested - - @staticmethod - def cmd(args: argparse.Namespace) -> None: - args.snapshot_dispatch(args) - - def subparser_init( - self, subparsers: argparse._SubParsersAction - ) -> FlexibleArgumentParser: - parser = subparsers.add_parser( - self.name, - help="Create, inspect, or restore an initialized vLLM snapshot.", - usage="vllm snapshot [options]", - ) - actions = parser.add_subparsers(required=True, dest="snapshot_action") - - create_parser = actions.add_parser( - "create", help="Create a snapshot from an initialized TP1 engine." - ) - if self.create_requested: - from vllm.entrypoints.launchers.cli_args import make_arg_parser - - create_parser = make_arg_parser(create_parser) - else: - create_parser.add_argument("model_tag", nargs="?") - create_parser.add_argument("--snapshot-dir", required=True) - create_parser.set_defaults(snapshot_dispatch=run_create) - - inspect_parser = actions.add_parser( - "inspect", help="Inspect a snapshot without restoring it." - ) - inspect_parser.add_argument("snapshot_dir") - inspect_parser.set_defaults(snapshot_dispatch=run_inspect) - - restore_parser = actions.add_parser( - "restore", help="Restore a same-host TP1 snapshot." - ) - restore_parser.add_argument("snapshot_dir") - restore_parser.add_argument("--host", default=None) - restore_parser.add_argument("--port", type=int, default=8000) - restore_parser.set_defaults(snapshot_dispatch=run_restore) - - return parser - - -def cmd_init(*, create_requested: bool = False) -> list[CLISubcommand]: - return [SnapshotSubcommand(create_requested=create_requested)] diff --git a/vllm/model_executor/hw_agnostic/__init__.py b/vllm/model_executor/hw_agnostic/__init__.py deleted file mode 100644 index 208f01a7cb..0000000000 --- a/vllm/model_executor/hw_agnostic/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/model_executor/hw_agnostic/custom_op.py b/vllm/model_executor/hw_agnostic/custom_op.py deleted file mode 100644 index ecb37681f2..0000000000 --- a/vllm/model_executor/hw_agnostic/custom_op.py +++ /dev/null @@ -1,314 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import functools -import inspect - -import torch -import torch.nn as nn - -from vllm.config import get_cached_compilation_config -from vllm.logger import init_logger -from vllm.model_executor.utils import maybe_disable_graph_partition -from vllm.platforms import current_platform - -logger = init_logger(__name__) - -# Dictionary of all custom ops (classes, indexed by registered name). -# To check if an op with a name is enabled, call .enabled() on the class. -# Examples: -# - MyOp.enabled() -# - op_registry["my_op"].enabled() -op_registry: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {} -op_registry_oot: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {} - - -def maybe_get_oot_by_class(class_type: type) -> type: - class_name = class_type.__name__ - if class_name in op_registry_oot: - return op_registry_oot[class_name] - return class_type - - -class PluggableLayer(nn.Module): - """Base class for pluggable layers. - - A PluggableLayer is a *module-composing* abstraction: it may instantiate other - ``torch.nn.Module`` objects as sub-layers, and its functionality depends on - these sub-layers following a generalized invocation sequence. Also, it is stateful - and may hold parameters or buffers. - - Unlike :class:`CustomOp`, PluggableLayer does NOT provide per-platform - ``forward_*`` dispatch. Instead, it supports out-of-tree (OOT) replacement - of the entire layer class at instantiation time, allowing customized - initialization and submodule composition. - """ - - def __new__(cls, *args, **kwargs): - try: - layer_class_name = cls.__name__ - except AttributeError: - raise TypeError( - f"Cannot instantiate '{cls.__name__}': its 'name' attribute " - f"was not set, possibly because it was not decorated with " - f"@PluggableLayer.register, or it's the PluggableLayer itself." - ) from None - - if layer_class_name not in op_registry_oot: - layer_cls_to_instantiate = cls - else: - layer_cls_to_instantiate = op_registry_oot[layer_class_name] - logger.debug( - "Instantiating pluggable layer: %s using %s", - layer_class_name, - str(layer_cls_to_instantiate), - ) - return super().__new__(layer_cls_to_instantiate) - - # Decorator to register pluggable layers. - @classmethod - def register(cls, name: str): - def decorator(op_cls): - assert name not in op_registry, f"Duplicate op name: {name}" - op_cls.name = name - op_registry[name] = op_cls - return op_cls - - return decorator - - # Decorator to register out-of-tree(oot) pluggable layers. - # For OOT pluggable layers: - # if in-tree layer class is registered with an oot_custom_layer, - # the oot_custom_layer will be used instead. - @classmethod - def register_oot(cls, _decorated_layer_cls=None, name: str | None = None): - def decorator(layer_cls): - reg_name = name if name is not None else cls.__name__ - assert reg_name not in op_registry_oot, f"Duplicate layer name: {reg_name}" - layer_cls.name = reg_name - op_registry_oot[reg_name] = layer_cls - return layer_cls - - if _decorated_layer_cls is None: - # Called with parentheses: @PluggableLayer.register_oot() - # or @PluggableLayer.register_oot(name="...") - return decorator - elif isinstance(_decorated_layer_cls, type): # Check if it's a class - # Called without parentheses: @PluggableLayer.register_oot - return decorator(_decorated_layer_cls) - else: - raise TypeError("Decorator can only be applied to classes.") - - -class CustomOp(nn.Module): - """Base class for custom ops. - Dispatches the forward method to the appropriate backend. - """ - - def __new__(cls, *args, **kwargs): - try: - op_name = cls.__name__ - except AttributeError: - raise TypeError( - f"Cannot instantiate '{cls.__name__}': its 'name' attribute " - f"was not set, possibly because it was not decorated with " - f"@CustomOp.register, or it's the CustomOp base class itself." - ) from None - - if op_name not in op_registry_oot: - op_cls_to_instantiate = cls - else: - op_cls_to_instantiate = op_registry_oot[op_name] - logger.debug( - "Instantiating custom op: %s using %s", - op_name, - str(op_cls_to_instantiate), - ) - return super().__new__(op_cls_to_instantiate) - - def __init__(self, *, enforce_enable: bool = False, compile_native: bool = False): - super().__init__() - self._enforce_enable = enforce_enable - self._forward_method = self.dispatch_forward(compile_native=compile_native) - - def forward(self, *args, **kwargs): - return self._forward_method(*args, **kwargs) - - def forward_native(self, *args, **kwargs): - """PyTorch-native implementation; OOT plugins override via - :meth:`forward_oot`.""" - raise NotImplementedError - - def forward_oot(self, *args, **kwargs): - # By default, OOT ops fall back to the PyTorch-native implementation. - return self.forward_native(*args, **kwargs) - - def dispatch_forward(self, compile_native: bool): - # The hw-agnostic CustomOp dispatches only between forward_native - # (in-tree, used on every platform) and forward_oot (overridden by - # an OOT plugin via register_oot). - compilation_config = get_cached_compilation_config() - - # NOTE(shen-shanshan): CustomOp object can be enforce enabled, e.g., - # enable device-specific kernels in ViT models when enabling graph - # mode. By default, it will follow the compilation_config to determine - # whether enable itself. - # This enforce_enable mechanism will be removed after we adding a - # separate compilation_config for multi-modal part. - enabled = self._enforce_enable or self.enabled() - if enabled: - compilation_config.enabled_custom_ops.update([self.__class__.name]) - else: - compilation_config.disabled_custom_ops.update([self.__class__.name]) - - if not enabled: - # Compile forward_native to avoid eager torch ops if inside - # opaque torch custom op (e.g. fused_moe, unified_attention, etc.) - return self.maybe_compile(self.forward_native, enable=compile_native) - - if current_platform.is_out_of_tree(): - return self.forward_oot - return self.forward_native - - def maybe_compile(self, fn, *, enable: bool = True): - """Compile fn if compilation enabled. - Useful for CustomOp instances called from within a torch custom op, - meaning the forward call is hidden from the model-level torch.compile. - - NOTE: this does not enable fusion across ops, so opaque custom ops - should still be unwrapped wherever possible. - """ - from vllm.config.compilation import CompilationMode - - # Do not compile if compilation disabled - if not enable: - return fn - - # Do not compile if global compilation disabled - compilation_config = get_cached_compilation_config() - if compilation_config.mode == CompilationMode.NONE: - return fn - - # If eager backend is used, do not compile either - if compilation_config.backend == "eager": - return fn - - compile_options = maybe_disable_graph_partition( - current_platform.simple_compile_backend - ) - backend = current_platform.simple_compile_backend - - dynamic_arg_dims = getattr(self.__class__, "_dynamic_arg_dims", None) - if dynamic_arg_dims is not None: - compiled_fn = torch.compile( - fn, - dynamic=False, - backend=backend, - options=compile_options, - ) - sig = inspect.signature(fn) - - @functools.wraps(fn) - def wrapper(*args, **kwargs): - bound = sig.bind(*args, **kwargs) - bound.apply_defaults() - for name, dims in dynamic_arg_dims.items(): - arg = bound.arguments.get(name) - if arg is not None and isinstance(arg, torch.Tensor): - dims_list = [dims] if isinstance(dims, int) else dims - for d in dims_list: - real_d = arg.ndim + d if d < 0 else d - torch._dynamo.mark_dynamic(arg, real_d) - return compiled_fn(*args, **kwargs) - - return wrapper - - # dynamic=True to avoid recompilations - return torch.compile( - fn, - dynamic=True, - backend=backend, - options=compile_options, - ) - - @classmethod - def enabled(cls) -> bool: - # if no name, then it was not registered - compilation_config = get_cached_compilation_config() - custom_ops = compilation_config.custom_ops - if not hasattr(cls, "name"): - logger.warning_once( - "Custom op %s was not registered, which means it won't appear " - "in the op registry. It will be enabled/disabled based on the " - "global settings.", - cls.__name__, - ) - return CustomOp.default_on() - - enabled = f"+{cls.name}" in custom_ops - disabled = f"-{cls.name}" in custom_ops - assert not (enabled and disabled), f"Cannot enable and disable {cls.name}" - - return (CustomOp.default_on() or enabled) and not disabled - - @staticmethod - def default_on() -> bool: - """Behavior controlled by `CompilationConfig.custom_ops`: On by default if - 'all', off by default if 'none'. - When PyTorch Inductor is used, 'none' is the default value, - otherwise 'all'. - """ - compilation_config = get_cached_compilation_config() - count_none = compilation_config.custom_ops.count("none") - count_all = compilation_config.custom_ops.count("all") - assert count_none + count_all == 1 - - return not count_none > 0 or count_all > 0 - - # Decorator to register custom ops. - @classmethod - def register( - cls, - name: str, - dynamic_arg_dims: dict[str, int | list[int]] | None = None, - ): - def decorator(op_cls): - assert name not in op_registry, f"Duplicate op name: {name}" - op_cls.name = name - op_cls._dynamic_arg_dims = dynamic_arg_dims - op_registry[name] = op_cls - return op_cls - - return decorator - - # Decorator to register out-of-tree(oot) custom ops. - # For OOT custom ops: - # if in-tree layer class is registered with an oot_custom_op layer, - # the oot_custom_op layer will be used instead. - # Example: - # - @UnquantizedFusedMoEMethod.register_oot - # class HPUUnquantizedFusedMoEMethod(UnquantizedFusedMoEMethod) - # or - # - @CustomOP.register_oot(name="UnquantizedFusedMoEMethod") - @classmethod - def register_oot(cls, _decorated_op_cls=None, name: str | None = None): - def decorator(op_cls): - reg_name = name if name is not None else cls.__name__ - assert reg_name not in op_registry_oot, f"Duplicate op name: {reg_name}" - op_cls.name = reg_name - op_registry_oot[reg_name] = op_cls - return op_cls - - if _decorated_op_cls is None: - # Called with parentheses: @CustomOP.register_oot() - # or @CustomOP.register_oot(name="...") - # So, _decorated_op_cls is None. - # We return the actual decorator function. - return decorator - elif isinstance(_decorated_op_cls, type): # Check if it's a class - # Called without parentheses: @CustomOP.register_oot - # The first argument is the class itself. - # We call the 'decorator' function immediately with the class. - return decorator(_decorated_op_cls) - else: - # Handle other unexpected cases if necessary - raise TypeError("Decorator can only be applied to classes.") diff --git a/vllm/model_executor/hw_agnostic/layers/__init__.py b/vllm/model_executor/hw_agnostic/layers/__init__.py deleted file mode 100644 index 208f01a7cb..0000000000 --- a/vllm/model_executor/hw_agnostic/layers/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/model_executor/hw_agnostic/layers/activation.py b/vllm/model_executor/hw_agnostic/layers/activation.py deleted file mode 100644 index c9282b4027..0000000000 --- a/vllm/model_executor/hw_agnostic/layers/activation.py +++ /dev/null @@ -1,27 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import torch -import torch.nn.functional as F - -from vllm.model_executor.hw_agnostic.custom_op import CustomOp - - -@CustomOp.register("silu_and_mul") -class SiluAndMul(CustomOp): - """SwiGLU: ``x -> silu(x[:d]) * x[d:]`` where ``d = x.shape[-1] // 2``.""" - - def forward_native(self, x: torch.Tensor) -> torch.Tensor: - d = x.shape[-1] // 2 - return F.silu(x[..., :d]) * x[..., d:] - - -# Activation-and-mul ops keyed by HF activation name. -_ACTIVATION_AND_MUL_REGISTRY = { - "silu": SiluAndMul, - "swish": SiluAndMul, -} - - -def get_act_and_mul_fn(act_fn_name: str) -> CustomOp: - """Build the hw-agnostic activation-and-mul op named `act_fn_name`.""" - return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name.lower()]() diff --git a/vllm/model_executor/hw_agnostic/layers/layernorm.py b/vllm/model_executor/hw_agnostic/layers/layernorm.py deleted file mode 100644 index d2c854afd3..0000000000 --- a/vllm/model_executor/hw_agnostic/layers/layernorm.py +++ /dev/null @@ -1,90 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import torch -import torch.nn as nn -from torch import Tensor - -from vllm.model_executor.hw_agnostic.custom_op import CustomOp - - -@CustomOp.register("rms_norm") -class RMSNorm(CustomOp): - """``x -> w * x / sqrt(E[x^2] + eps)``. With ``residual``, fuses - ``residual += x`` then RMSNorm and returns ``(normalized, residual)``.""" - - def __init__( - self, - hidden_size: int, - eps: float = 1e-6, - var_hidden_size: int | None = None, - has_weight: bool = True, - dtype: torch.dtype | None = None, - ) -> None: - super().__init__() - self.hidden_size = hidden_size - self.variance_epsilon = eps - self.variance_size_override = ( - None if var_hidden_size == hidden_size else var_hidden_size - ) - weight_dtype = dtype or torch.get_default_dtype() - self.has_weight = has_weight - weight = torch.ones(hidden_size, dtype=weight_dtype) - if has_weight: - self.weight = nn.Parameter(weight) - else: - self.register_buffer("weight", weight, persistent=False) - - def _rms_norm( - self, - x: Tensor, - weight: Tensor | None, - epsilon: float, - variance_size: int | None = None, - ) -> Tensor: - """Weighted root-mean-square layer normalization""" - orig_dtype = x.dtype - x = x.to(torch.float32) - x_var = x if variance_size is None else x[..., :variance_size] - variance = x_var.pow(2).mean(dim=-1, keepdim=True) - x = x * torch.rsqrt(variance + epsilon) - if weight is not None: - x = x.to(weight.dtype) * weight - return x.to(orig_dtype) - - def _fused_add_rms_norm( - self, - x: Tensor, - x_residual: Tensor, - weight: Tensor | None, - epsilon: float, - variance_size: int | None = None, - ) -> tuple[Tensor, Tensor]: - """Fused add and weighted root-mean-square layer normalization""" - orig_dtype = x.dtype - x = x.to(torch.float32) - x = x + x_residual.to(torch.float32) - x_residual = x.to(orig_dtype) - - x_var = x if variance_size is None else x[..., :variance_size] - variance = x_var.pow(2).mean(dim=-1, keepdim=True) - x = x * torch.rsqrt(variance + epsilon) - if weight is not None: - x = x.to(weight.dtype) * weight - return x.to(orig_dtype), x_residual - - def forward_native( - self, - x: torch.Tensor, - residual: torch.Tensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - weight = self.weight if self.has_weight else None - epsilon = self.variance_epsilon - variance_size = self.variance_size_override - - if residual is None: - return self._rms_norm(x, weight, epsilon, variance_size) - else: - return self._fused_add_rms_norm(x, residual, weight, epsilon, variance_size) - - def extra_repr(self) -> str: - return f"hidden_size={self.hidden_size}, eps={self.variance_epsilon}" diff --git a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py deleted file mode 100644 index 7046bfc991..0000000000 --- a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py +++ /dev/null @@ -1,190 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Manual fusion of tensor-parallel all-reduce with the following GemmaRMSNorm. - -Under tensor parallelism a ``RowParallelLinear`` (e.g. attention ``o_proj``) -produces a per-rank partial sum that is all-reduced, and the result is then fed -into a ``GemmaRMSNorm`` that adds the residual and normalizes. flashinfer ships a -kernel that fuses all-reduce + residual-add + RMSNorm into a single launch; this -helper drives it directly (no torch.compile pass) for models that run eager. - -Scope: attention output only, no quantization. When the flashinfer fast path is -not applicable (TP==1, flashinfer/NVSwitch unavailable, unsupported dtype, or an -oversize batch) it falls back to ``all_reduce`` + ``GemmaRMSNorm``, which is -numerically identical to the unfused model path. -""" - -import torch - -from vllm.distributed.communication_op import tensor_model_parallel_all_reduce -from vllm.distributed.parallel_state import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, - get_tp_group, -) -from vllm.model_executor.layers.layernorm import GemmaRMSNorm -from vllm.platforms import current_platform - -MiB = 1024 * 1024 - -# flashinfer fused all-reduce + RMSNorm is wired as a registered custom op in -# allreduce_rms_fusion; both that op and the workspace helpers only exist when -# flashinfer.comm.allreduce_fusion is importable. -try: - from vllm.compilation.passes.fusion.allreduce_rms_fusion import ( - flashinfer_trtllm_fused_allreduce_norm, - ) - from vllm.distributed.device_communicators.flashinfer_all_reduce import ( - flashinfer_comm, - get_fi_ar_workspace, - ) - - _AR_RESIDUAL_RMS_NORM = ( - flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm - if flashinfer_comm is not None - else None - ) -except (ImportError, AttributeError): - flashinfer_trtllm_fused_allreduce_norm = None # type: ignore[assignment] - get_fi_ar_workspace = None # type: ignore[assignment] - _AR_RESIDUAL_RMS_NORM = None - - -_FI_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16) - - -def _max_token_num(tp_size: int, hidden_size: int, dtype: torch.dtype) -> int | None: - """Workspace token budget for flashinfer fused all-reduce, or None if the - current world size / device is unsupported. Mirrors ``FlashInferAllReduce``.""" - from vllm.config.compilation import PassConfig - - max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(tp_size) - if not max_size_mb: - return None - element_size = torch.tensor([], dtype=dtype).element_size() - return int(max_size_mb * MiB) // (hidden_size * element_size) - - -def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool, int]: - """Whether the flashinfer fused path applies; returns (ok, max_token_num).""" - if ( - flashinfer_trtllm_fused_allreduce_norm is None - or get_fi_ar_workspace is None - or _AR_RESIDUAL_RMS_NORM is None - ): - return False, 0 - if ( - not hidden_states.is_cuda - or hidden_states.dim() != 2 - or not hidden_states.is_contiguous() - or hidden_states.dtype not in _FI_SUPPORTED_DTYPES - ): - return False, 0 - - num_tokens, hidden_size = hidden_states.shape - max_token_num = _max_token_num(tp_size, hidden_size, hidden_states.dtype) - if max_token_num is None or num_tokens > max_token_num: - return False, 0 - - # Lazily create / fetch the (globally cached) workspace; returns None on - # GPUs without NVSwitch, in which case we fall back gracefully. - workspace = get_fi_ar_workspace( - world_size=tp_size, - rank=get_tensor_model_parallel_rank(), - max_token_num=max_token_num, - hidden_dim=hidden_size, - dtype=hidden_states.dtype, - group=get_tp_group().cpu_group, - ) - if workspace is None: - return False, 0 - # The token-count bound above uses the whole workspace budget, but a backend - # may use only a fraction of it per call (mnnvl rotates through three Lamport - # buffers). Ask the workspace so we don't admit tensors that the kernel will - # reject with "The buffer size in the given workspace is insufficient". - if not workspace.is_buffer_size_sufficient( - tp_size=tp_size, - num_tokens=num_tokens, - hidden_dim=hidden_size, - dtype=hidden_states.dtype, - ): - return False, 0 - return True, max_token_num - - -def _can_use_aiter_fused_ar_rms(hidden_states: torch.Tensor) -> bool: - if not current_platform.is_rocm(): - return False - from vllm._aiter_ops import rocm_aiter_ops - - if not rocm_aiter_ops.is_custom_all_reduce_enabled(): - return False - if ( - hidden_states.dim() != 2 - or not hidden_states.is_contiguous() - or hidden_states.dtype not in _FI_SUPPORTED_DTYPES - ): - return False - aiter_ar = rocm_aiter_ops.get_aiter_allreduce() - if aiter_ar is None or aiter_ar.disabled: - return False - total_bytes = hidden_states.numel() * hidden_states.element_size() - if total_bytes > aiter_ar.effective_max_size(): - return False - if not aiter_ar.should_custom_ar(hidden_states): - return False - return aiter_ar.use_1stage_fused_ar_rms(hidden_states) - - -def fused_allreduce_gemma_rms_norm( - hidden_states: torch.Tensor, - residual: torch.Tensor, - norm: GemmaRMSNorm, -) -> tuple[torch.Tensor, torch.Tensor]: - """All-reduce ``hidden_states`` + add ``residual`` + GemmaRMSNorm, fused. - - ``hidden_states`` is the per-rank *partial* (un-reduced) output of a - row-parallel linear; ``norm`` is the GemmaRMSNorm applied right after. - Returns ``(normed_output, new_residual)``, equivalent to - ``norm(all_reduce(hidden_states), residual)``. - """ - tp_size = get_tensor_model_parallel_world_size() - if tp_size == 1: - # No all-reduce needed; identical to the unfused path. - return norm(hidden_states, residual) - - ok, max_token_num = _can_use_flashinfer(hidden_states, tp_size) - if ok: - norm_out = torch.empty_like(hidden_states) - # With norm_out provided, the kernel writes the new residual - # (all_reduce(hidden_states) + residual) into the hidden_states buffer - # and the normalized result into norm_out, leaving `residual` untouched. - flashinfer_trtllm_fused_allreduce_norm( - allreduce_in=hidden_states, - residual=residual, - rms_gamma=norm.weight, - rms_eps=norm.variance_epsilon, - world_size=tp_size, - weight_bias=1.0, # GemmaRMSNorm-style - launch_with_pdl=True, - fp32_acc=True, - max_token_num=max_token_num, - pattern_code=_AR_RESIDUAL_RMS_NORM, - norm_out=norm_out, - ) - return norm_out, hidden_states - - if _can_use_aiter_fused_ar_rms(hidden_states): - from vllm._aiter_ops import rocm_aiter_ops - - return rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()( - input_=hidden_states, - residual=residual, - weight=norm.weight, - epsilon=norm.variance_epsilon, - gemma_norm=True, - ) - - # Fallback: explicit all-reduce + GemmaRMSNorm (matches the unfused model). - reduced = tensor_model_parallel_all_reduce(hidden_states) - return norm(reduced, residual) diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py b/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py deleted file mode 100644 index 69a234b081..0000000000 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py +++ /dev/null @@ -1,857 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepGEMM MegaMoE experts: fused EP dispatch, expert GEMMs and combine. - -Each ``MEGA_MOE_*`` format is one row of the SM100 feature matrix. MX weights -and activations carry one UE8M0 scale per 1x32 block. NVFP4 weights carry one -E4M3 scale per 1x16 block plus a per-expert FP32 global scale; NVFP4 dispatch -scales each token by its own FP32 scale. -""" - -import math -from collections.abc import Callable -from dataclasses import dataclass -from inspect import signature -from typing import Any, ClassVar, Literal - -import torch -import torch.nn as nn - -import vllm.envs as envs -from vllm.config import VllmConfig -from vllm.distributed import get_ep_group -from vllm.distributed.eplb.eplb_state import EplbLayerState -from vllm.forward_context import get_forward_context, is_forward_context_available -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.prepare_megamoe import ( - prepare_megamoe_inputs, - prepare_megamoe_nvfp4_inputs, -) -from vllm.model_executor.layers.fused_moe.router.base_router import ( - eplb_map_to_physical_and_record, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - QuantKey, - kMxfp8Dynamic, - kNvfp4DynamicToken, -) -from vllm.model_executor.models.utils import extract_layer_index -from vllm.model_executor.utils import set_weight_attrs -from vllm.platforms import current_platform -from vllm.v1.worker.ubatching import dbo_current_ubatch_id - -logger = init_logger(__name__) - - -@dataclass(frozen=True) -class MegaMoEActivation: - """Gated activation fused into the MegaMoE L1 epilogue. - - ``swiglu`` computes ``gate * sigmoid(alpha * gate) * (up + beta)`` after - clamping gate from above and up symmetrically to ``clamp``; this covers - plain, DeepSeek-V4 (``clamp``) and GPT-OSS style (``alpha``/``beta``) - experts. ``situ`` is Kimi K3's SiTU with gate tanh scale ``alpha`` and - linear tanh scale ``beta``; ``beta=0`` leaves the linear branch unscaled. - """ - - name: Literal["swiglu", "situ"] = "swiglu" - clamp: float | None = None - alpha: float = 1.0 - beta: float = 0.0 - - def __post_init__(self) -> None: - if self.name not in ("swiglu", "situ"): - raise ValueError(f"Unsupported MegaMoE activation {self.name!r}.") - if not (math.isfinite(self.alpha) and math.isfinite(self.beta)): - raise ValueError(f"MegaMoE activation scales must be finite: {self}.") - if self.clamp is not None and not self.clamp >= 0: - raise ValueError(f"MegaMoE activation clamp must be >= 0: {self}.") - if self.name == "situ" and ( - self.clamp is not None or self.alpha <= 0 or self.beta < 0 - ): - raise ValueError( - "SiTU takes no clamp, a positive alpha and a non-negative beta; " - f"got {self}." - ) - - def kernel_kwargs(self) -> dict[str, Any]: - # Defaults stay implicit so DeepGEMM builds without alpha/beta still work. - kwargs: dict[str, Any] = {"activation_clamp": self.clamp} - if self.name != "swiglu": - kwargs["activation"] = self.name - if self.alpha != 1.0: - kwargs["activation_alpha"] = self.alpha - if self.beta != 0.0: - kwargs["activation_beta"] = self.beta - return kwargs - - -@dataclass(frozen=True) -class MegaMoEWeightSpec: - """Storage of one MegaMoE weight in the kernel layout.""" - - dtype: torch.dtype - # Weight values packed per storage element along K (2 for FP4). - values_per_element: int = 1 - # K values per block scale; None for unscaled weights. - scale_group_k: int | None = None - # Block scale dtype: UE8M0 (MX) or E4M3 (NVFP4). - scale_dtype: torch.dtype | None = None - - @property - def has_global_scale(self) -> bool: - """E4M3 block scales pair with an FP32 per-tensor scale.""" - return self.scale_dtype == torch.float8_e4m3fn - - -_MXFP4_WEIGHT = MegaMoEWeightSpec( - torch.int8, - values_per_element=2, - scale_group_k=32, - scale_dtype=torch.float8_e8m0fnu, -) -_MXFP8_WEIGHT = MegaMoEWeightSpec( - torch.float8_e4m3fn, scale_group_k=32, scale_dtype=torch.float8_e8m0fnu -) -_NVFP4_WEIGHT = MegaMoEWeightSpec( - torch.int8, - values_per_element=2, - scale_group_k=16, - scale_dtype=torch.float8_e4m3fn, -) -_BF16_WEIGHT = MegaMoEWeightSpec(torch.bfloat16) - - -@dataclass(frozen=True) -class MegaMoEWeightFormat: - """One routed-expert precision of the MegaMoE feature matrix.""" - - name: str - # DeepGEMM MMA tag; it also keys the symmetric-buffer layout. - mma_type: str - kernel: str - routed: MegaMoEWeightSpec - # Precisions a shared expert may have to run inside the same launch. - shared: tuple[MegaMoEWeightSpec, ...] - # Quantization of the dispatched hidden states; None dispatches BF16. - hidden_quant: QuantKey | None - activations: frozenset[str] - - @property - def uses_nvfp4_api(self) -> bool: - """``nvfp4_mega_moe`` has its own symmetric buffer, BLOCK_M and scales.""" - return self.kernel == "nvfp4_mega_moe" - - -MEGA_MOE_MXFP4 = MegaMoEWeightFormat( - name="mxfp4", - mma_type="fp8xfp4", - kernel="fp8_fp4_mega_moe", - routed=_MXFP4_WEIGHT, - shared=(_MXFP8_WEIGHT,), - hidden_quant=kMxfp8Dynamic, - activations=frozenset({"swiglu", "situ"}), -) -MEGA_MOE_MXFP8 = MegaMoEWeightFormat( - name="mxfp8", - mma_type="fp8xfp8", - kernel="fp8_fp4_mega_moe", - routed=_MXFP8_WEIGHT, - shared=(_MXFP8_WEIGHT,), - hidden_quant=kMxfp8Dynamic, - activations=frozenset({"swiglu", "situ"}), -) -MEGA_MOE_NVFP4 = MegaMoEWeightFormat( - name="nvfp4", - mma_type="nvfp4", - kernel="nvfp4_mega_moe", - routed=_NVFP4_WEIGHT, - shared=(_MXFP8_WEIGHT, _BF16_WEIGHT), - hidden_quant=kNvfp4DynamicToken, - activations=frozenset({"swiglu"}), -) -MEGA_MOE_BF16 = MegaMoEWeightFormat( - name="bf16", - mma_type="bf16xbf16", - kernel="bf16_mega_moe", - routed=_BF16_WEIGHT, - shared=(_BF16_WEIGHT,), - hidden_quant=None, - activations=frozenset({"swiglu"}), -) -MEGA_MOE_WEIGHT_FORMATS = { - fmt.name: fmt - for fmt in (MEGA_MOE_MXFP4, MEGA_MOE_MXFP8, MEGA_MOE_NVFP4, MEGA_MOE_BF16) -} - - -def check_deep_gemm_mega_moe_supported( - device: torch.device, hidden_size: int, intermediate_size: int -) -> None: - """Raise unless DeepGEMM MegaMoE can run on ``device`` with these shapes. - - Raises: - NotImplementedError: If the device architecture is unsupported. - ValueError: If the shapes are unsupported. - - """ - if not current_platform.is_device_capability_family( - 100, device_id=device.index or 0 - ): - raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") - if hidden_size % 128 != 0 or intermediate_size % 128 != 0: - raise ValueError( - "DeepGEMM MegaMoE requires hidden and intermediate sizes " - "to be multiples of 128." - ) - - -def _linear_weight_scale(linear: nn.Module) -> torch.Tensor | None: - scale = getattr(linear, "weight_scale", None) - return getattr(linear, "weight_scale_inv", None) if scale is None else scale - - -def _matches_weight_spec(linear: nn.Module, spec: MegaMoEWeightSpec) -> bool: - """Whether a checkpoint linear can be written into a ``spec`` weight.""" - weight, scale = linear.weight, _linear_weight_scale(linear) - if weight.dtype != spec.dtype or weight.dim() != 2: - return False - if spec.scale_group_k is None: - return scale is None - if scale is None or scale.dim() != 2: - return False - (rows, cols), (scale_rows, scale_cols) = weight.shape, scale.shape - scale_dtypes = {spec.scale_dtype} - if spec.scale_dtype == torch.float8_e8m0fnu: - scale_dtypes.add(torch.uint8) - return ( - scale.dtype in scale_dtypes - and rows % scale_rows == 0 - and cols % scale_cols == 0 - and (cols // scale_cols) % spec.scale_group_k == 0 - ) - - -class DeepGemmMegaMoEExperts(nn.Module): - """Routed experts, and optionally the shared expert, as one MegaMoE launch. - - Parameters are registered in the kernel layout and every checkpoint shard is - written straight into it, so loading (and reloading) needs no post-load - transform. Whether the shared expert runs inside the launch is decided here, - before any weight loads; a fused layer owns the shared-expert weights and - the caller drops its shared MLP. - - NVFP4 experts also take per-expert ``weight_scale_2`` (``w1`` and ``w3`` - must match: one L1 scale per expert) and ``w2`` ``input_scale``; the kernel - quantizes every L2 input with the max of the latter over all loaded - experts. Dispatch uses per-token scales, so ``w1``/``w3`` ``input_scale`` - are not loaded. - - Args: - num_experts: Physical experts across EP ranks, redundant included. - activation: Activation fused into the L1 epilogue. - weight_format: Routed-expert precision, a row of the feature matrix. - shared_experts: Shared MLP to fuse if it is unsharded and in one of - the ``weight_format.shared`` precisions; ``None`` never fuses. - - """ - - _symm_buffers: ClassVar[dict[tuple, Any]] = {} - _synchronized_ep_groups: ClassVar[set[tuple[int, int]]] = set() - - def __init__( - self, - vllm_config: VllmConfig, - *, - num_experts: int, - num_local_experts: int, - experts_start_idx: int, - top_k: int, - hidden_size: int, - intermediate_size: int, - activation: MegaMoEActivation, - weight_format: MegaMoEWeightFormat, - shared_experts: nn.Module | None = None, - num_logical_experts: int | None = None, - fast_math: bool = True, - prefix: str = "", - ): - super().__init__() - if activation.name not in weight_format.activations: - raise NotImplementedError( - f"DeepGEMM MegaMoE {weight_format.name} experts support " - f"{sorted(weight_format.activations)}, not {activation.name!r}." - ) - device = torch.get_default_device() - if device.type != current_platform.device_type: # meta under ipc_cache - device = torch.device( - current_platform.device_type, torch.accelerator.current_device_index() - ) - check_deep_gemm_mega_moe_supported(device, hidden_size, intermediate_size) - - self.prefix = prefix - self.capture_fn: Callable[[torch.Tensor], None] | None = None - self.num_experts = num_experts - self.num_local_experts = num_local_experts - self.experts_start_idx = experts_start_idx - self.experts_end_idx = experts_start_idx + num_local_experts - self.num_logical_experts = ( - num_logical_experts if num_logical_experts is not None else num_experts - ) - self.top_k = top_k - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens - self.weight_format = weight_format - self.activation = activation - self.fast_math = fast_math - self._kernel_kwargs = activation.kernel_kwargs() - self.eplb_state = EplbLayerState() - # L2 input scale of every logical expert; the NVFP4 kernel takes the max. - self._w2_input_scales = torch.zeros(self.num_logical_experts, device="cpu") - self.l2_activation_scale = 1.0 - # The first of each expert's w1/w3 weight_scale_2 until the other loads. - self._pending_l1_scales: dict[int, tuple[str, float]] = {} - - self.num_shared_experts, self.shared_spec = self._fusable_shared_experts( - shared_experts - ) - self._create_weights() - for name, param in self.named_parameters(recurse=False): - shared = name.startswith("shared_") - loader = self.shared_weight_loader if shared else self.weight_loader - set_weight_attrs(param, {"weight_loader": loader}) - if not shared and name.endswith("weight_scale"): - param.quant_method = "block" - - # Register in the static forward context so the custom-op wrapper - # can look up this module by name from within a torch.compile graph. - compilation_config = vllm_config.compilation_config - if prefix in compilation_config.static_forward_context: - raise ValueError(f"Duplicate layer name: {prefix}") - compilation_config.static_forward_context[prefix] = self - - @property - def has_fused_shared_experts(self) -> bool: - return self.num_shared_experts > 0 - - def _fusable_shared_experts( - self, shared_experts: nn.Module | None - ) -> tuple[int, MegaMoEWeightSpec | None]: - """How many shared experts the launch fuses, and in which precision.""" - if shared_experts is None: - return 0, None - gate_up, down = shared_experts.gate_up_proj, shared_experts.down_proj - hidden, inter = self.hidden_size, self.intermediate_size - shared_size = down.weight.shape[-1] - spec = next( - ( - spec - for spec in self.weight_format.shared - if all(_matches_weight_spec(lin, spec) for lin in (gate_up, down)) - ), - None, - ) - if ( - tuple(gate_up.weight.shape) != (2 * shared_size, hidden) - or tuple(down.weight.shape) != (hidden, shared_size) - or shared_size % inter - ): - reason = ( - f"needs an unsharded shared MLP over hidden size {hidden} with a " - f"multiple of {inter} intermediate; got {tuple(gate_up.weight.shape)}" - f" and {tuple(down.weight.shape)}" - ) - elif spec is None: - scale = _linear_weight_scale(gate_up) - supported = ", ".join( - f"{s.dtype}{' with 1x32 UE8M0 scales' if s.scale_group_k else ''}" - for s in self.weight_format.shared - ) - reason = ( - f"{self.weight_format.name} experts fuse {supported} shared " - f"experts; got weight {gate_up.weight.dtype} and scale " - f"{None if scale is None else (tuple(scale.shape), scale.dtype)}" - ) - elif not self._kernel_fuses_shared_experts(spec): - reason = ( - "the installed DeepGEMM Python API is older than the vLLM source. " - "Rebuild the vendored _deep_gemm_C extension to enable it" - ) - else: - return shared_size // inter, spec - logger.warning_once("Disabling MegaMoE shared-expert fusion: %s.", reason) - return 0, None - - def _kernel_fuses_shared_experts(self, spec: MegaMoEWeightSpec) -> bool: - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - fmt = self.weight_format - buffer_api = "get_symm_buffer_for_mega_moe" - block_m_api = "get_block_m_for_mega_moe" - if fmt.uses_nvfp4_api: - buffer_api = "NVFP4SymmBuffer" - block_m_api = "get_block_m_for_nvfp4_mega_moe" - try: - api = { - *signature(getattr(deep_gemm, buffer_api)).parameters, - *signature(getattr(deep_gemm, fmt.kernel)).parameters, - } - except (AttributeError, TypeError, ValueError): - return False - # Scaled shared inputs are laid out by the runtime BLOCK_M. - return { - "num_shared_experts", - "shared_l1_weights", - "shared_l2_weights", - } <= api and (spec.scale_group_k is None or hasattr(deep_gemm, block_m_api)) - - def _create_weights(self) -> None: - """Register the routed and fused shared-expert parameters. - - Routed parameters are indexed by local expert on dim 0, which EPLB moves. - """ - e, h, i = self.num_local_experts, self.hidden_size, self.intermediate_size - fmt = self.weight_format - self._register_weight("w13", (e, 2 * i, h), fmt.routed) - self._register_weight("w2", (e, h, i), fmt.routed) - if fmt.routed.has_global_scale: - # w13_weight_scale_2 is the kernel's L1 alpha (per-token dispatch). - for name in ("w13_weight_scale_2", "w2_weight_scale_2", "w2_input_scale"): - scale = torch.zeros(e, dtype=torch.float32) - self.register_parameter(name, nn.Parameter(scale, requires_grad=False)) - # w2 global weight scale times the L2 activation scale, per launch. - l2_alpha = torch.zeros(e, dtype=torch.float32) - self.register_buffer("_l2_alpha", l2_alpha, persistent=False) - if self.shared_spec is not None: - # The shared expert has the layout of one routed expert slice. - s = i * self.num_shared_experts - self._register_weight("shared_l1", (2 * s, h), self.shared_spec) - self._register_weight("shared_l2", (h, s), self.shared_spec) - - def _register_weight( - self, name: str, shape: tuple[int, ...], spec: MegaMoEWeightSpec - ) -> None: - def register(name: str, data: torch.Tensor) -> None: - self.register_parameter(name, nn.Parameter(data, requires_grad=False)) - - *outer, n, k = shape - register( - f"{name}_weight", - torch.zeros(*outer, n, k // spec.values_per_element, dtype=spec.dtype), - ) - if spec.scale_group_k is not None: - # Block scales, MN-major, four per int32 along K. - scale_k = k // (4 * spec.scale_group_k) - scale = torch.zeros(*outer, scale_k, n, dtype=torch.int32) - register(f"{name}_weight_scale", scale.transpose(-1, -2)) - - def _kernel_weights( - self, name: str - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - weight = getattr(self, f"{name}_weight") - scale = getattr(self, f"{name}_weight_scale", None) - return weight if scale is None else (weight, scale) - - @property - def l1_weights(self) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - return self._kernel_weights("w13") - - @property - def l2_weights(self) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - return self._kernel_weights("w2") - - def _map_global_expert_id(self, expert_id: int) -> list[int]: - """Local slots that logical expert ``expert_id`` occupies on this rank.""" - physical_ids: list[int] = [] - for p in range(self.experts_start_idx, self.experts_end_idx): - if p % self.num_logical_experts == expert_id: - physical_ids.append(p - self.experts_start_idx) - return physical_ids - - def weight_loader( - self, - param: nn.Parameter, - loaded_weight: torch.Tensor, - weight_name: str, - shard_id: str, - expert_id: int, - return_success: bool = False, - ) -> bool | None: - if not loaded_weight.is_meta: - self._record_global_scale(weight_name, loaded_weight, shard_id, expert_id) - local_expert_ids = self._map_global_expert_id(expert_id) - param_prefix = "w2_" if shard_id == "w2" else "w13_" - if not local_expert_ids or param_prefix not in weight_name: - return False if return_success else None - is_global_scale = weight_name.endswith(("weight_scale_2", "input_scale")) - is_scale = weight_name.endswith("weight_scale") - for local_expert_id in local_expert_ids: - if is_global_scale: - param.data[local_expert_id].copy_(loaded_weight.reshape(())) - else: - self._load_weight( - param.data[local_expert_id], loaded_weight, shard_id, is_scale - ) - return True if return_success else None - - def _record_global_scale( - self, weight_name: str, scale: torch.Tensor, shard_id: str, expert_id: int - ) -> None: - """Track NVFP4 global scales the kernel takes one of per launch or expert. - - Every rank sees the scales of all experts, so ranks agree and EPLB can - move experts. - """ - if weight_name.endswith("w2_input_scale"): - self._w2_input_scales[expert_id] = scale.max() - self.l2_activation_scale = float(self._w2_input_scales.max()) or 1.0 - elif weight_name.endswith("w13_weight_scale_2"): - # Compare w1 and w3 of the same load, not the live (maybe stale) slot. - prev = self._pending_l1_scales.pop(expert_id, None) - value = float(scale) - if prev is None or prev[0] == shard_id: - self._pending_l1_scales[expert_id] = (shard_id, value) - elif not math.isclose(value, prev[1], rel_tol=1e-5): - raise ValueError( - f"MegaMoE NVFP4 expert {expert_id}: {prev[0]} weight_scale_2 " - f"{prev[1]} != {shard_id} weight_scale_2 {value}; the kernel " - "takes one L1 scale per expert." - ) - - def load_shared_expert_weight(self, name: str, loaded_weight: torch.Tensor) -> str: - """Load one shared-expert checkpoint tensor; return the parameter name.""" - if ".w2." in name or ".down_proj." in name: - param_name = "shared_l2_weight" - elif ".w1." in name or ".w3." in name: - param_name = "shared_l1_weight" - else: - raise ValueError(f"Unexpected shared-expert tensor: {name}") - if "weight_scale" in name: - param_name += "_scale" - param = getattr(self, param_name) - param.weight_loader(param, loaded_weight, name) - return f"{name.partition('shared_experts.')[0]}experts.{param_name}" - - def shared_weight_loader( - self, param: nn.Parameter, loaded_weight: torch.Tensor, weight_name: str - ) -> None: - shard_id = ( - "w1" if ".w1." in weight_name else "w3" if ".w3." in weight_name else "w2" - ) - is_scale = "weight_scale" in weight_name - self._load_weight(param.data, loaded_weight, shard_id, is_scale) - - def _load_weight( - self, - dst: torch.Tensor, - src: torch.Tensor, - shard_id: str, - is_scale: bool, - ) -> None: - """Write checkpoint shard ``src`` into ``dst`` in the kernel layout. - - ``dst`` is one routed expert's slice or a shared-expert parameter. - """ - if shard_id not in ("w1", "w2", "w3"): - raise ValueError(f"Unsupported expert shard id: {shard_id}") - rows = dst.shape[0] // (1 if shard_id == "w2" else 2) - if is_scale: - src = self._pack_scale(src, rows, dst.shape[1] * 4) - else: - src = src.view(dst.dtype) - if tuple(src.shape) != (rows, dst.shape[1]): - raise ValueError( - f"MegaMoE {shard_id} weight shape {tuple(src.shape)} does not " - f"match the kernel shard ({rows}, {dst.shape[1]})" - ) - dst, src = self._kernel_layout_shard(dst, src, shard_id, is_scale) - dst.copy_(src) - - @staticmethod - def _pack_scale(sf: torch.Tensor, rows: int, cols: int) -> torch.Tensor: - """Expand block scales to the kernel granularity, four per int32 along K.""" - sf = sf.view(torch.uint8) - if rows % sf.shape[0] or cols % sf.shape[1]: - raise ValueError( - f"Scale shape {tuple(sf.shape)} does not tile ({rows}, {cols})" - ) - sf = sf.repeat_interleave(rows // sf.shape[0], 0) - sf = sf.repeat_interleave(cols // sf.shape[1], 1) - return sf.contiguous().clone().view(torch.int32) - - @staticmethod - def _kernel_layout_shard( - dst: torch.Tensor, - src: torch.Tensor, - shard_id: str, - is_scale: bool, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Pair a checkpoint shard with its destination in the MegaMoE layout.""" - k = src.shape[-1] - if shard_id == "w2": - if not is_scale: - return dst, src - # UTCCP: scale row 128*b + 32*j + i -> 128*b + 4*i + j. - return dst.view(-1, 32, 4, k).transpose(1, 2), src.reshape(-1, 4, 32, k) - half = 0 if shard_id == "w1" else 1 - if not is_scale: - # Interleave gate/up by 8 rows: r -> (r // 8) * 16 + 8 * half + r % 8. - return dst.view(-1, 2, 8, k)[:, half], src.reshape(-1, 8, k) - # Interleave + UTCCP: 64b + 16j + 8q + m -> 128b + 64q + 32half + 4m + j. - dst = dst.view(-1, 2, 2, 8, 4, k)[:, :, half].permute(0, 3, 1, 2, 4) - return dst, src.reshape(-1, 4, 2, 8, k) - - def get_symm_buffer(self): - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - group = get_ep_group().device_group - key = ( - id(group), - torch.accelerator.current_device_index(), - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - self.num_shared_experts, - None if self.shared_spec is None else self.shared_spec.dtype, - self.weight_format.mma_type, - self.activation.name, - ) - symm_buffer = self._symm_buffers.get(key) - if symm_buffer is None: - args = ( - group, - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - ) - if self.weight_format.uses_nvfp4_api: - symm_buffer = deep_gemm.NVFP4SymmBuffer( - *args, - num_shared_experts=self.num_shared_experts, - shared_dtype=torch.float8_e4m3fn - if self.shared_spec is None - else self.shared_spec.dtype, - ) - else: - symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( - *args, - num_shared_experts=self.num_shared_experts, - mma_type=self.weight_format.mma_type, - activation=self.activation.name, - ) - self._symm_buffers[key] = symm_buffer - return symm_buffer - - def get_block_m(self, num_tokens: int) -> int: - """BLOCK_M the MegaMoE scheduler runs ``num_tokens`` local tokens with.""" - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - args = ( - get_ep_group().world_size, - self.num_experts, - self.get_symm_buffer().num_max_tokens_per_rank, - num_tokens, - self.top_k, - ) - if self.weight_format.uses_nvfp4_api: - return deep_gemm.get_block_m_for_nvfp4_mega_moe(*args) - return deep_gemm.get_block_m_for_mega_moe(*args, self.weight_format.mma_type) - - def set_eplb_state( - self, - moe_layer_idx: int, - expert_load_view: torch.Tensor, - logical_to_physical_map: torch.Tensor, - logical_replica_count: torch.Tensor, - ) -> None: - self.eplb_state.set_layer_state( - moe_layer_idx, - expert_load_view, - logical_to_physical_map, - logical_replica_count, - ) - - def get_expert_weights(self) -> list[torch.Tensor]: - def _to_eplb_view(name: str, t: torch.Tensor) -> torch.Tensor: - """Return a (num_local_experts, -1) view with contiguous memory layout.""" - assert t.shape[0] == self.num_local_experts - if t.is_contiguous(): - return t.view(self.num_local_experts, -1) - elif t.dim() == 3 and t.stride(1) == 1 and t.stride(2) == t.shape[1]: - # scales have shape (E, M, N) with memory layout (E, N, M) - back = torch.transpose(t, 1, 2) - assert back.is_contiguous() - return back.view(self.num_local_experts, -1) - - raise AssertionError( - f"MegaMoE EPLB {name}: non-contiguous expert tensor with " - f"unexpected layout shape={tuple(t.shape)} " - f"stride={tuple(t.stride())} dtype={t.dtype}" - ) - - return [ - _to_eplb_view(name, param) - for name, param in self.named_parameters(recurse=False) - if not name.startswith("shared_") - ] - - def update_expert_map(self) -> None: - pass - - @property - def layer_id(self) -> int: - return extract_layer_index(self.prefix) - - def _synchronize_first_launch(self) -> None: - ep_group = get_ep_group() - key = (id(ep_group.cpu_group), torch.accelerator.current_device_index()) - if key in self._synchronized_ep_groups: - return - torch.accelerator.synchronize() - torch.distributed.barrier(group=ep_group.cpu_group) - self._synchronized_ep_groups.add(key) - - def forward( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - shared_experts: nn.Module | None = None, - ) -> torch.Tensor: - """Run the routed and shared experts on this rank's tokens. - - Args: - hidden_states: BF16 ``[num_tokens, hidden_size]`` activations. - topk_weights: ``[num_tokens, top_k]`` routing weights. - topk_ids: ``[num_tokens, top_k]`` logical expert ids. - shared_experts: Unfused shared MLP, run on ``hidden_states`` and - added to the output. Must be ``None`` when this layer fused the - shared expert at construction. - - Returns: - BF16 ``[num_tokens, hidden_size]`` routed plus shared expert output. - - """ - if shared_experts is not None and self.has_fused_shared_experts: - raise ValueError( - f"{self.prefix} fuses the shared expert into MegaMoE; " - "pass shared_experts=None." - ) - num_tokens = hidden_states.shape[0] - if num_tokens > self.max_num_tokens: - raise ValueError( - f"MegaMoE got {num_tokens} tokens, but the symmetric buffer was " - f"sized for {self.max_num_tokens}." - ) - self._synchronize_first_launch() - - is_padding = None - if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): - is_padding = get_forward_context().is_padding - if is_padding is not None: - is_padding = is_padding[:num_tokens] - - if self.capture_fn is not None: - self.capture_fn(topk_ids) - - # EPLB: map logical expert IDs to physical replicas and record load. - eplb_state = self.eplb_state - if eplb_state.logical_to_physical_map is not None: - assert eplb_state.expert_load_view is not None - assert eplb_state.logical_replica_count is not None - assert eplb_state.should_record_tensor is not None - if is_padding is not None: - topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids) - topk_ids = eplb_map_to_physical_and_record( - topk_ids=topk_ids, - expert_load_view=eplb_state.expert_load_view, - logical_to_physical_map=eplb_state.logical_to_physical_map, - logical_replica_count=eplb_state.logical_replica_count, - record_enabled=eplb_state.should_record_tensor, - num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ - dbo_current_ubatch_id() - ] - if eplb_state.num_unpadded_tokens_tensors is not None - else None, - ) - - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - fmt = self.weight_format - symm_buffer = self.get_symm_buffer() - shared_x_sf = None - shared_block_m = None - if self.shared_spec is not None and self.shared_spec.scale_group_k: - shared_x_sf = symm_buffer.shared_l1_acts_sf - shared_block_m = self.get_block_m(num_tokens) - kwargs = dict(self._kernel_kwargs) - if fmt.uses_nvfp4_api: - l2_activation_scale = self.l2_activation_scale - prepare_megamoe_nvfp4_inputs( - hidden_states, - topk_weights, - topk_ids, - symm_buffer.x[:num_tokens], - symm_buffer.x_sf[:num_tokens], - symm_buffer.x_scales[:num_tokens], - symm_buffer.topk_idx[:num_tokens], - symm_buffer.topk_weights[:num_tokens], - is_padding=is_padding, - shared_x=symm_buffer.shared_l1_acts[:num_tokens] - if self.has_fused_shared_experts - else None, - shared_x_sf=shared_x_sf, - shared_block_m=shared_block_m, - w2_global_scale=self.w2_weight_scale_2, - l2_alpha=self._l2_alpha, - l2_activation_scale=l2_activation_scale, - ) - kwargs.update( - l1_alpha=self.w13_weight_scale_2, - l2_alpha=self._l2_alpha, - l2_activation_scale=l2_activation_scale, - use_x_scales=True, - ) - else: - quantized = fmt.hidden_quant is not None - prepare_megamoe_inputs( - hidden_states, - topk_weights, - topk_ids, - symm_buffer.x[:num_tokens], - symm_buffer.x_sf[:num_tokens] if quantized else None, - symm_buffer.topk_idx[:num_tokens], - symm_buffer.topk_weights[:num_tokens], - is_padding=is_padding, - shared_x_sf=shared_x_sf, - shared_block_m=shared_block_m, - hidden_quant=fmt.hidden_quant, - ) - - if self.has_fused_shared_experts: - kwargs["shared_l1_weights"] = self._kernel_weights("shared_l1") - kwargs["shared_l2_weights"] = self._kernel_weights("shared_l2") - y = torch.empty_like(hidden_states, dtype=torch.bfloat16) - getattr(deep_gemm, fmt.kernel)( - y, - self.l1_weights, - self.l2_weights, - symm_buffer, - fast_math=self.fast_math, - **kwargs, - ) - if shared_experts is not None: - y += shared_experts(hidden_states) - return y - - -DeepGemmMegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] diff --git a/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py b/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py deleted file mode 100644 index 879c8c6998..0000000000 --- a/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py +++ /dev/null @@ -1,79 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from abc import abstractmethod -from collections.abc import Callable, Mapping -from typing import Any, ClassVar, Generic, TypeAlias, TypeVar - -from vllm.model_executor.warmup.jit_warmup import ( - VllmJitKernel, -) - -DEFAULT_CUTEDSL_COMPILE_OPTIONS = "--enable-tvm-ffi" - -CompileKeyT = TypeVar("CompileKeyT") -CuTeDSLLaunchSpec: TypeAlias = ( - tuple[CompileKeyT, tuple[Any, ...]] - | tuple[CompileKeyT, tuple[Any, ...], Any] - | tuple[ - CompileKeyT, - tuple[Any, ...], - Any, - Callable[[], Any], - ] -) - - -def cutedsl_fake_stream(*, use_tvm_ffi_env_stream: bool = True) -> Any: - import cutlass.cute as cute - - return cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=use_tvm_ffi_env_stream) - - -def compile_cutedsl( - entry: Callable[..., Any], - *args: Any, - options: str = DEFAULT_CUTEDSL_COMPILE_OPTIONS, - use_tvm_ffi_env_stream: bool = True, -) -> Any: - import cutlass.cute as cute - - return cute.compile( - entry, - *args, - cutedsl_fake_stream(use_tvm_ffi_env_stream=use_tvm_ffi_env_stream), - options=options, - ) - - -class VllmCuTeDSLJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]): - """CuTeDSL owner whose compiled executor is shared by warmup and runtime.""" - - kernel: ClassVar[Any] - bind_launch_inputs = False - - @abstractmethod - def warmup_inputs(self, compile_key: CompileKeyT) -> tuple[Any, ...]: - """Return fake arguments that compile one executor specialization.""" - raise NotImplementedError - - def compile(self, compile_key: CompileKeyT) -> None: - if compile_key in self._compiled_cache: - return - self._compiled_cache[compile_key] = compile_cutedsl( - self.kernel(compile_key), - *self.warmup_inputs(compile_key), - ) - - def launch( - self, - launch_spec: CuTeDSLLaunchSpec[CompileKeyT], - _inputs: Mapping[str, Any], - ) -> Any: - # (compile_key, args), optionally followed by output and epilogue. - compile_key, launch_args = launch_spec[:2] - executor = self._get_or_compile(compile_key) - result = executor(*launch_args) - if len(launch_spec) == 4: - return launch_spec[3]() - return launch_spec[2] if len(launch_spec) == 3 else result diff --git a/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py b/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py deleted file mode 100644 index 188f26fdcf..0000000000 --- a/vllm/model_executor/warmup/jit_warmup_tilelang_helper.py +++ /dev/null @@ -1,171 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Compile-only helpers for TileLang JIT warmup.""" - -from __future__ import annotations - -import logging -from abc import abstractmethod -from collections.abc import Callable, Iterator, Mapping, MutableMapping -from contextlib import contextmanager -from dataclasses import dataclass, field -from functools import wraps -from typing import Any, ClassVar, Generic, TypeAlias, TypeVar, cast - -import torch - -from vllm.model_executor.warmup.jit_warmup import VllmJitKernel - -CompileKeyT = TypeVar("CompileKeyT") -# (kernel_args, launch_args), optionally followed by keywords and a result. -TileLangLaunchSpec: TypeAlias = ( - tuple[tuple[Any, ...], tuple[Any, ...]] - | tuple[tuple[Any, ...], tuple[Any, ...], Mapping[str, Any] | None] - | tuple[tuple[Any, ...], tuple[Any, ...], Mapping[str, Any] | None, Any] -) - - -@contextmanager -def _quiet_tilelang_warmup_logs() -> Iterator[None]: - loggers = [ - logging.getLogger("tilelang.cache.kernel_cache"), - logging.getLogger("tilelang.jit.kernel"), - ] - previous_levels = [logger.level for logger in loggers] - try: - for logger in loggers: - logger.setLevel(logging.ERROR) - yield - finally: - for logger, level in zip(loggers, previous_levels): - logger.setLevel(level) - - -@dataclass(frozen=True) -class TileLangWarmupTensor: - """Minimal tensor-like object accepted by TileLang compile(). - - TileLang builds its cache key and TIR from tensor dtype, shape and stride. - This object deliberately has no storage, so compile-only warmup does not - allocate GPU memory and does not launch the kernel. - """ - - dtype: torch.dtype - shape: tuple[int, ...] = (1,) - strides: tuple[int, ...] | None = field(default=None) - - def stride(self) -> tuple[int, ...]: - if self.strides is not None: - return self.strides - - strides: list[int] = [] - stride = 1 - for size in reversed(self.shape): - strides.append(stride) - stride *= size - return tuple(reversed(strides)) - - def new_empty( - self, - shape: tuple[int, ...], - *, - dtype: torch.dtype | None = None, - device: Any = None, - ) -> TileLangWarmupTensor: - return TileLangWarmupTensor(dtype or self.dtype, shape) - - -def make_tilelang_warmup_tensor( - dtype: torch.dtype, - *shape: int, - strides: tuple[int, ...] | None = None, -) -> TileLangWarmupTensor: - return TileLangWarmupTensor(dtype=dtype, shape=tuple(shape), strides=strides) - - -def _tilelang_call_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: - call_kwargs = dict(kwargs) - tune_params = call_kwargs.pop("__tune_params", {}) - if isinstance(tune_params, dict): - call_kwargs.update(tune_params) - return call_kwargs - - -def compile_tilelang(jit_impl: Any, *args: Any, **kwargs: Any) -> None: - """Compile one TileLang specialization and populate its call cache. - - TileLang's ``compile()`` materializes the kernel without launching it. - We also store the compiled kernel in ``_kernel_cache`` using the same - parsed key as ``__call__`` so runtime does not report a cache miss for an - already materialized specialization. - """ - with _quiet_tilelang_warmup_logs(): - compiled = jit_impl.compile(*args, **kwargs) - func = getattr(jit_impl, "func", None) - parse_args = getattr(func, "parse_args", None) - cache = getattr(jit_impl, "_kernel_cache", None) - if not callable(parse_args) or not isinstance(cache, MutableMapping): - return - - key, _ = parse_args(*args, **_tilelang_call_kwargs(kwargs)) - cache[key] = compiled - - -class VllmTileLangJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]): - """TileLang owner whose runtime launch specification is reused for warmup.""" - - kernel: ClassVar[Any] - _warming_key: CompileKeyT | None = None - - @abstractmethod - def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]: - """Return runtime-shaped inputs that reproduce one compile key.""" - raise NotImplementedError - - def compile(self, compile_key: CompileKeyT) -> None: - self._warming_key = compile_key - try: - cast(Callable[..., Any], self)(**self.warmup_inputs(compile_key)) - finally: - self._warming_key = None - - def launch( - self, - jit_impl: Any, - args: tuple[Any, ...], - kwargs: Mapping[str, Any] | None = None, - ) -> Any: - call_kwargs = dict(kwargs or {}) - if self._warming_key is not None: - return compile_tilelang(jit_impl, *args, **call_kwargs) - return jit_impl(*args, **call_kwargs) - - -def kernel_launcher( - call_fn: Callable[..., TileLangLaunchSpec], -) -> Callable[..., Any]: - """Launch TileLang from declarative kernel and runtime argument tuples.""" - - @wraps(call_fn) - def wrapper( - self: VllmTileLangJitKernel[Any], - *args: Any, - **kwargs: Any, - ) -> Any: - launch_spec = call_fn(self, *args, **kwargs) - kernel_args, launch_args = launch_spec[:2] - if self._warming_key is not None and kernel_args not in ( - (), - (self._warming_key,), - ): - raise RuntimeError( - "TileLang warmup inputs produced a different compile key: " - f"expected {self._warming_key!r}, got kernel arguments " - f"{kernel_args!r}" - ) - launch_kwargs = launch_spec[2] if len(launch_spec) > 2 else None - output = launch_spec[3] if len(launch_spec) > 3 else None - result = self.launch(self.kernel(*kernel_args), launch_args, launch_kwargs) - return output if len(launch_spec) > 3 else result - - return wrapper diff --git a/vllm/models/common/ops/fused_allreduce_rms_norm.py b/vllm/models/common/ops/fused_allreduce_rms_norm.py deleted file mode 100644 index 443c8f3eec..0000000000 --- a/vllm/models/common/ops/fused_allreduce_rms_norm.py +++ /dev/null @@ -1,63 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Fused all-reduce + residual-add + RMSNorm for eager model paths. - -This recovers a fusion that vLLM's torch.compile passes would normally do but -that doesn't fire for models running eager (or under a breakable CUDA graph). -""" - -import torch - -from vllm.distributed import ( - get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, -) -from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import ( - _AR_RESIDUAL_RMS_NORM, - _can_use_flashinfer, - flashinfer_trtllm_fused_allreduce_norm, -) -from vllm.model_executor.layers.layernorm import RMSNorm - - -def fused_allreduce_rms_norm( - hidden_states: torch.Tensor, - residual: torch.Tensor, - norm: RMSNorm, -) -> tuple[torch.Tensor, torch.Tensor]: - """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer. - - ``hidden_states`` is the per-rank *partial* output of a row-parallel linear - run with ``reduce_results=False``; ``norm`` is the RMSNorm applied right - after. Returns ``(normed_output, new_residual)``, equivalent to - ``norm(all_reduce(hidden_states), residual)``. Falls back to an explicit - all-reduce + RMSNorm when the flashinfer fast path is unavailable. - """ - tp_size = get_tensor_model_parallel_world_size() - if tp_size == 1: - return norm(hidden_states, residual) - - if flashinfer_trtllm_fused_allreduce_norm is not None: - ok, max_token_num = _can_use_flashinfer(hidden_states, tp_size) - if ok: - norm_out = torch.empty_like(hidden_states) - # With norm_out provided, the kernel writes the new residual - # (all_reduce(hidden_states) + residual) into the hidden_states - # buffer and the normalized result into norm_out. - flashinfer_trtllm_fused_allreduce_norm( - allreduce_in=hidden_states, - residual=residual, - rms_gamma=norm.weight, - rms_eps=norm.variance_epsilon, - world_size=tp_size, - weight_bias=0.0, # standard RMSNorm (Gemma would use 1.0) - launch_with_pdl=True, - fp32_acc=True, - max_token_num=max_token_num, - pattern_code=_AR_RESIDUAL_RMS_NORM, - norm_out=norm_out, - ) - return norm_out, hidden_states - - reduced = tensor_model_parallel_all_reduce(hidden_states) - return norm(reduced, residual) diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index f2f8cbf7c3..b3c0bbb81e 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -195,7 +195,6 @@ class VideoLoader: VIDEO_LOADER_REGISTRY = VideoLoaderRegistry() -VIDEO_LOADER_REGISTRY.register_gpu_codec("deepstream") @VIDEO_LOADER_REGISTRY.register("opencv") diff --git a/vllm/multimodal/video_decoders/__init__.py b/vllm/multimodal/video_decoders/__init__.py index f4302d9e4a..f617b9e964 100644 --- a/vllm/multimodal/video_decoders/__init__.py +++ b/vllm/multimodal/video_decoders/__init__.py @@ -12,24 +12,15 @@ from .base import ( check_frame_pixel_limit, ) -VideoDecoderBackend = Literal["opencv", "torchcodec", "pynvvideocodec", "deepstream"] +VideoDecoderBackend = Literal["opencv", "pynvvideocodec"] _BACKEND_OPTION_DEFAULTS: dict[str, dict[str, Any]] = { "opencv": {}, - "torchcodec": { - "num_ffmpeg_threads": 0, - "seek_mode": "exact", - "device": "cpu", - }, PYNVVIDEOCODEC_VIDEO_BACKEND: { "hw_decoders": PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, "max_width": None, "max_height": None, }, - "deepstream": { - "pool_size": None, - "timeout_sec": 120.0, - }, } diff --git a/vllm/multimodal/video_decoders/deepstream.py b/vllm/multimodal/video_decoders/deepstream.py deleted file mode 100644 index b0fc0f7c42..0000000000 --- a/vllm/multimodal/video_decoders/deepstream.py +++ /dev/null @@ -1,151 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import threading -from typing import Any, ClassVar - -import numpy.typing as npt - -from vllm.logger import init_logger - -from .base import VideoSourceMetadata, VideoTargetMetadata, check_frame_pixel_limit - -logger = init_logger(__name__) - - -def decode_deepstream( - loader_cls, - data: bytes, - target: VideoTargetMetadata, - sampling_kwargs: dict, - *, - pool_size: int | None = None, - timeout_sec: float = 120.0, -) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]: - from nvidia.deepstream_videodecode import probe_metadata - - total_frames, original_fps, duration, _width, _height, codec = probe_metadata(data) - check_frame_pixel_limit(_width, _height) - source = loader_cls._prepare_source( - VideoSourceMetadata(total_frames, original_fps, duration) - ) - frame_idx = loader_cls.compute_frames_index_to_sample( - source=source, target=target, **sampling_kwargs - ) - frames, valid = DeepStreamVideoBackendMixin.decode_indices( - data, - frame_idx, - source, - codec=codec, - pool_size=pool_size, - timeout_sec=timeout_sec, - ) - return frames, source, frame_idx, valid - - -class DeepStreamVideoBackendMixin: - """NVIDIA DeepStream (NVDEC) GPU-decode codec utilities. - - Decoding runs on a shared pool of daemon threads inside one CUDA - context (see the ``nvidia-deepstream-videodecode-cu13`` package). The - container bytes are pushed into an ``appsrc`` GStreamer pipeline, so no - local file path is required — HTTP and base64 sources decode identically - to local files. - - Like the OpenCV mixin, this provides only the codec layer. - Frame *selection* lives in the loader's - ``compute_frames_index_to_sample`` and arrives here as an explicit - list of frame indices. - """ - - # Process-wide lazy decode pool, shared across all DeepStream backends. - _pool: ClassVar[Any] = None - _pool_lock: ClassVar[Any] = None - - @classmethod - def _get_pool(cls, pool_size: int | None = None): - """Lazy-initialize the shared decode pool on first use. - - ``pool_size`` (number of decode worker threads) comes from - ``--media-io-kwargs`` (``{"video": {"pool_size": N}}``); when unset it - defaults to the existing ``VLLM_MEDIA_LOADING_THREAD_COUNT`` so no - DeepStream-specific env var is needed. The pool is a process-wide - singleton, so the first decode's value wins. - """ - if cls._pool is not None: - return cls._pool - if cls._pool_lock is None: - cls._pool_lock = threading.Lock() - with cls._pool_lock: - if cls._pool is not None: - return cls._pool - import os - - from nvidia.deepstream_videodecode import DecodePool - - if pool_size is None: - pool_size = int(os.environ.get("VLLM_MEDIA_LOADING_THREAD_COUNT", 8)) - pool_size = max(1, min(int(pool_size), 16)) - logger.info( - "[DeepStream] initializing decode pool with %d workers", - pool_size, - ) - cls._pool = DecodePool(num_workers=pool_size) - return cls._pool - - @classmethod - def decode_indices( - cls, - data: bytes, - frame_indices: list[int], - source: VideoSourceMetadata, - codec: str = "", - pool_size: int | None = None, - timeout_sec: float = 120.0, - ) -> tuple[npt.NDArray, list[int]]: - """Decode the requested frame indices from raw container bytes. - - The whole stream is decoded; the pool keeps exactly the frames whose - decode-order index is in ``frame_indices`` (1:1, frame-exact) and - sends EOS once the last one is matched. - - ``codec`` (e.g. ``"h264"``/``"hevc"``) lets the pool keep its NVDEC - session warm across same-codec streams and rebuild only on a codec - change. Frames are returned as a CPU NHWC uint8 array so the - upstream multimodal parser sees the same shape as the other - backends. - """ - if not frame_indices: - raise ValueError("DeepStream backend received no frame indices") - - result = cls._get_pool(pool_size).decode( - data, - target_indices=frame_indices, - codec=codec, - max_frames=len(frame_indices), - timeout_sec=timeout_sec, - ) - if result.error: - raise ValueError(f"DeepStream decode failed: {result.error}") - if result.frames is None or result.n_kept == 0: - raise ValueError("DeepStream decode produced no frames") - - valid = frame_indices[: result.n_kept] - # GPU -> CPU NHWC uint8 at the codec boundary (one PCIe copy); keeps - # the array shape identical to the OpenCV backend. Copy into - # PINNED host memory (reused across calls by PyTorch's pinned caching - # allocator) so the D2H runs at full PCIe bandwidth (~13 GB/s) rather - # than the ~1 GB/s pageable path that plain ``.cpu()`` takes — ~12x - # faster for a 1080p x8 frame batch (~46ms -> ~4ms). ``numpy()`` keeps - # the pinned tensor alive via the array's base. - import torch - - gpu = result.frames - if gpu.is_cuda: - host = torch.empty(gpu.shape, dtype=gpu.dtype, pin_memory=True) - host.copy_(gpu, non_blocking=True) - torch.cuda.current_stream().synchronize() - arr = host.numpy() - else: - arr = gpu.numpy() - return arr, valid diff --git a/vllm/multimodal/video_decoders/torchcodec.py b/vllm/multimodal/video_decoders/torchcodec.py deleted file mode 100644 index 346f7b645b..0000000000 --- a/vllm/multimodal/video_decoders/torchcodec.py +++ /dev/null @@ -1,133 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from typing import Literal - -import numpy as np -import numpy.typing as npt -import torch - -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available - -from .base import ( - VideoSourceMetadata, - VideoTargetMetadata, - check_frame_pixel_limit, -) - -try: - from torchcodec.decoders import VideoDecoder -except (ImportError, RuntimeError): - VideoDecoder = PlaceholderModule("torchcodec").placeholder_attr( # type: ignore[assignment] - "decoders.VideoDecoder" - ) - -logger = init_logger(__name__) - - -def decode_torchcodec( - loader_cls, - data: bytes, - target: VideoTargetMetadata, - sampling_kwargs: dict, - *, - num_ffmpeg_threads: int = 0, - seek_mode: Literal["exact", "approximate"] = "exact", - device: str = "cpu", -) -> tuple[npt.NDArray | torch.Tensor, VideoSourceMetadata, list[int], list[int]]: - check_torchcodec_available() - decoder = TorchCodecVideoBackendMixin.make_torchcodec_decoder( - data, - num_ffmpeg_threads=num_ffmpeg_threads, - seek_mode=seek_mode, - device=device, - ) - check_frame_pixel_limit( - decoder.metadata.width or 0, - decoder.metadata.height or 0, - ) - source = loader_cls._prepare_source( - TorchCodecVideoBackendMixin.get_torchcodec_metadata(decoder) - ) - frame_idx = loader_cls.compute_frames_index_to_sample( - source=source, target=target, **sampling_kwargs - ) - frames, valid = TorchCodecVideoBackendMixin.decode_torchcodec_frames( - decoder, frame_idx, device=device - ) - return frames, source, frame_idx, valid - - -class TorchCodecVideoBackendMixin: - """TorchCodec (FFmpeg-backed, PyTorch-native) codec utilities. - - Builds a :class:`~torchcodec.decoders.VideoDecoder` over the in-memory - bytes and extracts the sampled indices with a single batched - ``get_frames_at`` call, while releasing the GIL during decode. - """ - - @staticmethod - def make_torchcodec_decoder( - data: bytes, - *, - num_ffmpeg_threads: int = 0, - seek_mode: Literal["exact", "approximate"] = "exact", - device: str = "cpu", - ) -> "VideoDecoder": - torch_device = torch.device(device) - if torch_device.type == "cuda" and not current_platform.is_cuda(): - raise ValueError( - f"torchcodec video decoding on device {device!r} requires " - "a CUDA-capable platform." - ) - elif torch_device.type not in ("cpu", "cuda"): - raise ValueError( - f"torchcodec video decoding only supports 'cpu' and 'cuda' " - f"devices, got {device!r}." - ) - # NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest - # of the pipeline expects, avoiding a transpose. - return VideoDecoder( - data, - dimension_order="NHWC", - num_ffmpeg_threads=num_ffmpeg_threads, - seek_mode=seek_mode, - device=device, - ) - - @staticmethod - def get_torchcodec_metadata(decoder: "VideoDecoder") -> VideoSourceMetadata: - md = decoder.metadata - total_frames = md.num_frames or 0 - fps = float(md.average_fps) if md.average_fps else 0.0 - duration = float(md.duration_seconds) if md.duration_seconds else 0.0 - if total_frames == 0 and duration > 0 and fps > 0: - total_frames = int(duration * fps) - return VideoSourceMetadata(total_frames, fps, duration) - - @staticmethod - def decode_torchcodec_frames( - decoder: "VideoDecoder", - frame_indices: list[int], - *, - device: str = "cpu", - ) -> tuple[npt.NDArray | torch.Tensor, list[int]]: - """Decode the requested indices in one batched, index-exact call.""" - if not frame_indices: - return np.empty((0,), dtype=np.uint8), [] - # Note: torchcodec releases the GIL for the entire call - batch = decoder.get_frames_at(frame_indices) - frames = batch.data - if ( - torch.device(device).type != "cpu" - and (status := getattr(decoder, "cpu_fallback", None)) is not None - and bool(status) - ): - logger.warning_once( - "torchcodec could not use NVDEC for this video and " - "decoded on CPU instead; check codec support and " - "libnvcuvid." - ) - return frames, list(frame_indices) diff --git a/vllm/profiler/layerwise_profile.py b/vllm/profiler/layerwise_profile.py deleted file mode 100644 index 1393e123fb..0000000000 --- a/vllm/profiler/layerwise_profile.py +++ /dev/null @@ -1,400 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import copy -from collections import defaultdict -from collections.abc import Callable -from dataclasses import asdict, dataclass, field -from typing import Any, Generic, TypeAlias, TypeVar - -from torch._C._autograd import DeviceType, _KinetoEvent, _ProfilerResult -from torch._C._profiler import _EventType, _ExperimentalConfig, _ProfilerEvent -from torch.autograd.profiler import FunctionEvent -from torch.profiler import ProfilerActivity, profile - -from vllm.profiler.utils import ( - TablePrinter, - event_has_module, - event_is_torch_op, - event_module_repr, - event_torch_op_stack_trace, - indent_string, -) -from vllm.utils.import_utils import PlaceholderModule - -try: - import pandas as pd -except ImportError: - pd = PlaceholderModule("pandas") - - -@dataclass -class _ModuleTreeNode: - event: _ProfilerEvent - parent: "_ModuleTreeNode | None" = None - children: list["_ModuleTreeNode"] = field(default_factory=list) - trace: str = "" - - @property - def is_leaf(self): - return self.event.children is None or len(self.event.children) == 0 - - @property - def is_torch_op(self): - return event_is_torch_op(self.event) - - @property - def is_cuda(self): - return ( - self.event.tag == _EventType.Kineto - and self.event.typed[1].device_type == DeviceType.CUDA - ) - - -@dataclass -class SummaryStatsEntry: - name: str - cuda_time_us: float - pct_cuda_time: float - invocations: int - - -@dataclass -class ModelStatsEntry: - name: str - cpu_time_us: float - cuda_time_us: float - pct_cuda_time: float - trace: str - - -StatsEntry: TypeAlias = ModelStatsEntry | SummaryStatsEntry -StatsEntryT = TypeVar("StatsEntryT", bound=StatsEntry) - - -@dataclass -class _StatsTreeNode(Generic[StatsEntryT]): - entry: StatsEntryT - children: list["_StatsTreeNode[StatsEntryT]"] = field(default_factory=list) - parent: "_StatsTreeNode[StatsEntryT] | None" = None - - -@dataclass -class LayerwiseProfileResults(profile): - _kineto_results: _ProfilerResult - _kineto_event_correlation_map: dict[int, list[_KinetoEvent]] = field(init=False) - _event_correlation_map: dict[int, list[FunctionEvent]] = field(init=False) - _module_tree: list[_ModuleTreeNode] = field(init=False) - _model_stats_tree: list[_StatsTreeNode[ModelStatsEntry]] = field(init=False) - _summary_stats_tree: list[_StatsTreeNode[SummaryStatsEntry]] = field(init=False) - - # profile metadata - num_running_seqs: int | None = None - - def __post_init__(self): - self._build_correlation_map() - self._build_module_tree() - self._build_stats_trees() - - def print_model_table(self, column_widths: dict[str, int] | None = None): - _column_widths = dict( - name=60, cpu_time_us=12, cuda_time_us=12, pct_cuda_time=12, trace=60 - ) - if column_widths: - _column_widths.update(**column_widths) - filtered_model_table = [ - (depth, row) - for depth, row in self._flatten_stats_tree(self._model_stats_tree) - if row.cuda_time_us > 0 or row.cpu_time_us > 0 - ] - TablePrinter(ModelStatsEntry, _column_widths).print_table( - self._indent_row_names_based_on_depth( - filtered_model_table, - indent_style=lambda indent: "|" + "-" * indent + " ", - ) - ) - - def print_summary_table(self, column_widths: dict[str, int] | None = None): - _column_widths = dict( - name=80, cuda_time_us=12, pct_cuda_time=12, invocations=15 - ) - if column_widths: - _column_widths.update(**column_widths) - filtered_summary_table = [ - (depth, row) - for depth, row in self._flatten_stats_tree(self._summary_stats_tree) - if row.cuda_time_us > 0 - ] - TablePrinter(SummaryStatsEntry, _column_widths).print_table( - self._indent_row_names_based_on_depth( - filtered_summary_table, - indent_style=lambda indent: "|" + "-" * indent + " ", - ) - ) - - def export_model_stats_table_csv(self, filename: str): - df = pd.DataFrame( - [asdict(row) for _, row in self._flatten_stats_tree(self._model_stats_tree)] - ) - df.to_csv(filename) - - def export_summary_stats_table_csv(self, filename: str): - df = pd.DataFrame( - [ - asdict(row) - for _, row in self._flatten_stats_tree(self._summary_stats_tree) - ] - ) - df.to_csv(filename) - - def convert_stats_to_dict(self) -> dict[str, Any]: - return { - "metadata": {"num_running_seqs": self.num_running_seqs}, - "summary_stats": self._convert_stats_tree_to_dict(self._summary_stats_tree), - "model_stats": self._convert_stats_tree_to_dict(self._model_stats_tree), - } - - @staticmethod - def _indent_row_names_based_on_depth( - depths_rows: list[tuple[int, StatsEntryT]], - indent_style: Callable[[int], str] | str = " ", - ): - indented_rows: list[StatsEntryT] = [] - for depth, row in depths_rows: - if row.cuda_time_us == 0: - continue - indented_row: StatsEntryT = copy.deepcopy(row) - indented_row.name = indent_string(indented_row.name, depth, indent_style) - indented_rows.append(indented_row) - return indented_rows - - def _build_correlation_map(self): - self._kineto_event_correlation_map = defaultdict(list) - for event in self._kineto_results.events(): - self._kineto_event_correlation_map[event.correlation_id()].append(event) - - def _build_module_tree(self): - self._module_tree = [] - event_tree = self._kineto_results.experimental_event_tree() - - def _df_traversal( - event: _ProfilerEvent, curr_node: _ModuleTreeNode | None = None - ): - # For the tensor parallel case for now only look at task 1 - if event.start_tid != 1: - return - - if event_has_module(event): - node = _ModuleTreeNode(event=event, parent=curr_node) - if curr_node: - curr_node.children.append(node) - else: - self._module_tree.append(node) - curr_node = node - - is_leaf = event.children is None or len(event.children) == 0 - if is_leaf and curr_node: - node = _ModuleTreeNode( - event=event, - parent=curr_node, - trace=event_torch_op_stack_trace( - event, until=lambda x: event_has_module(x) - ), - ) - curr_node.children.append(node) - curr_node = node - - for child in event.children: - _df_traversal(child, curr_node) - - for root in event_tree: - _df_traversal(root) - - def _get_kineto_gpu_event(self, node: _ModuleTreeNode): - if node.event.tag != _EventType.Kineto: - return None - correlated_kineto_events = self._kineto_event_correlation_map.get( - node.event.correlation_id, [] - ) - iterator = ( - x - for x in correlated_kineto_events - if x.device_type() == DeviceType.CUDA and x.name() == node.event.name - ) - return next(iterator, None) - - def _cumulative_cuda_time(self, node: _ModuleTreeNode): - """Return cuda time in microseconds.""" - - def _cumulative_cuda_time_recursive(node: _ModuleTreeNode): - if node.is_leaf and (gpu_kineto_event := self._get_kineto_gpu_event(node)): - return gpu_kineto_event.duration_ns() / 1000.0 - else: - cumulative_cuda_time = 0 - for child in node.children: - cumulative_cuda_time += _cumulative_cuda_time_recursive(child) - return cumulative_cuda_time - - return _cumulative_cuda_time_recursive(node) - - def _total_cuda_time(self): - return sum([self._cumulative_cuda_time(root) for root in self._module_tree]) - - def _build_stats_trees(self): - summary_dict: dict[tuple[str, ...], _StatsTreeNode[SummaryStatsEntry]] = {} - total_cuda_time = self._total_cuda_time() - - def pct_cuda_time(cuda_time_us): - return (cuda_time_us / total_cuda_time) * 100 - - def build_summary_stats_tree_df( - node: _ModuleTreeNode, - parent: _StatsTreeNode[SummaryStatsEntry] | None = None, - summary_trace: tuple[str, ...] = (), - ) -> _StatsTreeNode[SummaryStatsEntry] | None: - if event_has_module(node.event): - name = event_module_repr(node.event) - cuda_time_us = self._cumulative_cuda_time(node) - elif gpu_kineto_event := self._get_kineto_gpu_event(node): - name = gpu_kineto_event.name() - cuda_time_us = gpu_kineto_event.duration_ns() / 1000.0 - else: - return None - - summary_trace = summary_trace + (name,) - if summary_trace in summary_dict: - entry = summary_dict[summary_trace].entry - entry.cuda_time_us += cuda_time_us - entry.invocations += 1 - entry.pct_cuda_time = pct_cuda_time(entry.cuda_time_us) - else: - new_node = _StatsTreeNode( - entry=SummaryStatsEntry( - name=name, - cuda_time_us=cuda_time_us, - pct_cuda_time=pct_cuda_time(cuda_time_us), - invocations=1, - ), - parent=parent, - ) - if parent: - parent.children.append(new_node) - summary_dict[summary_trace] = new_node - - for child in node.children: - build_summary_stats_tree_df( - child, summary_dict[summary_trace], summary_trace - ) - - return summary_dict[summary_trace] - - self._summary_stats_tree = [] - for root in self._module_tree: - summary_node = build_summary_stats_tree_df(root) - if summary_node is not None: - self._summary_stats_tree.append(summary_node) - - def build_model_stats_tree_df( - node: _ModuleTreeNode, - parent: _StatsTreeNode[ModelStatsEntry] | None = None, - ) -> _StatsTreeNode[ModelStatsEntry] | None: - if event_has_module( - node.event, - ): - name = event_module_repr(node.event) - cuda_time_us = self._cumulative_cuda_time(node) - cpu_time_us = node.event.duration_time_ns / 1000 - trace = "" - elif gpu_kineto_event := self._get_kineto_gpu_event(node): - name = gpu_kineto_event.name() - cuda_time_us = gpu_kineto_event.duration_ns() / 1000.0 - cpu_time_us = 0 - trace = node.trace - else: - return None - - new_node = _StatsTreeNode( - entry=ModelStatsEntry( - name=name, - cpu_time_us=cpu_time_us, - cuda_time_us=cuda_time_us, - pct_cuda_time=pct_cuda_time(cuda_time_us), - trace=trace, - ), - parent=parent, - ) - if parent: - parent.children.append(new_node) - - for child in node.children: - build_model_stats_tree_df(child, new_node) - - return new_node - - self._model_stats_tree = [] - for root in self._module_tree: - model_node = build_model_stats_tree_df(root) - if model_node is not None: - self._model_stats_tree.append(model_node) - - def _flatten_stats_tree( - self, tree: list[_StatsTreeNode[StatsEntryT]] - ) -> list[tuple[int, StatsEntryT]]: - entries: list[tuple[int, StatsEntryT]] = [] - - def df_traversal(node: _StatsTreeNode[StatsEntryT], depth: int = 0): - entries.append((depth, node.entry)) - for child in node.children: - df_traversal(child, depth=depth + 1) - - for root in tree: - df_traversal(root) - - return entries - - def _convert_stats_tree_to_dict( - self, tree: list[_StatsTreeNode[StatsEntryT]] - ) -> list[dict[str, Any]]: - root_dicts: list[dict[str, Any]] = [] - - def df_traversal( - node: _StatsTreeNode[StatsEntryT], curr_json_list: list[dict[str, Any]] - ): - curr_json_list.append({"entry": asdict(node.entry), "children": []}) - for child in node.children: - df_traversal(child, curr_json_list[-1]["children"]) - - for root in tree: - df_traversal(root, root_dicts) - - return root_dicts - - -class layerwise_profile(profile): - def __init__(self, num_running_seqs: int | None = None): - """Layerwise profile constructor. - - Args: - num_running_seqs (Optional[int], optional): When given, - num_running_seqs will be passed to LayerProfileResults - for metadata update. Defaults to None. - - """ - super().__init__( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - record_shapes=True, - with_stack=True, - with_modules=True, - experimental_config=_ExperimentalConfig(verbose=True), - ) - - self.num_running_seqs = num_running_seqs - - def __enter__(self): - return super().__enter__() - - def __exit__(self, exc_type, exc_val, exc_tb): - super().__exit__(exc_type, exc_val, exc_tb) - self.results = LayerwiseProfileResults( - self.profiler.kineto_results, num_running_seqs=self.num_running_seqs - ) diff --git a/vllm/profiler/utils.py b/vllm/profiler/utils.py deleted file mode 100644 index 5f3c5eaaec..0000000000 --- a/vllm/profiler/utils.py +++ /dev/null @@ -1,141 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import dataclasses -from collections.abc import Callable - -from _typeshed import DataclassInstance -from torch._C._profiler import _EventType, _ProfilerEvent, _TensorMetadata - -# -# String / Print Manipulation -# - - -def trim_string_back(string: str, width: int) -> str: - if len(string) > width: - offset = len(string) - width + 3 - string = string[:-offset] - if len(string) > 3: - string = string + "..." - return string - - -class TablePrinter: - def __init__(self, row_cls: type[DataclassInstance], column_widths: dict[str, int]): - self.row_cls = row_cls - self.fieldnames = [x.name for x in dataclasses.fields(row_cls)] - self.column_widths = column_widths - assert set(self.column_widths.keys()) == set(self.fieldnames) - - def print_table(self, rows: list[DataclassInstance]): - self._print_header() - self._print_line() - for row in rows: - self._print_row(row) - - def _print_header(self): - for i, f in enumerate(self.fieldnames): - last = i == len(self.fieldnames) - 1 - col_width = self.column_widths[f] - print( - trim_string_back(f, col_width).ljust(col_width), - end=" | " if not last else "\n", - ) - - def _print_row(self, row): - assert isinstance(row, self.row_cls) - - for i, f in enumerate(self.fieldnames): - last = i == len(self.fieldnames) - 1 - col_width = self.column_widths[f] - val = getattr(row, f) - - val_str = "" - if isinstance(val, str): - val_str = trim_string_back(val, col_width).ljust(col_width) - elif type(val) in [float, int]: - val_str = f"{float(val):>.2f}".rjust(col_width) - else: - val_str = f"{val}".rjust(col_width) - print(val_str, end=" | " if not last else "\n") - - def _print_line(self): - total_col_width = 0 - for column_width in self.column_widths.values(): - total_col_width += column_width - print("=" * (total_col_width + 3 * (len(self.column_widths) - 1))) - - -def indent_string( - string: str, indent: int, indent_style: Callable[[int], str] | str = " " -) -> str: - if indent: - if isinstance(indent_style, str): - return indent_style * indent + string - else: - return indent_style(indent) + string - else: - return string - - -# -# _ProfilerEvent utils -# - - -def event_has_module(event: _ProfilerEvent) -> bool: - event_type, typed_event = event.typed - if event_type == _EventType.PyCall: - return typed_event.module is not None - return False - - -def event_is_torch_op(event: _ProfilerEvent) -> bool: - return event.tag == _EventType.TorchOp - - -def event_arg_repr(arg) -> str: - if arg is None or type(arg) in [float, int, bool, str]: - return f"{arg}" - elif isinstance(arg, list): - return f"[{', '.join([event_arg_repr(x) for x in arg])}]" - elif isinstance(arg, tuple): - return f"({', '.join([event_arg_repr(x) for x in arg])})" - else: - assert isinstance(arg, _TensorMetadata), f"Unsupported type: {type(arg)}" - sizes_str = ", ".join([str(x) for x in arg.sizes]) - return f"{str(arg.dtype).replace('torch.', '')}[{sizes_str}]" - - -def event_torch_op_repr(event: _ProfilerEvent) -> str: - assert event.tag == _EventType.TorchOp - args_str = ", ".join([event_arg_repr(x) for x in event.typed[1].inputs]) - return f"{event.name}({args_str})".replace("aten::", "") - - -def event_module_repr(event: _ProfilerEvent) -> str: - assert event_has_module(event) - module = event.typed[1].module - if module.parameters and len(module.parameters) > 0: - args_str = ", ".join( - [f"{x[0]}={event_arg_repr(x[1])}" for x in module.parameters] - ) - return f"{module.cls_name}({args_str})" - else: - return module.cls_name - - -def event_torch_op_stack_trace( - curr_event: _ProfilerEvent, until: Callable[[_ProfilerEvent], bool] -) -> str: - trace = "" - curr_event = curr_event.parent - while curr_event and not until(curr_event): - if event_is_torch_op(curr_event): - if len(trace) > 0: - trace += " <- " - trace += event_torch_op_repr(curr_event) - curr_event = curr_event.parent - - return trace diff --git a/vllm/renderers/inkling_encoding.py b/vllm/renderers/inkling_encoding.py deleted file mode 100644 index 38039e4232..0000000000 --- a/vllm/renderers/inkling_encoding.py +++ /dev/null @@ -1,383 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Inkling chat-encoding core. - -Pure implementation of Inkling chat rendering, kept deliberately free of -vLLM imports: it depends only on the :class:`InklingTextTokenizer` protocol -and speaks OpenAI-style message dicts. If a standalone Inkling -input-processing library becomes available, this module is the unit to -swap out — the swap point is marked in ``vllm/renderers/inkling.py``. - -Multimodal parts follow the contract of vLLM's Inkling multimodal processor -(``InklingMultiModalProcessor`` anchors prompt updates on the bare -content-kind marker and inserts the per-patch placeholder run itself): - -* image parts emit only ``<|content_image|>`` — no seed placeholder id; -* audio parts emit ``<|content_audio_input|><|audio_end|>`` — no seed - placeholder id between them. - -``reasoning_effort`` (a float in [0, 0.99], sourced from -``chat_template_kwargs`` only) renders the ``Thinking effort level:`` -system control block after the tool declarations and initial system messages. -""" - -from __future__ import annotations - -import json -from collections.abc import Iterator, Mapping, Sequence -from typing import Any, Protocol - -END_OF_TEXT = "<|endoftext|>" -MESSAGE_USER = "<|message_user|>" -MESSAGE_MODEL = "<|message_model|>" -MESSAGE_SYSTEM = "<|message_system|>" -MESSAGE_TOOL = "<|message_tool|>" -CONTENT_TEXT = "<|content_text|>" -CONTENT_IMAGE = "<|content_image|>" -CONTENT_MODEL_END_SAMPLING = "<|content_model_end_sampling|>" -CONTENT_THINKING = "<|content_thinking|>" -CONTENT_AUDIO_INPUT = "<|content_audio_input|>" -CONTENT_TOOL_ERROR = "<|content_tool_error|>" -CONTENT_XML = "<|content_xml|>" -CONTENT_INVOKE_TOOL_JSON = "<|content_invoke_tool_json|>" -CONTENT_INVOKE_TOOL_TEXT = "<|content_invoke_tool_text|>" -END_MESSAGE = "<|end_message|>" -AUDIO_END = "<|audio_end|>" - -ROLE_MESSAGE_TOKENS: dict[str, str] = { - "user": MESSAGE_USER, - "assistant": MESSAGE_MODEL, - "system": MESSAGE_SYSTEM, - # The developer role folds into the system role on the Inkling wire. - "developer": MESSAGE_SYSTEM, - "tool": MESSAGE_TOOL, -} - -# Alternate vocab spellings per semantic token. The Inkling HF tokenizer -# exposes some semantic slots as ``<|unused_NNNNNN|>`` tokens (notably -# CONTENT_XML = <|unused_200024|>), so ID resolution must try these -# spellings in order. -SPECIAL_TOKEN_SPELLINGS: dict[str, tuple[str, ...]] = { - MESSAGE_USER: (MESSAGE_USER,), - MESSAGE_MODEL: (MESSAGE_MODEL,), - MESSAGE_SYSTEM: (MESSAGE_SYSTEM,), - MESSAGE_TOOL: (MESSAGE_TOOL,), - CONTENT_TEXT: (CONTENT_TEXT,), - CONTENT_IMAGE: (CONTENT_IMAGE,), - CONTENT_MODEL_END_SAMPLING: (CONTENT_MODEL_END_SAMPLING,), - CONTENT_THINKING: (CONTENT_THINKING,), - CONTENT_AUDIO_INPUT: (CONTENT_AUDIO_INPUT,), - CONTENT_XML: (CONTENT_XML, "<|unused_200024|>"), - CONTENT_INVOKE_TOOL_JSON: (CONTENT_INVOKE_TOOL_JSON,), - END_MESSAGE: (END_MESSAGE,), - AUDIO_END: (AUDIO_END,), -} - - -class InklingTextTokenizer(Protocol): - """Structural tokenizer contract required by the renderer.""" - - def encode_text(self, text: str) -> list[int]: ... - - def encode_special(self, token: str) -> int: ... - - -# OpenAI content-part type spellings that mean image / audio (rendering only -# needs the kind, not the bytes — the bytes are handled by the MM processor). -_IMAGE_PART_TYPES = frozenset({"image", "input_image", "image_url"}) -_AUDIO_PART_TYPES = frozenset({"audio", "input_audio", "audio_url"}) - -_MAX_REASONING_EFFORT = 0.99 - - -def render_inkling_messages( - messages: Sequence[Mapping[str, Any]], - tokenizer: InklingTextTokenizer, - *, - add_generation_prompt: bool = True, - tools: Sequence[Mapping[str, Any]] | None = None, - reasoning_effort: float | None = None, -) -> list[int]: - """Render chat messages to Inkling input ids. - - PURE renderer: emits Inkling framing plus bare media markers; media - encoding and placeholder expansion happen later in the MM processor. - ``add_generation_prompt`` appends the assistant turn opener so the - model continues into the response. - """ - input_ids: list[int] = [] - tool_call_id_to_name: dict[str, str] = {} - - # Request-level tools plus per-developer-message tools (Rust renderer - # semantics) are declared in a single leading system block. - all_tools = list(tools or []) - for message in messages: - if message.get("role") == "developer": - all_tools.extend(message.get("tools") or []) - - if all_tools: - _append_message( - input_ids, - tokenizer, - "system", - "xml", - _tool_declare_json(all_tools), - author_name="tool_declare", - ) - - for message in messages: - role = _expect_role(message) - if role not in {"system", "developer"} and reasoning_effort is not None: - _append_reasoning_effort(input_ids, tokenizer, reasoning_effort) - reasoning_effort = None - - if role == "tool": - tool_name = message.get("name") or tool_call_id_to_name.get( - str(message.get("tool_call_id") or ""), "" - ) - _append_message( - input_ids, - tokenizer, - "tool", - "text", - _flatten_text_content(message.get("content")), - author_name=str(tool_name), - ) - continue - - if role == "assistant": - reasoning_content = message.get("reasoning") - if reasoning_content is None: - reasoning_content = message.get("reasoning_content") - if reasoning_content: - if not isinstance(reasoning_content, str): - raise TypeError( - "assistant reasoning_content must be a string for " - "Inkling rendering" - ) - _append_message( - input_ids, - tokenizer, - "assistant", - "thinking", - reasoning_content, - ) - - for kind, text in _iter_render_parts(message.get("content")): - _append_message(input_ids, tokenizer, role, kind, text) - - if role == "assistant": - for tool_call in message.get("tool_calls") or []: - name, args = _tool_call_name_and_args(tool_call) - tool_call_id = _as_mapping(tool_call).get("id") - if tool_call_id: - tool_call_id_to_name[str(tool_call_id)] = name - _append_message( - input_ids, - tokenizer, - "assistant", - "invoke_tool_json", - _tool_call_json(name, args), - author_name=name, - ) - - input_ids.append(tokenizer.encode_special(CONTENT_MODEL_END_SAMPLING)) - - if reasoning_effort is not None: - _append_reasoning_effort(input_ids, tokenizer, reasoning_effort) - - if add_generation_prompt: - input_ids.append(tokenizer.encode_special(MESSAGE_MODEL)) - return input_ids - - -def _append_reasoning_effort( - input_ids: list[int], - tokenizer: InklingTextTokenizer, - reasoning_effort: float, -) -> None: - _append_message( - input_ids, - tokenizer, - "system", - "text", - _thinking_effort_text(reasoning_effort), - ) - - -def _thinking_effort_text(reasoning_effort: Any) -> str: - if isinstance(reasoning_effort, bool) or not isinstance( - reasoning_effort, (int, float) - ): - raise TypeError( - "Inkling reasoning_effort must be a number in [0.0, 0.99], " - f"got {type(reasoning_effort).__name__}" - ) - value = float(reasoning_effort) - if not 0.0 <= value <= _MAX_REASONING_EFFORT: - raise ValueError( - f"Inkling reasoning_effort must be in [0.0, 0.99], got {value}" - ) - - effort_text = f"{value:.2f}".rstrip("0").rstrip(".") - if effort_text in {"0", "-0"}: - effort_text = "0.0" - return f"Thinking effort level: {effort_text}" - - -def _append_message( - input_ids: list[int], - tokenizer: InklingTextTokenizer, - role: str, - kind: str, - text: str, - *, - author_name: str | None = None, -) -> None: - input_ids.append(tokenizer.encode_special(ROLE_MESSAGE_TOKENS[role])) - if author_name: - input_ids.extend(tokenizer.encode_text(author_name)) - - if kind == "text": - input_ids.append(tokenizer.encode_special(CONTENT_TEXT)) - input_ids.extend(tokenizer.encode_text(text)) - elif kind == "image": - # Bare marker only: the MM processor replaces the marker with - # `marker + placeholder * num_patches` (see mm_preprocess.py). - input_ids.append(tokenizer.encode_special(CONTENT_IMAGE)) - elif kind == "audio": - input_ids.append(tokenizer.encode_special(CONTENT_AUDIO_INPUT)) - input_ids.append(tokenizer.encode_special(AUDIO_END)) - elif kind == "thinking": - input_ids.append(tokenizer.encode_special(CONTENT_THINKING)) - input_ids.extend(tokenizer.encode_text(text)) - elif kind == "xml": - input_ids.append(tokenizer.encode_special(CONTENT_XML)) - input_ids.extend(tokenizer.encode_text(text)) - elif kind == "invoke_tool_json": - input_ids.append(tokenizer.encode_special(CONTENT_INVOKE_TOOL_JSON)) - input_ids.extend(tokenizer.encode_text(text)) - else: - raise ValueError(f"unsupported Inkling render part kind: {kind!r}") - - input_ids.append(tokenizer.encode_special(END_MESSAGE)) - - -def _iter_render_parts(content: Any) -> Iterator[tuple[str, str]]: - """Yield (kind, text) per content part: kind in {text, image, audio}.""" - if content is None: - return - if isinstance(content, str): - if content: - yield ("text", content) - return - if not isinstance(content, Sequence) or isinstance(content, (bytes, bytearray)): - raise TypeError("message content must be a string or a sequence of parts") - for part in content: - if isinstance(part, str): - yield ("text", part) - continue - if not isinstance(part, Mapping): - raise TypeError(f"content part must be mapping, got {type(part).__name__}") - ptype = part.get("type") - if ptype in (None, "text", "input_text"): - text = part.get("text", "") - yield ("text", text if isinstance(text, str) else "") - elif ptype in _IMAGE_PART_TYPES: - yield ("image", "") - elif ptype in _AUDIO_PART_TYPES: - yield ("audio", "") - else: - raise ValueError(f"unsupported content part type: {ptype!r}") - - -def _flatten_text_content(content: Any) -> str: - """Flatten a tool-response content (string or text parts) to text.""" - if content is None: - return "" - if isinstance(content, str): - return content - parts: list[str] = [] - for kind, text in _iter_render_parts(content): - if kind != "text": - raise ValueError( - "Inkling tool response content must be text, " - f"got a part of kind {kind!r}" - ) - parts.append(text) - return "".join(parts) - - -def _expect_role(message: Mapping[str, Any]) -> str: - role = message.get("role") - if role not in ROLE_MESSAGE_TOKENS: - raise ValueError( - f"unsupported Inkling message role {role!r}; " - f"expected one of {sorted(ROLE_MESSAGE_TOKENS)}" - ) - return str(role) - - -def _as_mapping(value: Any) -> Mapping[str, Any]: - if isinstance(value, Mapping): - return value - if hasattr(value, "model_dump"): - dumped = value.model_dump() - if isinstance(dumped, Mapping): - return dumped - raise TypeError(f"expected mapping, got {type(value).__name__}") - - -def _canonical_json(value: Any) -> str: - return json.dumps( - _sort_json(value), - ensure_ascii=False, - allow_nan=False, - separators=(",", ":"), - ) - - -def _sort_json(value: Any) -> Any: - if isinstance(value, Mapping): - return {str(key): _sort_json(value[key]) for key in sorted(value)} - if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): - return [_sort_json(item) for item in value] - return value - - -def _tool_declare_json(tools: Sequence[Mapping[str, Any]]) -> str: - tool_specs = [] - for tool_value in tools: - tool = _as_mapping(tool_value) - function = _as_mapping(tool.get("function", {})) - tool_specs.append( - { - "description": function.get("description") or "", - "name": function["name"], - "parameters": function.get("parameters") or {}, - "type": tool.get("type", "function"), - } - ) - return _canonical_json(tool_specs) - - -def _tool_call_name_and_args(tool_call_value: Any) -> tuple[str, Mapping[str, Any]]: - tool_call = _as_mapping(tool_call_value) - function = _as_mapping(tool_call.get("function", {})) - name = function.get("name") - if not isinstance(name, str): - raise TypeError("tool call function name must be a string") - - raw_args = function.get("arguments") or {} - if isinstance(raw_args, str): - args = json.loads(raw_args) if raw_args.strip() else {} - else: - args = raw_args - if not isinstance(args, Mapping): - raise TypeError("tool call function arguments must decode to an object") - return name, args - - -def _tool_call_json(name: str, args: Mapping[str, Any]) -> str: - name_json = json.dumps(name, ensure_ascii=False, allow_nan=False) - return f'{{"name":{name_json},"args":{_canonical_json(args)}}}' diff --git a/vllm/snapshot/__init__.py b/vllm/snapshot/__init__.py deleted file mode 100644 index 8c4fde1282..0000000000 --- a/vllm/snapshot/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Create and restore explicitly compatible vLLM engine snapshots.""" diff --git a/vllm/snapshot/controller.py b/vllm/snapshot/controller.py deleted file mode 100644 index d5564fb863..0000000000 --- a/vllm/snapshot/controller.py +++ /dev/null @@ -1,132 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Create and restore transaction orchestration for local snapshots.""" - -import argparse -import shutil -import sys -from pathlib import Path - -from vllm.snapshot.manifest import ( - read_manifest, - validate_artifact_root, - validate_identity, -) -from vllm.snapshot.runtime import ( - LocalSnapshotTools, - SnapshotCreateError, - SnapshotRestoreError, - _error_detail, -) -from vllm.snapshot.types import Oracle, oracles_match - - -def _child_engine_argv( - args: argparse.Namespace, engine_argv: tuple[str, ...] | None -) -> tuple[str, ...]: - """Build the child argv from the first argv source that is present. - - In order: an explicit ``engine_argv`` from the caller, the process argv - tail when invoked as ``vllm snapshot create ...``, then the bare model from - the parsed args. Whichever source wins, snapshot options are stripped and - sleep mode is appended. - """ - if not engine_argv: - process_argv = tuple(sys.argv[1:]) - if process_argv[:2] == ("snapshot", "create"): - engine_argv = process_argv[2:] - else: - model = getattr(args, "model_tag", None) or getattr(args, "model", None) - if not model: - raise SnapshotCreateError("snapshot create requires a model argument") - engine_argv = (str(model),) - child_argv: list[str] = [] - iterator = iter(engine_argv) - for item in iterator: - option, separator, _ = item.partition("=") - option = option.replace("_", "-") - if len(option) > 2 and "--snapshot-dir".startswith(option): - if not separator: - next(iterator, None) - else: - child_argv.append(item) - if "--enable-sleep-mode" not in child_argv: - child_argv.append("--enable-sleep-mode") - return tuple(child_argv) - - -def create_snapshot( - args: argparse.Namespace, - *, - engine_argv: tuple[str, ...] | None = None, - tools: LocalSnapshotTools | None = None, -) -> None: - target = Path(args.snapshot_dir).absolute() - toolset = tools or LocalSnapshotTools() - toolset.preflight("create", target) - if target.exists() or target.is_symlink(): - raise SnapshotCreateError(f"snapshot target already exists: {target}") - validate_artifact_root(target, creating=True) - target.parent.mkdir(mode=0o700, parents=False, exist_ok=True) - validate_artifact_root(target.parent, creating=False) - target.mkdir(mode=0o700) - validate_artifact_root(target, creating=False) - published = False - root_pid: int | None = None - try: - child_argv = _child_engine_argv(args, engine_argv) - root_pid = toolset.launch_child(target, child_argv) - oracle = toolset.wait_ready(target, root_pid) - inventory = toolset.inventory(root_pid, target) - toolset.dump(target, inventory) - toolset.verify_dead(inventory) - manifest = toolset.make_manifest(args, child_argv, inventory, oracle, target) - toolset.publish(target, manifest) - published = True - except BaseException: - if root_pid is not None: - toolset.abort_create(root_pid) - raise - finally: - # Runs even when abort_create raised, so no partial artifact is left. - if not published and target.exists(): - shutil.rmtree(target) - - -def restore_snapshot( - args: argparse.Namespace, *, tools: LocalSnapshotTools | None = None -) -> None: - artifact = Path(args.snapshot_dir).absolute() - toolset = tools or LocalSnapshotTools() - toolset.preflight("restore", artifact) - manifest = read_manifest(artifact) - validate_identity(manifest, toolset.current_identity(manifest.gpu_uuid)) - - root_pid: int | None = None - try: - root_pid = toolset.restore(artifact, manifest) - toolset.release(artifact, args.host, args.port) - toolset.wait_listener(root_pid, args.host, args.port) - actual = toolset.request_oracle(args.host, args.port, manifest) - expected = Oracle( - token_ids=manifest.oracle_token_ids, - text=manifest.oracle_text, - sampled_token_logprob=manifest.oracle_sampled_token_logprob, - ) - if not oracles_match(expected, actual): - raise SnapshotRestoreError( - f"snapshot oracle mismatch: expected={expected!r}, actual={actual!r}" - ) - toolset.complete_restore(root_pid) - except BaseException as error: - if root_pid is not None: - try: - toolset.cleanup(root_pid) - except BaseException as cleanup_error: - raise SnapshotRestoreError( - f"snapshot restore failed: {_error_detail(error)}; " - f"cleanup failed: {_error_detail(cleanup_error)}" - ) from error - if isinstance(error, SnapshotRestoreError): - raise - raise SnapshotRestoreError(str(error)) from error diff --git a/vllm/snapshot/manifest.py b/vllm/snapshot/manifest.py deleted file mode 100644 index c550ee80c5..0000000000 --- a/vllm/snapshot/manifest.py +++ /dev/null @@ -1,276 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import json -import math -import os -import stat -from pathlib import Path -from typing import Annotated, Any, Literal, Self - -from pydantic import ( - BaseModel, - ConfigDict, - Field, - StrictFloat, - StrictInt, - StrictStr, - ValidationError, - field_validator, - model_validator, -) - - -class SnapshotCompatibilityError(RuntimeError): - """The snapshot identity does not match the requested runtime.""" - - -class SnapshotSecurityError(RuntimeError): - """The snapshot artifact does not meet the private-path contract.""" - - -_MAX_PID = 2**31 - 1 -_PositivePid = Annotated[int, Field(gt=0, le=_MAX_PID)] -_OracleTokenId = Annotated[int, Field(ge=0)] - - -class SnapshotRuntimeIdentity(BaseModel): - model_config = ConfigDict(extra="forbid", frozen=True, strict=True) - - vllm_version: str - python_version: str - torch_version: str - cuda_runtime: str - driver_version: str - criu_version: str - cuda_checkpoint_sha256: str - kernel_release: str - host_id: str - gpu_name: str - gpu_uuid: str - environment: tuple[tuple[str, str], ...] - - -class SnapshotManifest(SnapshotRuntimeIdentity): - schema_version: Literal[1] - boundary: Literal["post-engine-init-reloadable-state-released"] - created_at: str - artifact_bytes: Annotated[int, Field(ge=0)] - model: str - served_model_name: str - model_revision: str - tokenizer_revision: str - engine_argv: tuple[str, ...] - process_tree: Annotated[tuple[_PositivePid, ...], Field(min_length=1)] - cuda_holders: Annotated[tuple[_PositivePid, ...], Field(min_length=1)] - external_cache_files: tuple[tuple[str, str], ...] = () - oracle_token_ids: Annotated[ - tuple[_OracleTokenId, ...], Field(min_length=1, max_length=1) - ] - oracle_text: str - oracle_sampled_token_logprob: float - - @field_validator("schema_version", mode="before") - @classmethod - def _require_schema_version_one(cls, value: Any) -> Any: - if type(value) is not int or value != 1: - raise ValueError("schema version must be integer 1") - return value - - @field_validator("oracle_sampled_token_logprob", mode="before") - @classmethod - def _require_finite_float(cls, value: Any) -> Any: - if type(value) is not float or not math.isfinite(value): - raise ValueError("sampled token log probability must be a finite float") - return value - - @model_validator(mode="after") - def _validate_process_relationships(self) -> Self: - if len(set(self.process_tree)) != len(self.process_tree): - raise ValueError("process_tree contains duplicate PIDs") - if len(set(self.cuda_holders)) != len(self.cuda_holders): - raise ValueError("cuda_holders contains duplicate PIDs") - if not set(self.cuda_holders).issubset(self.process_tree): - raise ValueError("cuda_holders must be a subset of process_tree") - return self - - -class ReleaseMarker(BaseModel): - """Release marker the controller writes; unknown keys are ignored.""" - - release: Literal[True] - host: StrictStr | None = None - port: Annotated[StrictInt, Field(ge=1, le=65535)] - - @field_validator("release", mode="before") - @classmethod - def _require_json_true(cls, value: Any) -> Any: - # A literal accepts 1 as True; the barrier requires JSON true itself. - if value is not True: - raise ValueError("release must be JSON true") - return value - - -class ReadyMarker(BaseModel): - """Ready marker the child writes; unknown keys are ignored.""" - - token_ids: tuple[StrictInt, ...] - text: StrictStr - sampled_token_logprob: StrictFloat | StrictInt - - -def validate_identity( - expected: SnapshotRuntimeIdentity, actual: SnapshotRuntimeIdentity -) -> None: - """Require an exact match for every field that can affect compatibility.""" - # ``expected`` may be a manifest, so compare the identity fields only. - fields = set(SnapshotRuntimeIdentity.model_fields) - expected_values = expected.model_dump(include=fields) - actual_values = actual.model_dump(include=fields) - differing = [ - name - for name in SnapshotRuntimeIdentity.model_fields - if expected_values[name] != actual_values[name] - ] - if differing: - raise SnapshotCompatibilityError(f"snapshot mismatch: {', '.join(differing)}") - - -def _existing_path_chain(path: Path): - current = path - while True: - if current.exists() or current.is_symlink(): - yield current - if current.parent == current: - break - current = current.parent - - -def _is_trusted_sticky_directory(path_stat: os.stat_result) -> bool: - return ( - path_stat.st_uid == 0 - and stat.S_ISDIR(path_stat.st_mode) - and bool(path_stat.st_mode & stat.S_ISVTX) - ) - - -def validate_artifact_root(path: Path, *, creating: bool) -> None: - """Validate that an artifact path cannot be replaced by another user.""" - path = Path(os.path.abspath(path)) - if path.is_symlink(): - raise SnapshotSecurityError(f"snapshot path is a symlink: {path}") - if not creating and not path.exists(): - raise SnapshotSecurityError(f"snapshot path does not exist: {path}") - - effective_uid = os.geteuid() - for component in _existing_path_chain(path): - component_stat = component.lstat() - if stat.S_ISLNK(component_stat.st_mode): - raise SnapshotSecurityError( - f"snapshot path contains a symlink: {component}" - ) - if component_stat.st_uid not in (0, effective_uid): - raise SnapshotSecurityError( - f"snapshot path has an untrusted owner: {component}" - ) - if component_stat.st_mode & ( - stat.S_IWGRP | stat.S_IWOTH - ) and not _is_trusted_sticky_directory(component_stat): - raise SnapshotSecurityError( - f"snapshot path has a group- or world-writable ancestor: {component}" - ) - - if path.exists(): - path_stat = path.lstat() - if not stat.S_ISDIR(path_stat.st_mode): - raise SnapshotSecurityError(f"snapshot path is not a directory: {path}") - if path_stat.st_uid != effective_uid: - raise SnapshotSecurityError( - f"snapshot directory is not owned by the current user: {path}" - ) - if stat.S_IMODE(path_stat.st_mode) & 0o077: - raise SnapshotSecurityError( - f"snapshot directory must have mode 0700: {path}" - ) - - -def _validation_path(error: ValidationError) -> str: - detail = error.errors(include_url=False)[0] - location = detail["loc"] - if not location: - return "JSON" if detail["type"] == "json_invalid" else "root" - return ".".join(map(str, location)) - - -def _fsync_directory(path: Path) -> None: - directory_fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)) - try: - os.fsync(directory_fd) - finally: - os.close(directory_fd) - - -def _write_json_atomic(path: Path, payload: object, *, overwrite: bool = False) -> None: - temporary = path.with_name(f"{path.name}.tmp") - if overwrite: - temporary.unlink(missing_ok=True) - elif path.exists() or path.is_symlink(): - raise FileExistsError(path) - flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0) - descriptor = os.open(temporary, flags, 0o600) - try: - with os.fdopen(descriptor, "wb") as output: - output.write( - json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() - ) - output.flush() - os.fsync(output.fileno()) - os.replace(temporary, path) - _fsync_directory(path.parent) - except BaseException: - temporary.unlink(missing_ok=True) - raise - - -def write_manifest_atomic(path: Path, manifest: SnapshotManifest) -> None: - """Write a new private manifest and make its directory entry durable.""" - path = Path(path) - validate_artifact_root(path, creating=False) - destination = path / "manifest.json" - try: - _write_json_atomic(destination, manifest.model_dump(mode="json")) - except FileExistsError as error: - raise SnapshotSecurityError( - f"manifest already exists: {destination}" - ) from error - - -def read_manifest(path: Path) -> SnapshotManifest: - """Read a snapshot published by atomically installing its manifest.""" - path = Path(path) - validate_artifact_root(path, creating=False) - manifest_path = path / "manifest.json" - if manifest_path.is_symlink(): - raise SnapshotSecurityError(f"manifest is a symlink: {manifest_path}") - flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) - try: - file_descriptor = os.open(manifest_path, flags) - except FileNotFoundError as error: - raise SnapshotSecurityError( - f"manifest does not exist: {manifest_path}" - ) from error - try: - with os.fdopen(file_descriptor, encoding="utf-8") as manifest_file: - payload = manifest_file.read() - except UnicodeDecodeError as error: - raise SnapshotCompatibilityError("invalid snapshot manifest: JSON") from error - try: - return SnapshotManifest.model_validate_json(payload, strict=True) - except ValidationError as error: - raise SnapshotCompatibilityError( - f"invalid snapshot manifest: {_validation_path(error)}" - ) from error - - -def inspect_snapshot(path: Path) -> dict[str, object]: - return read_manifest(path).model_dump(mode="json") diff --git a/vllm/snapshot/runtime.py b/vllm/snapshot/runtime.py deleted file mode 100644 index 83c70e4738..0000000000 --- a/vllm/snapshot/runtime.py +++ /dev/null @@ -1,1318 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Host runtime operations for local snapshots.""" - -import argparse -import ast -import hashlib -import importlib.metadata -import importlib.util -import json -import os -import platform -import select -import shutil -import signal -import socket -import stat -import struct -import subprocess -import sys -import tempfile -import time -import urllib.request -from collections.abc import Iterable -from contextlib import suppress -from dataclasses import dataclass -from datetime import UTC, datetime -from pathlib import Path -from typing import NoReturn - -import regex as re -from pydantic import ValidationError - -from vllm.logger import init_logger -from vllm.snapshot.manifest import ( - ReadyMarker, - SnapshotManifest, - SnapshotRuntimeIdentity, - _fsync_directory, - _validation_path, - _write_json_atomic, - validate_artifact_root, - write_manifest_atomic, -) -from vllm.snapshot.types import Oracle - -logger = init_logger(__name__) - - -class SnapshotCreateError(RuntimeError): - """Snapshot creation did not produce a complete artifact.""" - - -class SnapshotRestoreError(RuntimeError): - """Snapshot restore failed.""" - - -_COMMON_CRIU_OPTIONS = ( - "--shell-job", - "--ext-unix-sk", - "--tcp-established", - "--link-remap", - "--file-locks", -) -_CRIU_DIAGNOSTIC_LIMIT = 8192 -_TCP_TABLES = ( - ("AF_INET", Path("/proc/net/tcp")), - ("AF_INET6", Path("/proc/net/tcp6")), -) - -_CACHE_REMEDY = "recreate the snapshot or restore with the same cache state" -# Recorded instead of a digest for files whose content is volatile by design. -_CACHE_LOCKED = "locked" -# Recorded instead of a digest for a file capture found but could not read. -_CACHE_UNREADABLE = "unreadable" - - -def _generated_cache_prefixes() -> tuple[str, ...]: - """Path prefixes of the generated caches an engine process keeps open. - - vLLM redirects Inductor and Triton under ``VLLM_CACHE_ROOT`` when it - compiles, but a process can still hold their upstream defaults open. - A cache root is resolved and separator-terminated so that it still matches - the kernel-canonical descriptor targets under a symlinked root, and so that - a sibling directory sharing its name does not. Inductor's default location - is a bare prefix instead, because its directory carries a per-user suffix. - """ - import vllm.envs as envs - - inductor_root = os.environ.get("TORCHINDUCTOR_CACHE_DIR") - roots = tuple( - f"{os.path.realpath(os.path.expanduser(root)).rstrip(os.sep)}{os.sep}" - for root in ( - envs.VLLM_CACHE_ROOT, - os.environ.get("TRITON_CACHE_DIR") or "~/.triton/cache", - "~/.cache/flashinfer", - *((inductor_root,) if inductor_root else ()), - ) - ) - if inductor_root: - return roots - return roots + ( - os.path.join(os.path.realpath(tempfile.gettempdir()), "torchinductor_"), - ) - - -def _error_detail(error: BaseException) -> str: - message = str(error) - return f"{type(error).__name__}: {message}" if message else type(error).__name__ - - -def _close_pidfds(handles: Iterable[tuple[int, int]]) -> None: - for _pid, pidfd in handles: - with suppress(OSError): - os.close(pidfd) - - -@dataclass(frozen=True) -class ProcessInventory: - root_pid: int - process_tree: tuple[int, ...] - cuda_holders: tuple[int, ...] - gpu_uuid: str - external_cache_files: tuple[tuple[str, str], ...] = () - - -@dataclass(frozen=True) -class _TcpSocketRecord: - family: str - local_raw: str - remote_raw: str - inode: int - - -def _process_comm(pid: int) -> str: - try: - return (Path("/proc") / str(pid) / "comm").read_text().strip() - except OSError: - return "unknown" - - -def _decode_endpoint(family: str, raw: str) -> str: - """Render a ``/proc/net/tcp`` ``HEXADDR:HEXPORT`` field as ``ip:port``. - - The kernel prints the address as little-endian 32-bit words: one for IPv4, - four in order for IPv6. - """ - hex_address, _, hex_port = raw.partition(":") - port = int(hex_port, 16) - if family == "AF_INET": - packed = struct.pack(" None: - owned = tuple(record for record in records if record.inode in socket_owners) - endpoints = { - (record.family, record.local_raw, record.remote_raw) for record in owned - } - offenders = tuple( - record - for record in owned - if (record.family, record.remote_raw, record.local_raw) not in endpoints - ) - if not offenders: - return - named = "; ".join( - f"pid {socket_owners[record.inode]} " - f"({_process_comm(socket_owners[record.inode])}) " - f"{_decode_endpoint(record.family, record.local_raw)} -> " - f"{_decode_endpoint(record.family, record.remote_raw)}" - for record in offenders[:5] - ) - remaining = f" (+{len(offenders) - 5} more)" if len(offenders) > 5 else "" - raise SnapshotCreateError( - f"snapshot tree has an external established TCP connection: {named}{remaining}" - ) - - -class LocalSnapshotTools: - """Host-owned CRIU and cuda-checkpoint implementation.""" - - def __init__(self) -> None: - self.criu = shutil.which("criu") or "criu" - self.cuda_checkpoint = shutil.which("cuda-checkpoint") or "cuda-checkpoint" - self.nvidia_smi = shutil.which("nvidia-smi") or "nvidia-smi" - plugin_dir = os.environ.get("CRIU_CUDA_PLUGIN_DIR", "") - self.plugin_dir = Path(plugin_dir) if plugin_dir else None - self.shm_dir = Path("/dev/shm") - self.timeout_s = float(os.environ.get("VLLM_SNAPSHOT_TIMEOUT_S", "900")) - if not 0 < self.timeout_s < float("inf"): - raise ValueError("VLLM_SNAPSHOT_TIMEOUT_S must be a positive finite number") - self._children: dict[int, subprocess.Popen[bytes]] = {} - self._restored_processes: dict[int, tuple[tuple[int, int], ...]] = {} - - def _privileged(self) -> list[str]: - return [] if os.geteuid() == 0 else ["sudo", "-n"] - - def _run( - self, - command: list[str], - *, - timeout: float | None = None, - ) -> subprocess.CompletedProcess[str]: - return subprocess.run( - command, - check=True, - capture_output=True, - text=True, - timeout=timeout or self.timeout_s, - ) - - def preflight(self, action: str, artifact: Path) -> None: - if platform.system() != "Linux" or platform.machine() != "x86_64": - raise RuntimeError("snapshot requires Linux x86_64") - if action == "restore" and ( - not callable(getattr(os, "pidfd_open", None)) - or not callable(getattr(signal, "pidfd_send_signal", None)) - ): - raise RuntimeError("snapshot requires Linux pidfd support") - for command in (self.criu, self.cuda_checkpoint, self.nvidia_smi): - if shutil.which(command) is None: - raise RuntimeError(f"snapshot dependency not found: {command}") - if self.plugin_dir is None: - raise RuntimeError("CRIU_CUDA_PLUGIN_DIR is required") - if not (self.plugin_dir / "cuda_plugin.so").is_file(): - raise RuntimeError(f"CRIU CUDA plugin not found in {self.plugin_dir}") - if os.geteuid() != 0: - self._run(["sudo", "-n", "true"], timeout=5) - if action == "restore": - validate_artifact_root(artifact, creating=False) - else: - validate_artifact_root(artifact, creating=True) - - def launch_child( - self, - workdir: Path, - engine_argv: tuple[str, ...], - ) -> int: - command = [ - sys.executable, - "-m", - "vllm.snapshot.server", - "--ready-file", - str(workdir / "ready.json"), - "--release-file", - str(workdir / "release.json"), - "--release-timeout-s", - str(self.timeout_s), - "--", - *engine_argv, - ] - child_environment = os.environ.copy() - child_environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") - child_environment.setdefault("NCCL_IB_DISABLE", "1") - with (workdir / "child.log").open("wb") as log_file: - process = subprocess.Popen( - command, - env=child_environment, - stdin=subprocess.DEVNULL, - stdout=log_file, - stderr=subprocess.STDOUT, - start_new_session=True, - ) - self._children[process.pid] = process - return process.pid - - def wait_ready(self, workdir: Path, root_pid: int) -> Oracle: - ready_file = workdir / "ready.json" - deadline = time.monotonic() + self.timeout_s - waitid = os.waitid # type: ignore[attr-defined] - while time.monotonic() < deadline: - if ready_file.is_file() and ready_file.stat().st_size: - try: - marker = ReadyMarker.model_validate_json(ready_file.read_text()) - except ValidationError as error: - raise SnapshotCreateError( - f"snapshot ready marker invalid: {_validation_path(error)}" - ) from error - return Oracle( - token_ids=marker.token_ids, - text=marker.text, - sampled_token_logprob=marker.sampled_token_logprob, - ) - process = self._children.get(root_pid) - if process is not None and waitid( - os.P_PID, - root_pid, - os.WEXITED | os.WNOHANG | os.WNOWAIT, - ): - log = (workdir / "child.log").read_text(errors="replace") - raise SnapshotCreateError(f"snapshot child exited before ready:\n{log}") - try: - os.kill(root_pid, 0) - except ProcessLookupError as error: - log = (workdir / "child.log").read_text(errors="replace") - raise SnapshotCreateError( - f"snapshot child exited before ready:\n{log}" - ) from error - time.sleep(0.05) - raise SnapshotCreateError("snapshot child did not become ready before timeout") - - def _tree_pids(self, root_pid: int) -> tuple[int, ...]: - rows = self._run(["ps", "-eo", "pid=,ppid="]).stdout.splitlines() - children: dict[int, list[int]] = {} - for row in rows: - pid, parent_pid = map(int, row.split()) - children.setdefault(parent_pid, []).append(pid) - stack = [root_pid] - found: list[int] = [] - while stack: - pid = stack.pop() - if pid in found: - continue - found.append(pid) - stack.extend(children.get(pid, [])) - return tuple(found) - - def _cuda_process_rows(self) -> tuple[str, ...]: - output = self._run( - [ - self.nvidia_smi, - "--query-compute-apps=pid,gpu_uuid", - "--format=csv,noheader,nounits", - ] - ).stdout - return tuple(line for line in output.splitlines() if line.strip()) - - def _cuda_pids(self, rows: tuple[str, ...]) -> tuple[int, ...]: - return tuple(sorted({int(line.split(",", 1)[0].strip()) for line in rows})) - - def _descriptor_targets(self, pid: int) -> tuple[str, ...]: - descriptor_dir = Path("/proc") / str(pid) / "fd" - try: - descriptors = tuple(descriptor_dir.iterdir()) - except FileNotFoundError as error: - raise SnapshotCreateError( - f"snapshot process exited during descriptor inventory: {pid}" - ) from error - targets: list[str] = [] - for descriptor in descriptors: - try: - targets.append(os.readlink(descriptor)) - except FileNotFoundError: - continue - return tuple(targets) - - def _descriptor_inventory( - self, process_tree: tuple[int, ...] - ) -> tuple[tuple[int, ...], dict[int, int]]: - """Collect io_uring holders and each socket inode's owning pid.""" - io_uring_pids: list[int] = [] - socket_owners: dict[int, int] = {} - for pid in process_tree: - targets = self._descriptor_targets(pid) - if "anon_inode:[io_uring]" in targets: - io_uring_pids.append(pid) - for target in targets: - if target.startswith("/dev/infiniband/"): - raise SnapshotCreateError( - "snapshot process has open InfiniBand/RDMA state that " - f"CRIU cannot capture: pid {pid}, {target}. If NCCL owns " - "it, use NCCL_IB_DISABLE=1. Close other RDMA clients " - "before snapshot creation." - ) - if target.startswith("socket:[") and target.endswith("]"): - # Shared descriptors name one holder; any of them locates - # the connection for the operator. - socket_owners.setdefault(int(target[len("socket:[") : -1]), pid) - return tuple(io_uring_pids), socket_owners - - def _tcp_records(self) -> tuple[_TcpSocketRecord, ...]: - records: list[_TcpSocketRecord] = [] - for family, table in _TCP_TABLES: - if not table.exists(): - continue - for row in table.read_text().splitlines()[1:]: - fields = row.split() - if fields[3] != "01": - continue - records.append( - _TcpSocketRecord( - family=family, - local_raw=fields[1], - remote_raw=fields[2], - inode=int(fields[9]), - ) - ) - return tuple(records) - - def _cache_fingerprint(self, path: Path) -> str | None: - """Digest a cache file, or mark the ones whose content is volatile.""" - try: - if not stat.S_ISREG(path.stat().st_mode): - return None - except OSError: - return None - if path.name.endswith(".lock"): - return _CACHE_LOCKED - try: - return self._sha256(path) - except OSError: - # Restore can only check that this one still exists, so record the - # downgrade rather than passing it off as a volatile file. - return _CACHE_UNREADABLE - - def _external_cache_files( - self, process_tree: tuple[int, ...], artifact: Path - ) -> tuple[tuple[str, str], ...]: - """Fingerprint the generated-cache files the captured tree holds open. - - CRIU reopens these by path on restore, so a JIT cache that rotates - between capture and restore fails the restore permanently. Recording - them lets restore fail early and name the file that changed. Files - inside the artifact are excluded: they belong to the snapshot itself. - Deleted targets are excluded too, because CRIU remaps those links. - """ - prefixes = _generated_cache_prefixes() - artifact_prefix = f"{os.path.abspath(artifact)}{os.sep}" - candidates: set[str] = set() - for pid in process_tree: - for target in self._descriptor_targets(pid): - if target.endswith(" (deleted)") or target.startswith(artifact_prefix): - continue - if target.endswith(".lock") or target.startswith(prefixes): - candidates.add(target) - recorded = ( - (path, self._cache_fingerprint(Path(path))) for path in sorted(candidates) - ) - return tuple( - (path, fingerprint) - for path, fingerprint in recorded - if fingerprint is not None - ) - - def inventory(self, root_pid: int, artifact: Path) -> ProcessInventory: - process_tree = self._tree_pids(root_pid) - cuda_rows = self._cuda_process_rows() - cuda_holders = tuple( - pid for pid in self._cuda_pids(cuda_rows) if pid in process_tree - ) - if not cuda_holders: - raise SnapshotCreateError("snapshot tree has no CUDA-holding process") - io_uring_pids, socket_owners = self._descriptor_inventory(process_tree) - if io_uring_pids: - raise SnapshotCreateError( - "snapshot tree owns io_uring state that CRIU cannot dump; set " - "kernel.io_uring_disabled=1 before starting an unprivileged " - "vLLM process, or use =2 to disable it host-wide " - f"(pids: {', '.join(map(str, io_uring_pids))})" - ) - _validate_tcp_connections(self._tcp_records(), socket_owners) - gpu_uuid = self._gpu_uuid_for_pids(cuda_holders, cuda_rows) - return ProcessInventory( - root_pid=root_pid, - process_tree=process_tree, - cuda_holders=cuda_holders, - gpu_uuid=gpu_uuid, - external_cache_files=self._external_cache_files(process_tree, artifact), - ) - - def _gpu_uuid_for_pids(self, pids: tuple[int, ...], rows: tuple[str, ...]) -> str: - wanted = set(pids) - matches: set[str] = set() - try: - for line in rows: - pid, gpu_uuid = (item.strip() for item in line.split(",", 1)) - if int(pid) in wanted: - matches.add(gpu_uuid) - except (TypeError, ValueError) as error: - raise SnapshotCreateError( - "could not identify the GPU used by the snapshot process tree" - ) from error - if len(matches) != 1: - raise SnapshotCreateError( - "snapshot process tree must use exactly one identifiable GPU" - ) - return matches.pop() - - def _criu( - self, - action: str, - artifact: Path, - arguments: list[str], - ) -> None: - if self.plugin_dir is None: - raise RuntimeError("CRIU_CUDA_PLUGIN_DIR is required") - env = os.environ.copy() - env["PATH"] = f"{Path(self.cuda_checkpoint).parent}:{env['PATH']}" - env["VLLM_SNAPSHOT_IGNORED_FD_MAP"] = str(artifact / "ignored-nvidia-fds.tsv") - env["RIG_CRIU_IGNORED_FD_MAP"] = env["VLLM_SNAPSHOT_IGNORED_FD_MAP"] - command = [ - *self._privileged(), - "env", - f"PATH={env['PATH']}", - f"VLLM_SNAPSHOT_IGNORED_FD_MAP={env['VLLM_SNAPSHOT_IGNORED_FD_MAP']}", - f"RIG_CRIU_IGNORED_FD_MAP={env['RIG_CRIU_IGNORED_FD_MAP']}", - self.criu, - action, - "-L", - str(self.plugin_dir), - *arguments, - ] - try: - self._run(command) - except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as error: - self._log_criu_diagnostics(error, action, artifact) - raise - - def _log_criu_diagnostics( - self, - error: subprocess.CalledProcessError | subprocess.TimeoutExpired, - action: str, - artifact: Path, - ) -> None: - def bounded(value: str | bytes) -> str: - if isinstance(value, bytes): - return value[-_CRIU_DIAGNOSTIC_LIMIT:].decode(errors="replace") - return value[-_CRIU_DIAGNOSTIC_LIMIT:] - - for name in ("stdout", "stderr"): - value = getattr(error, name, None) - if value: - logger.error("CRIU %s (tail):\n%s", name, bounded(value)) - - log_path = artifact / "images" / f"{action}.log" - try: - result = self._run( - [ - *self._privileged(), - "tail", - "-c", - str(_CRIU_DIAGNOSTIC_LIMIT), - "--", - str(log_path), - ], - timeout=5, - ) - if result.stdout: - logger.error("CRIU %s.log (tail):\n%s", action, bounded(result.stdout)) - except BaseException as diagnostic_error: - logger.error( - "CRIU %s.log unavailable: %s", - action, - bounded(_error_detail(diagnostic_error)), - ) - - def _record_child_log_size(self, artifact: Path) -> None: - child_log = artifact / "child.log" - size_path = artifact / "child.log.snapshot-size" - payload = f"{child_log.stat().st_size}\n".encode() - descriptor = os.open( - size_path, - os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), - 0o600, - ) - with os.fdopen(descriptor, "wb") as output: - output.write(payload) - output.flush() - os.fsync(output.fileno()) - - def _link_remap_names(self) -> set[str]: - try: - entries = tuple(self.shm_dir.iterdir()) - except OSError as error: - raise SnapshotCreateError( - f"could not inventory CRIU link remaps in {self.shm_dir}: " - f"{_error_detail(error)}" - ) from error - return { - path.name - for path in entries - if path.name.startswith("link_remap.") - and path.name.removeprefix("link_remap.").isdigit() - } - - def _capture_link_remaps(self, artifact: Path, names: set[str]) -> None: - if not names: - return - destination_dir = artifact / "link-remaps" - destination_dir.mkdir(mode=0o700) - for name in sorted(names): - source = self.shm_dir / name - source_fd = os.open( - source, - os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), - ) - try: - metadata = os.fstat(source_fd) - if not stat.S_ISREG(metadata.st_mode) or metadata.st_uid != os.getuid(): - raise SnapshotCreateError(f"invalid CRIU link remap: {source}") - destination_fd = os.open( - destination_dir / name, - os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), - 0o600, - ) - with os.fdopen(destination_fd, "wb") as destination: - with os.fdopen(os.dup(source_fd), "rb") as contents: - shutil.copyfileobj(contents, destination) - destination.flush() - os.fsync(destination.fileno()) - finally: - os.close(source_fd) - source.unlink() - - def _rollback_link_remaps(self, created: dict[str, tuple[int, int]]) -> None: - for name, identity in created.items(): - target = self.shm_dir / name - try: - metadata = target.lstat() - except FileNotFoundError: - continue - if (metadata.st_dev, metadata.st_ino) == identity: - target.unlink() - - def _restore_link_remaps(self, artifact: Path) -> dict[str, tuple[int, int]]: - source_dir = artifact / "link-remaps" - if not source_dir.exists(): - return {} - created: dict[str, tuple[int, int]] = {} - try: - for source in sorted(source_dir.iterdir()): - if ( - not source.name.startswith("link_remap.") - or not source.name.removeprefix("link_remap.").isdigit() - or source.is_symlink() - or not source.is_file() - ): - raise SnapshotRestoreError(f"invalid saved link remap: {source}") - payload = source.read_bytes() - target = self.shm_dir / source.name - try: - descriptor = os.open( - target, - os.O_WRONLY - | os.O_CREAT - | os.O_EXCL - | getattr(os, "O_NOFOLLOW", 0), - 0o600, - ) - except FileExistsError as error: - try: - descriptor = os.open( - target, - os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), - ) - except OSError: - raise SnapshotRestoreError( - f"conflicting CRIU link remap exists: {target}" - ) from error - with os.fdopen(descriptor, "rb") as existing: - metadata = os.fstat(existing.fileno()) - if ( - not stat.S_ISREG(metadata.st_mode) - or metadata.st_uid != os.getuid() - or existing.read() != payload - ): - raise SnapshotRestoreError( - f"conflicting CRIU link remap exists: {target}" - ) from error - continue - metadata = os.fstat(descriptor) - created[source.name] = (metadata.st_dev, metadata.st_ino) - with os.fdopen(descriptor, "wb") as output: - output.write(payload) - output.flush() - os.fsync(output.fileno()) - except BaseException: - self._rollback_link_remaps(created) - raise - return created - - def _reset_child_log(self, artifact: Path) -> None: - size_path = artifact / "child.log.snapshot-size" - descriptor = os.open( - size_path, - os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), - ) - with os.fdopen(descriptor, "r") as source: - size = int(source.read().strip()) - if size < 0: - raise SnapshotRestoreError("captured child log size is invalid") - - child_log = artifact / "child.log" - descriptor = os.open( - child_log, - os.O_WRONLY | getattr(os, "O_NOFOLLOW", 0), - ) - try: - metadata = os.fstat(descriptor) - if not stat.S_ISREG(metadata.st_mode) or metadata.st_size < size: - raise SnapshotRestoreError("captured child log is invalid") - os.ftruncate(descriptor, size) - os.fsync(descriptor) - finally: - os.close(descriptor) - - def dump(self, workdir: Path, inventory: ProcessInventory) -> None: - images = workdir / "images" - images.mkdir(mode=0o700) - fd_map = workdir / "ignored-nvidia-fds.tsv" - fd_map.touch(mode=0o600) - remaps_before = self._link_remap_names() - try: - self._criu( - "dump", - workdir, - [ - "--tree", - str(inventory.root_pid), - "--images-dir", - str(images), - "--log-file", - "dump.log", - *_COMMON_CRIU_OPTIONS, - ], - ) - self._capture_link_remaps( - workdir, - self._link_remap_names() - remaps_before, - ) - except BaseException as primary: - try: - remaps_after = self._link_remap_names() - except BaseException as cleanup_error: - raise SnapshotCreateError( - f"snapshot dump failed: {_error_detail(primary)}; " - "remap cleanup inventory failed: " - f"{_error_detail(cleanup_error)}" - ) from primary - for name in remaps_after - remaps_before: - (self.shm_dir / name).unlink(missing_ok=True) - raise - self._record_child_log_size(workdir) - - def verify_dead(self, inventory: ProcessInventory) -> None: - process = self._children.get(inventory.root_pid) - if process is not None: - process.wait(timeout=10) - for pid in inventory.process_tree: - try: - os.kill(pid, 0) - except ProcessLookupError: - continue - raise SnapshotCreateError(f"process survived CRIU dump: {pid}") - self._children.pop(inventory.root_pid, None) - - def _version(self, command: list[str]) -> str: - try: - result = self._run(command, timeout=10) - except subprocess.CalledProcessError as error: - return (error.stdout or error.stderr).splitlines()[0] - return result.stdout.splitlines()[0] if result.stdout else "unknown" - - def _sha256(self, path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as source: - for chunk in iter(lambda: source.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - def _plugin_identity(self) -> tuple[tuple[str, str], ...]: - if self.plugin_dir is None: - return () - return tuple( - (f"snapshot_plugin_sha256:{path.name}", self._sha256(path)) - for path in sorted(self.plugin_dir.glob("*.so")) - ) - - def _torch_identity(self) -> tuple[str, str]: - torch_version = importlib.metadata.version("torch") - torch_spec = importlib.util.find_spec("torch") - locations = torch_spec and torch_spec.submodule_search_locations - if not locations: - raise RuntimeError("installed torch package could not be located") - version_path = Path(next(iter(locations))) / "version.py" - # Parse the version file rather than importing torch: snapshot create's - # CLI path stays import-light (vllm/entrypoints/cli/snapshot.py). - syntax = ast.parse(version_path.read_text(), filename=str(version_path)) - cuda_runtime = "" - for statement in syntax.body: - target = None - value = None - if isinstance(statement, ast.Assign) and len(statement.targets) == 1: - target = statement.targets[0] - value = statement.value - elif isinstance(statement, ast.AnnAssign): - target = statement.target - value = statement.value - if isinstance(target, ast.Name) and target.id == "cuda" and value: - parsed = ast.literal_eval(value) - cuda_runtime = str(parsed or "") - break - return torch_version, cuda_runtime - - def _gpu_identity(self, gpu_uuid: str) -> tuple[str, str, str]: - output = self._run( - [ - self.nvidia_smi, - "--query-gpu=name,uuid,driver_version", - "--format=csv,noheader,nounits", - f"--id={gpu_uuid}", - ] - ).stdout.strip() - name, uuid, driver = (item.strip() for item in output.split(",", 2)) - return name, uuid, driver - - def _artifact_bytes(self, workdir: Path) -> int: - return sum(path.stat().st_size for path in workdir.rglob("*") if path.is_file()) - - def _environment_identity(self) -> tuple[tuple[str, str], ...]: - prefixes = ("VLLM_", "CUDA_", "NCCL_", "TORCH_", "TRITON_") - environment = os.environ.copy() - environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") - environment.setdefault("NCCL_IB_DISABLE", "1") - selected = tuple( - sorted( - (key, hashlib.sha256(key.encode() + b"\0" + value.encode()).hexdigest()) - for key, value in environment.items() - if key.startswith(prefixes) - and key not in {"VLLM_API_KEY", "VLLM_SNAPSHOT_TIMEOUT_S"} - ) - ) - return selected + self._plugin_identity() - - def make_manifest( - self, - args: argparse.Namespace, - engine_argv: tuple[str, ...], - inventory: ProcessInventory, - oracle: Oracle, - workdir: Path, - ) -> SnapshotManifest: - identity = self.current_identity(inventory.gpu_uuid) - redacted_argv = list(engine_argv) - redact_api_keys = False - for index, argument in enumerate(redacted_argv): - option, separator, _ = argument.partition("=") - normalized_option = option.replace("_", "-") - # Parsed argv can contain unambiguous long-option abbreviations. - if len(normalized_option) > 2 and "--api-key".startswith(normalized_option): - redact_api_keys = not separator - if separator: - redacted_argv[index] = f"{option}=***" - elif len(normalized_option) > 2 and "--hf-token".startswith( - normalized_option - ): - if separator: - redacted_argv[index] = f"{option}=***" - elif index + 1 < len(redacted_argv): - redacted_argv[index + 1] = "***" - redact_api_keys = False - elif argument.startswith("-"): - redact_api_keys = False - elif redact_api_keys: - redacted_argv[index] = "***" - engine_argv = tuple(redacted_argv) - revision = str(getattr(args, "revision", None) or "") - tokenizer_revision = str(getattr(args, "tokenizer_revision", None) or revision) - source_model = str(getattr(args, "model_tag", None) or args.model) - from vllm.config.model import get_served_model_name - - served_model_name = get_served_model_name( - source_model, getattr(args, "served_model_name", None) - ) - return SnapshotManifest( - schema_version=1, - boundary="post-engine-init-reloadable-state-released", - created_at=datetime.now(UTC).isoformat(), - artifact_bytes=self._artifact_bytes(workdir), - **identity.model_dump(), - model=source_model, - served_model_name=str(served_model_name), - model_revision=revision, - tokenizer_revision=tokenizer_revision, - engine_argv=engine_argv, - process_tree=inventory.process_tree, - cuda_holders=inventory.cuda_holders, - external_cache_files=inventory.external_cache_files, - oracle_token_ids=oracle.token_ids, - oracle_text=oracle.text, - oracle_sampled_token_logprob=oracle.sampled_token_logprob, - ) - - def publish(self, workdir: Path, manifest: SnapshotManifest) -> None: - write_manifest_atomic(workdir, manifest) - _fsync_directory(workdir.parent) - - def current_identity(self, gpu_uuid: str) -> SnapshotRuntimeIdentity: - gpu_name, gpu_uuid, driver_version = self._gpu_identity(gpu_uuid) - torch_version, cuda_runtime = self._torch_identity() - host_id = Path("/etc/machine-id").read_text().strip() - return SnapshotRuntimeIdentity( - vllm_version=importlib.metadata.version("vllm"), - python_version=platform.python_version(), - torch_version=torch_version, - cuda_runtime=cuda_runtime, - driver_version=driver_version, - criu_version=self._version([self.criu, "--version"]), - cuda_checkpoint_sha256=self._sha256(Path(self.cuda_checkpoint)), - kernel_release=platform.release(), - host_id=host_id, - gpu_name=gpu_name, - gpu_uuid=gpu_uuid, - environment=self._environment_identity(), - ) - - def _read_restored_pid(self, pidfile: Path) -> int: - try: - payload = self._run([*self._privileged(), "cat", str(pidfile)]).stdout - except Exception as error: - raise SnapshotRestoreError("restored PID file is missing") from error - match = re.fullmatch(r"([1-9][0-9]{0,9})\n?", payload) - if match is None: - raise SnapshotRestoreError("restored PID file is invalid") - restored_pid = int(match.group(1)) - if restored_pid > 2**31 - 1: - raise SnapshotRestoreError("restored PID file is invalid") - return restored_pid - - def _process_state(self, pid: int) -> tuple[int, int, int, int]: - payload = (Path("/proc") / str(pid) / "stat").read_text() - command_end = payload.rfind(")") - fields = payload[command_end + 2 :].split() - if command_end < 0 or len(fields) < 20: - raise ValueError(f"invalid process state for PID {pid}") - return (int(fields[1]), int(fields[2]), int(fields[3]), int(fields[19])) - - def _process_states(self) -> dict[int, tuple[int, int, int, int]]: - states: dict[int, tuple[int, int, int, int]] = {} - for entry in Path("/proc").iterdir(): - if not entry.name.isdigit(): - continue - pid = int(entry.name) - try: - states[pid] = self._process_state(pid) - except (OSError, ValueError): - continue - return states - - def _process_command(self, pid: int) -> tuple[str, ...]: - payload = (Path("/proc") / str(pid) / "cmdline").read_bytes() - return tuple( - os.fsdecode(argument) for argument in payload.split(b"\0") if argument - ) - - def _pidfd_open(self, pid: int) -> int: - pidfd_open = getattr(os, "pidfd_open", None) - if not callable(pidfd_open): - raise SnapshotRestoreError("snapshot requires Linux pidfd support") - return pidfd_open(pid) - - def _pidfd_exited(self, pidfd: int) -> bool: - poller = select.poll() - poller.register(pidfd, select.POLLIN | select.POLLHUP | select.POLLERR) - return any(self._pidfd_event_is_exit(event) for _fd, event in poller.poll(0)) - - def _pidfd_event_is_exit(self, event: int) -> bool: - if event & (select.POLLERR | select.POLLNVAL): - raise SnapshotRestoreError("restored process pidfd poll failed") - return bool(event & (select.POLLIN | select.POLLHUP)) - - def _pin_restored_tree( - self, - artifact: Path, - process_tree: tuple[int, ...], - cuda_holders: tuple[int, ...], - ) -> None: - root_pid = process_tree[0] - expected_pids = set(process_tree) - expected_cuda_holders = set(cuda_holders) - if not expected_cuda_holders or not expected_cuda_holders.issubset( - expected_pids - ): - raise SnapshotRestoreError( - "captured CUDA holders are outside the captured process tree" - ) - if root_pid in self._restored_processes: - raise SnapshotRestoreError("restored PID already has a pinned transaction") - handles: list[tuple[int, int]] = [] - try: - for pid in process_tree: - handles.append((pid, self._pidfd_open(pid))) - handles_by_pid = dict(handles) - if any(self._pidfd_exited(pidfd) for _pid, pidfd in handles): - raise SnapshotRestoreError( - "restored process exited during verification" - ) - - root_state = self._process_state(root_pid) - if root_state[1:3] != (root_pid, root_pid): - raise SnapshotRestoreError( - "restored root is not the session and process-group leader" - ) - command = self._process_command(root_pid) - release_file = str(artifact / "release.json") - if not any( - command[index : index + 2] == ("-m", "vllm.snapshot.server") - for index in range(len(command) - 1) - ) or not any( - command[index : index + 2] == ("--release-file", release_file) - for index in range(len(command) - 1) - ): - raise SnapshotRestoreError( - "restored root command does not match the snapshot barrier" - ) - if self._process_state(root_pid) != root_state: - raise SnapshotRestoreError("restored PID changed during verification") - if self._pidfd_exited(handles[0][1]): - raise SnapshotRestoreError("restored root exited during verification") - - first_states = self._process_states() - live_session = { - pid for pid, state in first_states.items() if state[2] == root_pid - } - if not expected_cuda_holders.issubset(live_session): - raise SnapshotRestoreError( - "restored session is missing captured CUDA holders" - ) - if live_session != expected_pids: - raise SnapshotRestoreError( - "restored session does not match the captured process tree" - ) - if first_states[root_pid] != root_state: - raise SnapshotRestoreError("restored PID changed during verification") - for pid in process_tree[1:]: - state = first_states[pid] - pidfd = handles_by_pid[pid] - if self._process_state(pid) != state or self._pidfd_exited(pidfd): - raise SnapshotRestoreError( - "restored process tree changed during verification" - ) - - second_states = self._process_states() - second_session = { - pid for pid, state in second_states.items() if state[2] == root_pid - } - if second_session != expected_pids or any( - second_states[pid] != first_states[pid] for pid in process_tree - ): - raise SnapshotRestoreError( - "restored process tree changed during verification" - ) - if any(self._pidfd_exited(pidfd) for _pid, pidfd in handles): - raise SnapshotRestoreError( - "restored process exited during verification" - ) - except BaseException as error: - _close_pidfds(handles) - if isinstance(error, SnapshotRestoreError): - raise - if not isinstance(error, Exception): - raise - raise SnapshotRestoreError( - "restored process tree could not be pinned" - ) from error - self._restored_processes[root_pid] = tuple(handles) - - def _abort_failed_restore( - self, - artifact: Path, - root_pid: int, - staged_remaps: dict[str, tuple[int, int]], - primary: BaseException, - ) -> NoReturn: - secondary: list[str] = [] - try: - _write_json_atomic( - artifact / "release.json", {"release": False}, overwrite=True - ) - except BaseException as error: - secondary.append(f"abort marker failed: {_error_detail(error)}") - # The marker above is only a request. Terminate and wait for the tree we - # own, before releasing the shared state it may still be using. A tree - # that could not be pinned is not ours to signal. - if root_pid in self._restored_processes: - try: - self.cleanup(root_pid) - except BaseException as error: - secondary.append(f"restore cleanup failed: {_error_detail(error)}") - try: - self._rollback_link_remaps(staged_remaps) - except BaseException as error: - secondary.append(f"remap rollback failed: {_error_detail(error)}") - if secondary: - raise SnapshotRestoreError( - f"snapshot restore failed: {_error_detail(primary)}; " - + "; ".join(secondary) - ) from primary - raise primary - - def _verify_external_cache_files(self, manifest: SnapshotManifest) -> None: - """Reject a rotated generated cache before CRIU reopens it by path. - - Like the PID probe above this reports a stale host early rather than - guarding it: CRIU stays authoritative for a cache replaced afterwards. - """ - for path, fingerprint in manifest.external_cache_files: - target = Path(path) - if not target.is_file(): - raise SnapshotRestoreError( - f"external cache file vanished since capture: {path}; " - f"{_CACHE_REMEDY}" - ) - if fingerprint in (_CACHE_LOCKED, _CACHE_UNREADABLE): - continue - try: - current = self._sha256(target) - except OSError as error: - raise SnapshotRestoreError( - f"external cache file is unreadable since capture: {path}; " - f"{_CACHE_REMEDY}" - ) from error - if current != fingerprint: - raise SnapshotRestoreError( - f"external cache file changed since capture: {path}; " - f"{_CACHE_REMEDY}" - ) - - def restore(self, artifact: Path, manifest: SnapshotManifest) -> int: - # This is necessarily a best-effort TOCTOU check. CRIU restore and the - # rollback below remain authoritative if a PID is claimed after it. - for pid in manifest.process_tree: - try: - os.kill(pid, 0) - except ProcessLookupError: - continue - except PermissionError: - pass - except Exception as error: - raise SnapshotRestoreError( - f"captured PID availability probe failed: {pid}" - ) from error - raise SnapshotRestoreError(f"captured PID is already occupied: {pid}") - self._verify_external_cache_files(manifest) - expected_root_pid = manifest.process_tree[0] - (artifact / "release.json").unlink(missing_ok=True) - pidfile = artifact / "restored.pid" - pidfile.unlink(missing_ok=True) - self._reset_child_log(artifact) - staged_remaps = self._restore_link_remaps(artifact) - try: - self._criu( - "restore", - artifact, - [ - "--images-dir", - str(artifact / "images"), - "--log-file", - "restore.log", - *_COMMON_CRIU_OPTIONS, - "--restore-detached", - "--pidfile", - str(pidfile), - ], - ) - # Pin before reading the PID file so that a failure between restore - # and the read cannot orphan the tree. - self._pin_restored_tree( - artifact, manifest.process_tree, manifest.cuda_holders - ) - restored_pid = self._read_restored_pid(pidfile) - if restored_pid != expected_root_pid: - raise SnapshotRestoreError( - "restored PID does not match the captured process tree" - ) - except BaseException as error: - self._abort_failed_restore( - artifact, expected_root_pid, staged_remaps, error - ) - return restored_pid - - def release(self, artifact: Path, host: str | None, port: int) -> None: - if not 1 <= port <= 65535: - raise SnapshotRestoreError(f"snapshot restore port is invalid: {port}") - family = socket.AF_INET6 if host and ":" in host else socket.AF_INET - with socket.socket(family=family, type=socket.SOCK_STREAM) as listener_probe: - listener_probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - try: - listener_probe.bind((host or "", port)) - except OSError as error: - address = f"{host or '0.0.0.0'}:{port}" - raise SnapshotRestoreError( - f"snapshot restore address is already in use: {address}" - ) from error - - _write_json_atomic( - artifact / "release.json", - {"release": True, "host": host, "port": port}, - overwrite=True, - ) - - def _connect_host(self, host: str | None) -> str: - if not host or host == "0.0.0.0": - return "127.0.0.1" - if host == "::": - return "::1" - return host - - def wait_listener(self, root_pid: int, host: str | None, port: int) -> None: - deadline = time.monotonic() + self.timeout_s - connect_host = self._connect_host(host) - while time.monotonic() < deadline: - try: - with socket.create_connection((connect_host, port), timeout=1): - return - except OSError as error: - handles = self._restored_processes.get(root_pid) - if handles is None or self._pidfd_exited(handles[0][1]): - raise SnapshotRestoreError( - "restored process exited before binding HTTP" - ) from error - time.sleep(0.1) - raise SnapshotRestoreError("restored HTTP listener did not become ready") - - def request_oracle( - self, host: str | None, port: int, manifest: SnapshotManifest - ) -> Oracle: - connect_host = self._connect_host(host) - url_host = f"[{connect_host}]" if ":" in connect_host else connect_host - payload = json.dumps( - { - "model": manifest.served_model_name, - "prompt": "The capital of France is", - "min_tokens": 1, - "max_tokens": 1, - "temperature": 0, - "seed": 0, - "logprobs": 0, - "return_token_ids": True, - } - ).encode() - request = urllib.request.Request( - f"http://{url_host}:{port}/v1/completions", - data=payload, - headers={"Content-Type": "application/json"}, - ) - with urllib.request.urlopen(request, timeout=120) as response: - body = json.load(response) - choice = body["choices"][0] - try: - return Oracle( - token_ids=tuple(choice["token_ids"]), - text=choice["text"], - sampled_token_logprob=choice["logprobs"]["token_logprobs"][0], - ) - except ValueError as error: - raise SnapshotRestoreError( - "snapshot HTTP canary must return exactly one finite token logprob" - ) from error - except (IndexError, KeyError, TypeError) as error: - raise SnapshotRestoreError( - "snapshot HTTP canary response is missing sampled token logprob" - ) from error - - def cleanup(self, root_pid: int) -> None: - handles = self._restored_processes.get(root_pid) - if handles is None: - raise SnapshotRestoreError("no pinned restore transaction for cleanup") - pidfd_send_signal = getattr(signal, "pidfd_send_signal", None) - if not callable(pidfd_send_signal): - raise SnapshotRestoreError("restored process cleanup is incomplete") - try: - for _pid, pidfd in handles: - with suppress(ProcessLookupError): - pidfd_send_signal(pidfd, signal.SIGTERM, None, 0) - time.sleep(min(0.5, self.timeout_s)) - for _pid, pidfd in handles: - with suppress(ProcessLookupError): - pidfd_send_signal(pidfd, signal.SIGKILL, None, 0) - except BaseException as error: - raise SnapshotRestoreError( - f"restored process cleanup is incomplete: {_error_detail(error)}" - ) from error - - poller = select.poll() - for _pid, pidfd in handles: - poller.register(pidfd, select.POLLIN | select.POLLHUP | select.POLLERR) - pending = {pidfd for _pid, pidfd in handles} - deadline = time.monotonic() + self.timeout_s - while pending: - remaining = deadline - time.monotonic() - if remaining <= 0: - raise SnapshotRestoreError("restored process cleanup is incomplete") - for pidfd, event in poller.poll(max(1, int(remaining * 1000))): - if not self._pidfd_event_is_exit(event): - continue - pending.discard(pidfd) - with suppress(KeyError, OSError): - poller.unregister(pidfd) - self._restored_processes.pop(root_pid) - _close_pidfds(handles) - - def complete_restore(self, root_pid: int) -> None: - handles = self._restored_processes.pop(root_pid, None) - if handles is None: - raise SnapshotRestoreError("no pinned restore transaction to complete") - _close_pidfds(handles) - - def abort_create(self, root_pid: int) -> None: - process = self._children.get(root_pid) - if process is None: - return - with suppress(ProcessLookupError): - os.killpg(root_pid, signal.SIGKILL) - try: - process.wait(timeout=10) - except subprocess.TimeoutExpired as error: - raise SnapshotCreateError( - f"snapshot child process group {root_pid} survived SIGKILL" - ) from error - self._children.pop(root_pid, None) diff --git a/vllm/snapshot/server.py b/vllm/snapshot/server.py deleted file mode 100644 index 2dc0f3384a..0000000000 --- a/vllm/snapshot/server.py +++ /dev/null @@ -1,252 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import argparse -import asyncio -import json -import os -import stat -import sys -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -from pydantic import ValidationError - -from vllm.snapshot.manifest import ReleaseMarker, _validation_path, _write_json_atomic -from vllm.snapshot.types import Oracle, oracles_match - -_CANARY_PROMPT = "The capital of France is" - - -class SnapshotCanaryError(RuntimeError): - """The initialized engine did not produce a valid snapshot oracle.""" - - -class SnapshotBarrierError(RuntimeError): - """The controller supplied an invalid snapshot release marker.""" - - -@dataclass(frozen=True) -class ListenerConfig: - host: str | None - port: int - - -@dataclass(frozen=True) -class ControlArgs: - ready_file: Path - release_file: Path - release_timeout_s: float - - -def write_ready_atomic(path: Path, oracle: Oracle) -> None: - _write_json_atomic( - path, - { - "sampled_token_logprob": oracle.sampled_token_logprob, - "token_ids": oracle.token_ids, - "text": oracle.text, - }, - ) - - -def read_release_marker(path: Path) -> ListenerConfig: - path = Path(path) - path_stat = path.lstat() - if stat.S_ISLNK(path_stat.st_mode): - raise SnapshotBarrierError("release marker must not be a symlink") - try: - payload = json.loads(path.read_text()) - except (OSError, json.JSONDecodeError) as error: - raise SnapshotBarrierError("release marker is not valid JSON") from error - try: - marker = ReleaseMarker.model_validate(payload) - except ValidationError as error: - raise SnapshotBarrierError( - f"release marker is invalid: {_validation_path(error)}" - ) from error - return ListenerConfig(host=marker.host, port=marker.port) - - -def parse_control_args(argv: list[str]) -> tuple[ControlArgs, list[str]]: - parser = argparse.ArgumentParser(add_help=False) - parser.add_argument("--ready-file", type=Path, required=True) - parser.add_argument("--release-file", type=Path, required=True) - parser.add_argument("--release-timeout-s", type=float, default=900.0) - control_args, remaining = parser.parse_known_args(argv) - if remaining and remaining[0] == "--": - remaining = remaining[1:] - if control_args.release_timeout_s <= 0: - raise ValueError("release timeout must be positive") - return ( - ControlArgs( - ready_file=control_args.ready_file, - release_file=control_args.release_file, - release_timeout_s=control_args.release_timeout_s, - ), - remaining, - ) - - -async def _release_reloadable_state(engine: Any) -> None: - """Discard model and KV state before the process image is captured.""" - await engine.sleep(level=2) - - -async def _restore_reloadable_state(engine: Any) -> None: - """Rebuild state discarded by ``_release_reloadable_state``.""" - await engine.wake_up(tags=["weights"]) - await engine.collective_rpc("reload_weights") - await engine.wake_up(tags=["kv_cache"]) - - -def oracle_from_request_output(request_output: Any) -> Oracle: - try: - candidate = request_output.outputs[0] - token_ids = tuple(candidate.token_ids) - (sampled_token_id,) = token_ids - return Oracle( - token_ids=token_ids, - text=candidate.text, - sampled_token_logprob=candidate.logprobs[0][sampled_token_id].logprob, - ) - except ValueError as error: - raise SnapshotCanaryError( - "snapshot canary must produce exactly one finite token logprob" - ) from error - except (AttributeError, IndexError, KeyError, TypeError) as error: - raise SnapshotCanaryError( - "snapshot canary did not return sampled token logprob" - ) from error - - -async def run_engine_canary(engine: Any) -> Oracle: - from vllm import SamplingParams - - final_output = None - sampling_params = SamplingParams( - temperature=0, - min_tokens=1, - max_tokens=1, - seed=0, - logprobs=0, - ) - async for output in engine.generate( - _CANARY_PROMPT, - sampling_params, - request_id="vllm-snapshot-canary", - ): - final_output = output - if final_output is None: - raise SnapshotCanaryError("canary generation returned no output") - return oracle_from_request_output(final_output) - - -def detach_snapshot_streams() -> None: - """Remove launch-log descriptors from the reusable process image.""" - sys.stdout.flush() - sys.stderr.flush() - sink = os.open(os.devnull, os.O_WRONLY) - try: - os.dup2(sink, sys.stdout.fileno()) - os.dup2(sink, sys.stderr.fileno()) - finally: - os.close(sink) - - -async def wait_for_release_marker( - path: Path, - *, - timeout_s: float, - poll_interval_s: float = 0.05, -) -> ListenerConfig: - if timeout_s <= 0: - raise ValueError("release timeout must be positive") - if poll_interval_s <= 0: - raise ValueError("release poll interval must be positive") - remaining_s = timeout_s - while not path.exists(): - if remaining_s <= 0: - raise TimeoutError(f"release marker not found before timeout: {path}") - delay_s = min(poll_interval_s, remaining_s) - await asyncio.sleep(delay_s) - remaining_s -= delay_s - return read_release_marker(path) - - -def parse_vllm_args(argv: list[str]) -> Any: - from vllm.entrypoints.launchers.cli_args import ( - make_arg_parser, - validate_parsed_serve_args, - ) - from vllm.utils.argparse_utils import FlexibleArgumentParser - - parser = FlexibleArgumentParser(prog="python -m vllm.snapshot.server") - args = make_arg_parser(parser).parse_args(argv) - if args.model_tag is not None: - args.model = args.model_tag - if args.grpc or args.headless: - raise ValueError("snapshot server requires the HTTP frontend") - if args.api_server_count not in (None, 1): - raise ValueError("snapshot server supports one API server") - args.api_server_count = None - validate_parsed_serve_args(args) - return args - - -async def run_vllm_snapshot_child(control: ControlArgs, args: Any) -> None: - from vllm.entrypoints.launchers.api_server.entry import ( - build_and_serve, - build_async_engine_client, - ) - from vllm.entrypoints.launchers.launcher import setup_server - - args.enable_sleep_mode = True - - async with build_async_engine_client(args) as engine: - oracle = await run_engine_canary(engine) - try: - await _release_reloadable_state(engine) - await _restore_reloadable_state(engine) - rehearsal_oracle = await run_engine_canary(engine) - if not oracles_match(oracle, rehearsal_oracle): - raise SnapshotCanaryError("snapshot rehearsal changed canary output") - await _release_reloadable_state(engine) - except SnapshotCanaryError: - raise - except Exception as error: - raise SnapshotCanaryError("snapshot rehearsal failed") from error - - detach_snapshot_streams() - write_ready_atomic(control.ready_file, oracle) - listener = await wait_for_release_marker( - control.release_file, - timeout_s=control.release_timeout_s, - ) - args.host = listener.host - args.port = listener.port - await _restore_reloadable_state(engine) - listen_address, sock = setup_server(args, reuse_port=False) - try: - shutdown_task = await build_and_serve( - engine, - listen_address, - sock, - args, - ) - await shutdown_task - finally: - sock.close() - - -def main(argv: list[str] | None = None) -> None: - control, remaining = parse_control_args(sys.argv[1:] if argv is None else argv) - args = parse_vllm_args(remaining) - import uvloop - - uvloop.run(run_vllm_snapshot_child(control, args)) - - -if __name__ == "__main__": - main() diff --git a/vllm/snapshot/types.py b/vllm/snapshot/types.py deleted file mode 100644 index 236ba12436..0000000000 --- a/vllm/snapshot/types.py +++ /dev/null @@ -1,37 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import math -from dataclasses import dataclass - - -@dataclass(frozen=True) -class Oracle: - token_ids: tuple[int, ...] - text: str - sampled_token_logprob: float - - def __post_init__(self) -> None: - if len(self.token_ids) != 1: - raise ValueError("snapshot oracle must contain exactly one token") - value = self.sampled_token_logprob - if ( - isinstance(value, bool) - or not isinstance(value, (int, float)) - or not math.isfinite(value) - ): - raise ValueError("sampled token logprob must be finite") - object.__setattr__(self, "sampled_token_logprob", float(value)) - - -def oracles_match(expected: Oracle, actual: Oracle) -> bool: - return ( - expected.token_ids == actual.token_ids - and expected.text == actual.text - and math.isclose( - expected.sampled_token_logprob, - actual.sampled_token_logprob, - rel_tol=0.0, - abs_tol=1e-3, - ) - ) diff --git a/vllm/transformers_utils/processors/deepseek_ocr.py b/vllm/transformers_utils/processors/deepseek_ocr.py deleted file mode 100644 index 55af43a709..0000000000 --- a/vllm/transformers_utils/processors/deepseek_ocr.py +++ /dev/null @@ -1,459 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# adapted from https://github.com/deepseek-ai/DeepSeek-OCR/blob/main/DeepSeek-OCR-master/DeepSeek-OCR-vllm/process/image_process.py -# and https://github.com/deepseek-ai/DeepSeek-OCR-2/blob/main/DeepSeek-OCR2-master/DeepSeek-OCR2-vllm/process/image_process.py -import math -from typing import Literal - -import torch -import torchvision.transforms as T -from PIL import Image, ImageOps -from transformers import BatchFeature, LlamaTokenizerFast -from transformers.processing_utils import ProcessorMixin - -# TODO(Isotr0py): change modes for variants -# see: https://github.com/deepseek-ai/DeepSeek-OCR/blob/8cf003d38821fa1b19c73da3bd1b0dc262ea8136/DeepSeek-OCR-master/DeepSeek-OCR-vllm/config.py#L1-L6 -# Tiny: base_size = 512, image_size = 512, crop_mode = False -# Small: base_size = 640, image_size = 640, crop_mode = False -# Base: base_size = 1024, image_size = 1024, crop_mode = False -# Large: base_size = 1280, image_size = 1280, crop_mode = False -# Gundam: base_size = 1024, image_size = 640, crop_mode = True -BASE_SIZE = 1024 -IMAGE_SIZE = 640 -CROP_MODE = True - -# TODO(Isotr0py): Expose as mm_kwargs -MIN_CROPS = 2 -MAX_CROPS = 6 # max:9; If your GPU memory is small, it is recommended to set it to 6. - - -def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size): - best_ratio_diff = float("inf") - best_ratio = (1, 1) - area = width * height - for ratio in target_ratios: - target_aspect_ratio = ratio[0] / ratio[1] - ratio_diff = abs(aspect_ratio - target_aspect_ratio) - if ratio_diff < best_ratio_diff: - best_ratio_diff = ratio_diff - best_ratio = ratio - elif ratio_diff == best_ratio_diff: - if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]: - best_ratio = ratio - return best_ratio - - -def calculate_aspect_ratios( - min_num: int = MIN_CROPS, max_num: int = MAX_CROPS -) -> list[tuple[int, int]]: - target_ratios: set[tuple[int, int]] = set( - (i, j) - for n in range(min_num, max_num + 1) - for i in range(1, n + 1) - for j in range(1, n + 1) - if i * j <= max_num and i * j >= min_num - ) - sorted_target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1]) - return sorted_target_ratios - - -def count_tiles( - orig_width, - orig_height, - min_num=MIN_CROPS, - max_num=MAX_CROPS, - image_size=640, - use_thumbnail=False, -): - aspect_ratio = orig_width / orig_height - - # calculate the existing image aspect ratio - target_ratios = calculate_aspect_ratios(min_num, max_num) - - # find the closest aspect ratio to the target - target_aspect_ratio = find_closest_aspect_ratio( - aspect_ratio, target_ratios, orig_width, orig_height, image_size - ) - - return target_aspect_ratio - - -def dynamic_preprocess( - image, min_num=MIN_CROPS, max_num=MAX_CROPS, image_size=640, use_thumbnail=False -): - orig_width, orig_height = image.size - aspect_ratio = orig_width / orig_height - - # calculate the existing image aspect ratio - target_ratios = calculate_aspect_ratios(min_num, max_num) - - # find the closest aspect ratio to the target - target_aspect_ratio = find_closest_aspect_ratio( - aspect_ratio, target_ratios, orig_width, orig_height, image_size - ) - - # calculate the target width and height - target_width = image_size * target_aspect_ratio[0] - target_height = image_size * target_aspect_ratio[1] - blocks = target_aspect_ratio[0] * target_aspect_ratio[1] - - # resize the image - resized_img = image.resize((target_width, target_height)) - processed_images = [] - for i in range(blocks): - box = ( - (i % (target_width // image_size)) * image_size, - (i // (target_width // image_size)) * image_size, - ((i % (target_width // image_size)) + 1) * image_size, - ((i // (target_width // image_size)) + 1) * image_size, - ) - # split the image - split_img = resized_img.crop(box) - processed_images.append(split_img) - assert len(processed_images) == blocks - if use_thumbnail and len(processed_images) != 1: - thumbnail_img = image.resize((image_size, image_size)) - processed_images.append(thumbnail_img) - return processed_images, target_aspect_ratio - - -class ImageTransform: - def __init__( - self, - mean: tuple[float, float, float] = (0.5, 0.5, 0.5), - std: tuple[float, float, float] = (0.5, 0.5, 0.5), - normalize: bool = True, - ): - self.mean = mean - self.std = std - self.normalize = normalize - - transform_pipelines = [T.ToTensor()] - - if normalize: - transform_pipelines.append(T.Normalize(mean, std)) - - self.transform = T.Compose(transform_pipelines) - - def __call__(self, pil_img: Image.Image): - x = self.transform(pil_img) - return x - - -class DeepseekOCRProcessor(ProcessorMixin): - tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast") - attributes = ["tokenizer"] - - def __init__( - self, - tokenizer: LlamaTokenizerFast, - patch_size: int = 16, - downsample_ratio: int = 4, - image_mean: tuple[float, float, float] = (0.5, 0.5, 0.5), - image_std: tuple[float, float, float] = (0.5, 0.5, 0.5), - normalize: bool = True, - image_token: str = "", - pad_token: str = "<｜▁pad▁｜>", - add_special_token: bool = False, - sft_format: str = "deepseek", - mask_prompt: bool = True, - ignore_id: int = -100, - image_size: int = IMAGE_SIZE, - base_size: int = BASE_SIZE, - strategy: Literal["v1", "v2"] = "v1", - max_crops: int = MAX_CROPS, - **kwargs, - ): - self.image_size = image_size - self.base_size = base_size - self.max_crops = max_crops - - # image token calculation strategy for - # Deepseek-OCR and Deepseek-OCR-2 - self.strategy = strategy - assert strategy in ["v1", "v2"], "Only 'v1' and 'v2' strategies are supported." - - self.patch_size = 16 - self.image_mean = image_mean - self.image_std = image_std - self.normalize = normalize - self.downsample_ratio = 4 - - self.image_transform = ImageTransform( - mean=image_mean, std=image_std, normalize=normalize - ) - - self.tokenizer = tokenizer - self.tokenizer.padding_side = "left" # must set this，padding side with make a difference in batch inference # noqa: E501 - - # add the pad_token as special token to use 'tokenizer.pad_token' - # and 'tokenizer.pad_token_id' - if self.tokenizer.pad_token is None: - self.tokenizer.add_special_tokens({"pad_token": pad_token}) - - # add image token - self.image_token_id = self.tokenizer.vocab.get(image_token) - self.image_token = image_token - self.pad_token = pad_token - self.add_special_token = add_special_token - self.sft_format = sft_format - self.mask_prompt = mask_prompt - self.ignore_id = ignore_id - - super().__init__( - tokenizer, - **kwargs, - ) - - @property - def bos_id(self): - return self.tokenizer.bos_token_id - - @property - def eos_id(self): - return self.tokenizer.eos_token_id - - @property - def pad_id(self): - return self.tokenizer.pad_token_id - - def encode(self, text: str, bos: bool = True, eos: bool = False): - t = self.tokenizer.encode(text, add_special_tokens=False) - if bos: - t = [self.bos_id] + t - if eos: - t = t + [self.eos_id] - return t - - def decode(self, t: list[int], **kwargs) -> str: - return self.tokenizer.decode(t, **kwargs) - - def process_one( - self, - prompt: str, - images: list[Image.Image], - crop_mode: bool = CROP_MODE, - ): - """Args: - prompt (str): the formatted prompt; - images (List[ImageType]): the list of images; - crop_mode (bool): if True, then crop the image; - - Returns: - outputs (BaseProcessorOutput): the output of the processor, - - input_ids (torch.LongTensor): [N + image tokens] - - target_ids (torch.LongTensor): [N + image tokens] - - pixel_values (torch.FloatTensor): [n_patches, 3, H, W] - - image_id (int): the id of the image token - - num_image_tokens (List[int]): the number of image tokens - - """ - if prompt is None or images is None: - raise ValueError("prompt and images must be used at the same time.") - - sft_format = prompt - - ( - input_ids, - pixel_values, - images_crop, - images_seq_mask, - images_spatial_crop, - num_image_tokens, - _, - ) = self.tokenize_with_images( - conversation=sft_format, - images=images, - bos=True, - eos=True, - cropping=crop_mode, - ) - - prepare = BatchFeature( - data=dict( - input_ids=input_ids, - pixel_values=pixel_values, - images_crop=images_crop, - images_seq_mask=images_seq_mask, - images_spatial_crop=images_spatial_crop, - num_image_tokens=num_image_tokens, - ), - tensor_type="pt", - ) - return prepare - - def __call__( - self, - *, - prompt: str, - images: list[Image.Image], - crop_mode: bool = CROP_MODE, - **kwargs, - ): - prepare = self.process_one( - prompt=prompt, - images=images, - crop_mode=crop_mode, - ) - - return prepare - - def tokenize_with_images( - self, - conversation: str, - images: list[Image.Image], - bos: bool = True, - eos: bool = True, - cropping: bool = True, - ): - """Tokenize text with tags.""" - num_image_tags = conversation.count(self.image_token) - if num_image_tags != len(images): - raise ValueError( - f"Number of {self.image_token!r} tokens in prompt " - f"({num_image_tags}) does not match number of images " - f"({len(images)})." - ) - text_splits = conversation.split(self.image_token) - images_list, images_crop_list, images_seq_mask, images_spatial_crop = ( - [], - [], - [], - [], - ) - image_shapes = [] - num_image_tokens = [] - tokenized_str = [] - for text_sep, image in zip(text_splits, images): - tokenized_sep = self.encode(text_sep, bos=False, eos=False) - tokenized_str += tokenized_sep - images_seq_mask += [False] * len(tokenized_sep) - - image_shapes.append(image.size) - - images_crop_raw = [] - if image.size[0] <= self.image_size and image.size[1] <= self.image_size: - crop_ratio = [1, 1] - elif cropping: - images_crop_raw, crop_ratio = dynamic_preprocess( - image, image_size=self.image_size, max_num=self.max_crops - ) - else: - crop_ratio = [1, 1] - - if not cropping: - image = image.resize((self.image_size, self.image_size)) - - global_view = ImageOps.pad( - image, - (self.base_size, self.base_size), - color=tuple(int(x * 255) for x in self.image_transform.mean), - ) - images_list.append(self.image_transform(global_view)) - - num_width_tiles, num_height_tiles = crop_ratio - images_spatial_crop.append([num_width_tiles, num_height_tiles]) - - if num_width_tiles > 1 or num_height_tiles > 1: - for cropped_image in images_crop_raw: - images_crop_list.append(self.image_transform(cropped_image)) - - num_queries = math.ceil( - (self.image_size // self.patch_size) / self.downsample_ratio - ) - num_queries_base = math.ceil( - (self.base_size // self.patch_size) / self.downsample_ratio - ) - - num_tokens_base = ( - (num_queries_base * (num_queries_base + 1)) - if self.strategy == "v1" - else num_queries_base * num_queries_base - ) - tokenized_image = [self.image_token_id] * num_tokens_base - - tokenized_image += [self.image_token_id] - if num_width_tiles > 1 or num_height_tiles > 1: - num_tokens_per_row = ( - num_queries * num_width_tiles + 1 - if self.strategy == "v1" - else num_queries * num_width_tiles - ) - local_row = [self.image_token_id] * num_tokens_per_row - tokenized_image += local_row * (num_queries * num_height_tiles) - tokenized_str += tokenized_image - images_seq_mask += [True] * len(tokenized_image) - num_image_tokens.append(len(tokenized_image)) - - """process the last text split""" - tokenized_sep = self.encode(text_splits[-1], bos=False, eos=False) - tokenized_str += tokenized_sep - images_seq_mask += [False] * len(tokenized_sep) - - """add the bos and eos tokens""" - if bos: - tokenized_str = [self.bos_id] + tokenized_str - images_seq_mask = [False] + images_seq_mask - if eos: - tokenized_str = tokenized_str + [self.eos_id] - images_seq_mask = images_seq_mask + [False] - - assert len(tokenized_str) == len(images_seq_mask), ( - f"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} " - f"is not equal to images_seq_mask's length {len(images_seq_mask)}." - ) - - masked_tokenized_str = [] - for token_index in tokenized_str: - if token_index != self.image_token_id: - masked_tokenized_str.append(token_index) - else: - masked_tokenized_str.append(self.ignore_id) - - assert ( - len(tokenized_str) == len(images_seq_mask) == len(masked_tokenized_str) - ), ( - f"tokenized_str's length {len(tokenized_str)}, " - f"input_ids' length {len(masked_tokenized_str)}, " - f"images_seq_mask's length {len(images_seq_mask)}, are not equal." - ) - - input_ids = torch.LongTensor(tokenized_str) - target_ids = torch.LongTensor(masked_tokenized_str) - images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool) - - # set input_ids < 0 | input_ids == self.image_token_id as ignore_id - target_ids[(input_ids < 0) | (input_ids == self.image_token_id)] = ( - self.ignore_id - ) - input_ids[input_ids < 0] = self.pad_id - - # Remove the ending eos token - assert input_ids[-1] == self.eos_id - input_ids = input_ids[:-1] - target_ids = target_ids[:-1] - images_seq_mask = images_seq_mask[:-1] - - if len(images_list) == 0: - pixel_values = torch.zeros((0, 3, self.base_size, self.base_size)) - images_spatial_crop = torch.zeros((0, 2), dtype=torch.long) - images_crop = torch.zeros((0, 3, self.image_size, self.image_size)) - else: - pixel_values = torch.stack(images_list, dim=0) - images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long) - if images_crop_list: - images_crop = torch.stack(images_crop_list, dim=0) - else: - images_crop = torch.zeros((0, 3, self.image_size, self.image_size)) - - input_ids = input_ids.unsqueeze(0) - - return ( - input_ids, - pixel_values, - images_crop, - images_seq_mask, - images_spatial_crop, - num_image_tokens, - image_shapes, - ) diff --git a/vllm/transformers_utils/processors/internvl.py b/vllm/transformers_utils/processors/internvl.py deleted file mode 100644 index 84976cb178..0000000000 --- a/vllm/transformers_utils/processors/internvl.py +++ /dev/null @@ -1,585 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# adapted from https://huggingface.co/OpenGVLab/InternVL2-4B/blob/main/modeling_internvl_chat.py -# -------------------------------------------------------- -# InternVL -# Copyright (c) 2023 OpenGVLab -# Licensed under The MIT License [see LICENSE for details] -# -------------------------------------------------------- - -import numpy.typing as npt -import torch -import torchvision.transforms as T -from PIL import Image -from transformers import ( - BaseVideoProcessor, - BatchFeature, - ImageProcessingMixin, - TensorType, -) -from transformers.processing_utils import ProcessorMixin - -from vllm.multimodal.image import convert_image_mode -from vllm.multimodal.processing import PromptUpdateDetails, cached_encode -from vllm.tokenizers.hf import HfTokenizer - -IMAGENET_MEAN = (0.485, 0.456, 0.406) -IMAGENET_STD = (0.229, 0.224, 0.225) - - -# adapted from https://huggingface.co/OpenGVLab/InternVL2-1B -def build_transform(input_size: int): - MEAN, STD = IMAGENET_MEAN, IMAGENET_STD - return T.Compose( - [ - T.Lambda(lambda img: convert_image_mode(img, "RGB")), - T.Resize( - (input_size, input_size), interpolation=T.InterpolationMode.BICUBIC - ), - T.ToTensor(), - T.Normalize(mean=MEAN, std=STD), - ] - ) - - -# adapted from https://huggingface.co/OpenGVLab/InternVL2-1B -def find_closest_aspect_ratio( - aspect_ratio: float, - target_ratios: list[tuple[int, int]], - *, - width: int, - height: int, - image_size: int, -) -> tuple[int, int]: - best_ratio_diff = float("inf") - best_ratio = (1, 1) - area = width * height - for ratio in target_ratios: - target_aspect_ratio = ratio[0] / ratio[1] - ratio_diff = abs(aspect_ratio - target_aspect_ratio) - if ratio_diff < best_ratio_diff: - best_ratio_diff = ratio_diff - best_ratio = ratio - elif ratio_diff == best_ratio_diff: - if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]: - best_ratio = ratio - return best_ratio - - -def resolve_internvl_min_max_num( - *, - min_dynamic_patch: int, - max_dynamic_patch: int, - dynamic_image_size: bool, - use_thumbnail: bool, -) -> tuple[int, int]: - min_dynamic_patch = min_dynamic_patch if dynamic_image_size else 1 - max_dynamic_patch = max_dynamic_patch if dynamic_image_size else 1 - - if use_thumbnail and max_dynamic_patch != 1: - max_dynamic_patch += 1 - - return min_dynamic_patch, max_dynamic_patch - - -def get_internvl_target_ratios( - min_num: int, - max_num: int, -) -> list[tuple[int, int]]: - target_ratios = { - (i, j) - for n in range(min_num, max_num + 1) - for i in range(1, n + 1) - for j in range(1, n + 1) - if min_num <= i * j <= max_num - } - return sorted(target_ratios, key=lambda x: x[0] * x[1]) - - -def calculate_internvl_targets( - *, - orig_width: int, - orig_height: int, - target_ratios: list[tuple[int, int]], - image_size: int, - use_thumbnail: bool, -) -> tuple[int, int, int]: - aspect_ratio = orig_width / orig_height - - # find the closest aspect ratio to the target - target_aspect_ratio = find_closest_aspect_ratio( - aspect_ratio, - target_ratios, - width=orig_width, - height=orig_height, - image_size=image_size, - ) - - # calculate the target width and height - target_width = image_size * target_aspect_ratio[0] - target_height = image_size * target_aspect_ratio[1] - blocks = target_aspect_ratio[0] * target_aspect_ratio[1] - - # add thumbnail image if num_blocks != 1 - if use_thumbnail and blocks != 1: - blocks += 1 - - return blocks, target_width, target_height - - -# adapted from https://huggingface.co/OpenGVLab/InternVL2-1B -def dynamic_preprocess_internvl( - image: Image.Image, - *, - target_ratios: list[tuple[int, int]], - image_size: int, - use_thumbnail: bool, -) -> list[Image.Image]: - orig_width, orig_height = image.size - - # calculate the number of blocks without thumbnail - blocks, target_width, target_height = calculate_internvl_targets( - orig_width=orig_width, - orig_height=orig_height, - target_ratios=target_ratios, - image_size=image_size, - use_thumbnail=False, - ) - - # resize the image - resized_img = image.resize((target_width, target_height)) - processed_images = [] - for i in range(blocks): - box = ( - (i % (target_width // image_size)) * image_size, - (i // (target_width // image_size)) * image_size, - ((i % (target_width // image_size)) + 1) * image_size, - ((i // (target_width // image_size)) + 1) * image_size, - ) - # split the image - split_img = resized_img.crop(box) - processed_images.append(split_img) - - assert len(processed_images) == blocks - - if use_thumbnail and len(processed_images) != 1: - thumbnail_img = image.resize((image_size, image_size)) - processed_images.append(thumbnail_img) - - return processed_images - - -# adapted from https://huggingface.co/OpenGVLab/InternVL2-1B -def image_to_pixel_values_internvl( - image: Image.Image, - *, - input_size: int, - min_num: int, - max_num: int, - use_thumbnail: bool, -) -> torch.Tensor: - target_ratios = get_internvl_target_ratios(min_num, max_num) - - transform = build_transform(input_size=input_size) - images = dynamic_preprocess_internvl( - image, - target_ratios=target_ratios, - image_size=input_size, - use_thumbnail=use_thumbnail, - ) - - pixel_values = torch.stack([transform(image) for image in images]) - return pixel_values - - -# adapted from https://huggingface.co/OpenGVLab/InternVL2-1B -def video_to_pixel_values_internvl( - video: npt.NDArray, - *, - input_size: int, - min_num: int, - max_num: int, - use_thumbnail: bool, -) -> torch.Tensor: - target_ratios = get_internvl_target_ratios(min_num, max_num) - - transform = build_transform(input_size=input_size) - frames_list = list[Image.Image]() - for frame in video: - pil_frame = dynamic_preprocess_internvl( - Image.fromarray(frame, mode="RGB"), - target_ratios=target_ratios, - image_size=input_size, - use_thumbnail=use_thumbnail, - ) - assert len(pil_frame) == 1 - frames_list.extend(pil_frame) - - pixel_values = torch.stack([transform(image) for image in frames_list]) - return pixel_values - - -class InternVLImageProcessor(ImageProcessingMixin): - def __init__( - self, - image_size: int, - min_dynamic_patch: int, - max_dynamic_patch: int, - dynamic_image_size: bool, - use_thumbnail: bool, - ) -> None: - self.image_size = image_size - self.min_dynamic_patch = min_dynamic_patch - self.max_dynamic_patch = max_dynamic_patch - self.dynamic_image_size = dynamic_image_size - self.use_thumbnail = use_thumbnail - - def resolve_min_max_num( - self, - *, - min_dynamic_patch: int | None = None, - max_dynamic_patch: int | None = None, - dynamic_image_size: bool | None = None, - use_thumbnail: bool | None = None, - ) -> tuple[int, int]: - if min_dynamic_patch is None: - min_dynamic_patch = self.min_dynamic_patch - if max_dynamic_patch is None: - max_dynamic_patch = self.max_dynamic_patch - if dynamic_image_size is None: - dynamic_image_size = self.dynamic_image_size - if use_thumbnail is None: - use_thumbnail = self.use_thumbnail - - return resolve_internvl_min_max_num( - min_dynamic_patch=min_dynamic_patch, - max_dynamic_patch=max_dynamic_patch, - dynamic_image_size=dynamic_image_size, - use_thumbnail=use_thumbnail, - ) - - def _images_to_pixel_values_lst( - self, - images: list[Image.Image], - min_dynamic_patch: int | None = None, - max_dynamic_patch: int | None = None, - dynamic_image_size: bool | None = None, - ) -> list[torch.Tensor]: - if min_dynamic_patch is None: - min_dynamic_patch = self.min_dynamic_patch - if max_dynamic_patch is None: - max_dynamic_patch = self.max_dynamic_patch - if dynamic_image_size is None: - dynamic_image_size = self.dynamic_image_size - - min_num, max_num = resolve_internvl_min_max_num( - min_dynamic_patch=min_dynamic_patch, - max_dynamic_patch=max_dynamic_patch, - dynamic_image_size=dynamic_image_size, - use_thumbnail=False, # Applied in image_to_pixel_values - ) - - return [ - image_to_pixel_values_internvl( - image, - input_size=self.image_size, - min_num=min_num, - max_num=max_num, - use_thumbnail=self.use_thumbnail, - ) - for image in images - ] - - def __call__( - self, - images: Image.Image | list[Image.Image], - *, - min_dynamic_patch: int | None = None, - max_dynamic_patch: int | None = None, - dynamic_image_size: bool | None = None, - return_tensors: str | TensorType | None = None, - **kwargs, - ) -> BatchFeature: - images_lst = [images] if not isinstance(images, list) else images - - pixel_values_lst = self._images_to_pixel_values_lst( - images_lst, - min_dynamic_patch=min_dynamic_patch, - max_dynamic_patch=max_dynamic_patch, - dynamic_image_size=dynamic_image_size, - ) - - image_inputs = { - "pixel_values_flat": torch.cat(pixel_values_lst), - "image_num_patches": torch.tensor([len(item) for item in pixel_values_lst]), - } - return BatchFeature(image_inputs, tensor_type=return_tensors) - - -class InternVLVideoProcessor(BaseVideoProcessor): - def __init__( - self, - image_size: int, - ) -> None: - self.image_size = image_size - - def _videos_to_pixel_values_lst( - self, - videos: list[npt.NDArray], - ) -> list[torch.Tensor]: - return [ - video_to_pixel_values_internvl( - video, - input_size=self.image_size, - min_num=1, - max_num=1, - use_thumbnail=False, - ) - for video in videos - ] - - def __call__( - self, - videos: npt.NDArray | list[npt.NDArray], - *, - return_tensors: str | TensorType | None = None, - **kwargs, - ) -> BatchFeature: - videos_lst = [videos] if not isinstance(videos, list) else videos - - pixel_values_lst = self._videos_to_pixel_values_lst(videos_lst) - - image_inputs = { - "pixel_values_flat_video": torch.cat(pixel_values_lst), - "video_num_patches": torch.tensor([len(item) for item in pixel_values_lst]), - } - return BatchFeature(image_inputs, tensor_type=return_tensors) - - -class InternVLProcessor(ProcessorMixin): - """This model doesn't define its own HF processor, - so we implement our own one here. - - The code to insert image tokens is based on: - https://huggingface.co/OpenGVLab/InternVL2-1B/blob/main/modeling_internvl_chat.py#L252 - - Code for video processing is adapted from video example: - https://huggingface.co/OpenGVLab/InternVL3-1B#inference-with-transformers - """ - - attributes = ["image_processor", "tokenizer", "video_processor"] - - def __init__( - self, - image_processor: InternVLImageProcessor, - tokenizer: HfTokenizer, - video_processor: InternVLVideoProcessor | None = None, - *, - image_seq_length: int, - start_image_token: str = "", - end_image_token: str = "", - ctx_image_token: str = "", - ctx_video_token: str | None = None, - ) -> None: - self.image_processor = image_processor - self.tokenizer = tokenizer - self.video_processor = video_processor - - self.image_seq_length = image_seq_length - self.start_image_token = start_image_token - self.end_image_token = end_image_token - self.ctx_image_token = ctx_image_token - self.ctx_video_token = ctx_video_token - - self.start_image_token_id = tokenizer.convert_tokens_to_ids(start_image_token) - self.end_image_token_id = tokenizer.convert_tokens_to_ids(end_image_token) - self.ctx_image_token_id = tokenizer.convert_tokens_to_ids(ctx_image_token) - self.ctx_video_token_id = ( - None - if ctx_video_token is None - else tokenizer.convert_tokens_to_ids(ctx_video_token) - ) - - def resolve_target_ratios( - self, - *, - min_dynamic_patch: int | None = None, - max_dynamic_patch: int | None = None, - dynamic_image_size: bool | None = None, - use_thumbnail: bool | None = None, - ) -> list[tuple[int, int]]: - min_num, max_num = self.image_processor.resolve_min_max_num( - min_dynamic_patch=min_dynamic_patch, - max_dynamic_patch=max_dynamic_patch, - dynamic_image_size=dynamic_image_size, - use_thumbnail=use_thumbnail, - ) - - return get_internvl_target_ratios(min_num, max_num) - - def get_num_image_tokens( - self, - *, - image_width: int, - image_height: int, - ) -> int: - image_processor = self.image_processor - target_ratios = self.resolve_target_ratios( - use_thumbnail=False, # Applied in calculate_targets - ) - - num_patches, _, _ = calculate_internvl_targets( - orig_width=image_width, - orig_height=image_height, - image_size=image_processor.image_size, - target_ratios=target_ratios, - use_thumbnail=image_processor.use_thumbnail, - ) - - return num_patches * self.image_seq_length - - def get_image_repl( - self, - num_patches: int | None, - num_features: int | None = None, - ) -> PromptUpdateDetails: - if num_patches is None: - assert num_features is not None - else: - num_features = num_patches * self.image_seq_length - - repl_features = self.ctx_image_token * num_features - repl_full = self.start_image_token + repl_features + self.end_image_token - - full_ids = cached_encode(self.tokenizer, repl_full, add_special_tokens=False) - - return PromptUpdateDetails.select_token_id(full_ids, self.ctx_image_token_id) - - def get_video_repl(self, num_patches: int) -> PromptUpdateDetails: - assert self.ctx_video_token is not None - - repl_features = self.ctx_video_token * self.image_seq_length - repl_features_with_sep = ( - self.start_image_token + repl_features + self.end_image_token - ) - # num_patches is equal to num_frames - repl_full = "".join( - [f"Frame{i + 1}: {repl_features_with_sep}" for i in range(num_patches)] - ) - - full_ids = cached_encode(self.tokenizer, repl_full, add_special_tokens=False) - assert self.ctx_video_token_id is not None - - return PromptUpdateDetails.select_token_id(full_ids, self.ctx_video_token_id) - - def __call__( - self, - text: str | list[str] | None = None, - images: Image.Image | list[Image.Image] | None = None, - videos: npt.NDArray | list[npt.NDArray] | None = None, - *, - min_dynamic_patch: int | None = None, - max_dynamic_patch: int | None = None, - dynamic_image_size: bool | None = None, - return_tensors: str | TensorType | None = None, - **kwargs, - ) -> BatchFeature: - if images is not None: - image_inputs = self.image_processor( - images=images, - min_dynamic_patch=min_dynamic_patch, - max_dynamic_patch=max_dynamic_patch, - dynamic_image_size=dynamic_image_size, - return_tensors=return_tensors, - ) - image_num_patches = image_inputs["image_num_patches"] - else: - image_inputs = {} - image_num_patches = [] - - if videos is not None: - if self.video_processor is None: - raise ValueError("This model does not support video inputs") - - video_inputs = self.video_processor( - videos=videos, - return_tensors=return_tensors, - ) - video_num_patches = video_inputs["video_num_patches"] - else: - video_inputs = {} - video_num_patches = [] - - if text is not None: - if not isinstance(text, list): - text = [text] - - if image_inputs: - image_token = "" - image_index = 0 - processed_text = list[str]() - replace_strings = list[str]() - - for prompt in text: - new_prompt = prompt - - while image_token in new_prompt: - new_prompt = new_prompt.replace(image_token, "", 1) - image_repl = self.get_image_repl(image_num_patches[image_index]) - # Convert the token IDs back to text for the - # string-based replacement below - replace_strings.append( - self.tokenizer.decode( - image_repl.full, skip_special_tokens=False - ) - ) - image_index += 1 - - while "" in new_prompt: - replace_str = replace_strings.pop(0) - new_prompt = new_prompt.replace("", replace_str, 1) - - processed_text.append(new_prompt) - - text = processed_text - - if video_inputs: - video_token = "