chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,90 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from tests.v1.engine.utils import (
|
||||
FULL_STRINGS,
|
||||
NUM_PROMPT_LOGPROBS_UNDER_TEST,
|
||||
NUM_SAMPLE_LOGPROBS_UNDER_TEST,
|
||||
PROMPT_LEN,
|
||||
TOKENIZER_NAME,
|
||||
DummyOutputProcessorTestVectors,
|
||||
generate_dummy_prompt_logprobs_tensors,
|
||||
generate_dummy_sample_logprobs,
|
||||
)
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
|
||||
from ...distributed.conftest import publisher_config, random_port # noqa: F401
|
||||
|
||||
EngineCoreSampleLogprobsType = list[tuple[torch.Tensor, torch.Tensor]]
|
||||
EngineCorePromptLogprobsType = tuple[torch.Tensor, torch.Tensor]
|
||||
|
||||
|
||||
def _build_test_vectors_no_logprobs() -> DummyOutputProcessorTestVectors:
|
||||
"""Generate output processor dummy test vectors, without logprobs
|
||||
|
||||
Returns:
|
||||
DummyOutputProcessorTestVectors instance with no logprobs
|
||||
"""
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
|
||||
vllm_config = EngineArgs(model=TOKENIZER_NAME).create_engine_config()
|
||||
# Tokenize prompts under test & create dummy generated tokens
|
||||
prompt_tokens = [tokenizer(text).input_ids[:PROMPT_LEN] for text in FULL_STRINGS]
|
||||
generation_tokens = [
|
||||
tokenizer(text).input_ids[PROMPT_LEN:] for text in FULL_STRINGS
|
||||
]
|
||||
# Generate prompt strings
|
||||
prompt_strings = [
|
||||
tokenizer.decode(prompt_tokens, skip_special_tokens=True)
|
||||
for prompt_tokens in prompt_tokens
|
||||
]
|
||||
prompt_strings_len = [len(prompt_string) for prompt_string in prompt_strings]
|
||||
return DummyOutputProcessorTestVectors(
|
||||
tokenizer=tokenizer,
|
||||
vllm_config=vllm_config,
|
||||
full_tokens=[tokenizer(text).input_ids for text in FULL_STRINGS],
|
||||
prompt_tokens=prompt_tokens,
|
||||
generation_tokens=generation_tokens,
|
||||
prompt_strings=prompt_strings,
|
||||
prompt_strings_len=prompt_strings_len,
|
||||
generation_strings=[
|
||||
text[prompt_len:]
|
||||
for text, prompt_len in zip(FULL_STRINGS, prompt_strings_len)
|
||||
],
|
||||
prompt_logprobs=[],
|
||||
generation_logprobs=[],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dummy_test_vectors() -> DummyOutputProcessorTestVectors:
|
||||
"""Generate output processor dummy test vectors, with logprobs
|
||||
|
||||
Returns:
|
||||
DummyOutputProcessorTestVectors instance with logprobs
|
||||
"""
|
||||
# Build dummy test vectors without logprobs
|
||||
dtv = _build_test_vectors_no_logprobs()
|
||||
# Inject logprobs into dummy test vectors
|
||||
# data structure
|
||||
dtv.generation_logprobs = [
|
||||
generate_dummy_sample_logprobs(
|
||||
sampled_tokens_list=tokens_list,
|
||||
num_logprobs=NUM_SAMPLE_LOGPROBS_UNDER_TEST,
|
||||
tokenizer=dtv.tokenizer,
|
||||
)
|
||||
for tokens_list in dtv.generation_tokens
|
||||
]
|
||||
dtv.prompt_logprobs = [
|
||||
generate_dummy_prompt_logprobs_tensors(
|
||||
prompt_tokens_list=tokens_list,
|
||||
num_logprobs=NUM_PROMPT_LOGPROBS_UNDER_TEST,
|
||||
tokenizer=dtv.tokenizer,
|
||||
)
|
||||
for tokens_list in dtv.prompt_tokens
|
||||
]
|
||||
return dtv
|
||||
@@ -0,0 +1,336 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Test for the fix in PR #29987: Eagerly abort cancelled final-step requests.
|
||||
|
||||
This test verifies that when a request is aborted during its final execution
|
||||
step (when it would naturally complete), it is properly marked as aborted
|
||||
rather than being treated as normally completed.
|
||||
|
||||
The test uses a dummy KV connector to verify that the connector receives
|
||||
the correct finish status (FINISHED_ABORTED, not FINISHED_LENGTH_CAPPED).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import SamplingParams
|
||||
from vllm.config import KVTransferConfig, VllmConfig
|
||||
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
|
||||
KVConnectorBase_V1,
|
||||
KVConnectorMetadata,
|
||||
KVConnectorRole,
|
||||
)
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
from vllm.utils.torch_utils import set_default_torch_num_threads
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
from vllm.v1.kv_cache_interface import KVCacheConfig
|
||||
from vllm.v1.request import Request
|
||||
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip(reason="V1 currently only supported on CUDA.", allow_module_level=True)
|
||||
|
||||
TEXT_PROMPT = "Hello"
|
||||
|
||||
|
||||
class DummyKVConnectorMetadata(KVConnectorMetadata):
|
||||
"""Dummy metadata for the test connector."""
|
||||
|
||||
def __init__(self):
|
||||
self.requests: list = []
|
||||
|
||||
|
||||
class DummyKVConnector(KVConnectorBase_V1):
|
||||
"""
|
||||
Dummy KV connector that captures request finish statuses to a file.
|
||||
This is used to verify the fix - without the fix, a request aborted
|
||||
during its final step would be captured as FINISHED_LENGTH_CAPPED
|
||||
instead of FINISHED_ABORTED.
|
||||
|
||||
The connector runs in a separate process, so we write statuses to a file
|
||||
that can be read by the test process.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
role: KVConnectorRole,
|
||||
kv_cache_config: KVCacheConfig,
|
||||
):
|
||||
super().__init__(vllm_config, role, kv_cache_config)
|
||||
# Get the status file path from extra config
|
||||
extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config or {}
|
||||
self.status_file = extra_config.get("status_file")
|
||||
# Log that we were initialized
|
||||
if self.status_file:
|
||||
try:
|
||||
with open(self.status_file, "a") as f:
|
||||
f.write(f"INIT:{role.name}\n")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def get_num_new_matched_tokens(
|
||||
self,
|
||||
request: Request,
|
||||
num_computed_tokens: int,
|
||||
) -> tuple[int | None, bool]:
|
||||
return (0, False)
|
||||
|
||||
def update_state_after_alloc(
|
||||
self,
|
||||
request: Request,
|
||||
blocks: Any,
|
||||
num_external_tokens: int,
|
||||
):
|
||||
pass
|
||||
|
||||
def build_connector_meta(
|
||||
self, scheduler_output: SchedulerOutput
|
||||
) -> KVConnectorMetadata:
|
||||
return DummyKVConnectorMetadata()
|
||||
|
||||
def request_finished(
|
||||
self,
|
||||
request: Request,
|
||||
block_ids: list[int],
|
||||
) -> tuple[bool, dict[str, Any] | None]:
|
||||
"""Capture the request status when finished by writing to a file."""
|
||||
if self.status_file:
|
||||
try:
|
||||
with open(self.status_file, "a") as f:
|
||||
# Write the status name (e.g., "FINISHED_ABORTED")
|
||||
f.write(f"{request.status.name}\n")
|
||||
except Exception as e:
|
||||
# Log but don't fail - this is just test instrumentation
|
||||
print(f"[DummyKVConnector] Failed to write status: {e}")
|
||||
return False, None
|
||||
|
||||
def start_load_kv(self, forward_context: Any, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
def wait_for_layer_load(self, layer_name: str) -> None:
|
||||
pass
|
||||
|
||||
def save_kv_layer(
|
||||
self,
|
||||
layer_name: str,
|
||||
kv_layer: Any,
|
||||
attn_metadata: Any,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
def wait_for_save(self):
|
||||
pass
|
||||
|
||||
|
||||
# Register the dummy connector
|
||||
KVConnectorFactory.register_connector(
|
||||
"DummyKVConnector", __name__, DummyKVConnector.__name__
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("async_scheduling", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_abort_during_final_step(async_scheduling: bool):
|
||||
"""
|
||||
Test that a request aborted during its final execution step is treated as
|
||||
aborted rather than completed.
|
||||
|
||||
This test:
|
||||
1. Monkeypatches execute_model to wait for a file to be deleted
|
||||
2. Configures a dummy KV connector to capture finish statuses
|
||||
3. Starts a request with max_tokens=1 (will complete on first decode step)
|
||||
4. Aborts the request, then deletes the file to unblock execute_model
|
||||
5. Verifies the KV connector received FINISHED_ABORTED not FINISHED_LENGTH_CAPPED
|
||||
|
||||
See https://github.com/vllm-project/vllm/pull/29987.
|
||||
|
||||
Without the fix, the KV connector would see FINISHED_LENGTH_CAPPED because
|
||||
update_from_output() would mark the request as completed before processing
|
||||
the abort. This causes KV cache blocks to not be freed properly in
|
||||
disaggregated prefill scenarios.
|
||||
|
||||
With the fix, _process_aborts_queue() runs before update_from_output(), so the
|
||||
abort takes precedence and the KV connector sees FINISHED_ABORTED.
|
||||
"""
|
||||
|
||||
# Create three temporary files:
|
||||
# 1. ready_file: deleted by execute_model to signal it has started
|
||||
# 2. block_file: execute_model waits for this to be deleted
|
||||
# 3. status_file: KV connector writes finish statuses here
|
||||
with tempfile.NamedTemporaryFile(delete=False) as f:
|
||||
ready_file = Path(f.name)
|
||||
with tempfile.NamedTemporaryFile(delete=False) as f2:
|
||||
block_file = Path(f2.name)
|
||||
with tempfile.NamedTemporaryFile(delete=False, mode="w") as f3:
|
||||
status_file = Path(f3.name)
|
||||
|
||||
try:
|
||||
# Get the original execute_model method
|
||||
from vllm.v1.worker.gpu_worker import Worker
|
||||
|
||||
original_execute_model = Worker.execute_model
|
||||
|
||||
def execute_model_with_wait(self, scheduler_output):
|
||||
# V2's `gpu_worker.compile_or_warm_up_model` calls
|
||||
# `warmup_kernels(...)` during engine init, which itself calls
|
||||
# `Worker.execute_model` three times (prefill / decode / cleanup)
|
||||
# to JIT compile triton kernels. None of those carry the test's
|
||||
# request id, so we only stall when our actual request is being
|
||||
# processed.
|
||||
scheduled = scheduler_output.num_scheduled_tokens or {}
|
||||
finished = scheduler_output.finished_req_ids or set()
|
||||
|
||||
def is_target_request(req_ids):
|
||||
return any(
|
||||
rid == request_id or rid.startswith(f"{request_id}-")
|
||||
for rid in req_ids
|
||||
)
|
||||
|
||||
if is_target_request(scheduled) or is_target_request(finished):
|
||||
# Signal that execute_model has been called by deleting ready_file
|
||||
if ready_file.exists():
|
||||
ready_file.unlink()
|
||||
|
||||
# Wait for the block file to be deleted (triggered from test after
|
||||
# abort). This runs in the worker process (after fork), so we poll
|
||||
# the filesystem.
|
||||
while block_file.exists():
|
||||
time.sleep(0.01)
|
||||
return original_execute_model(self, scheduler_output)
|
||||
|
||||
# Patch execute_model to inject the wait
|
||||
# This happens before the worker process is forked, so the patch applies there
|
||||
with patch.object(Worker, "execute_model", execute_model_with_wait):
|
||||
request_id = "test-abort-final-step"
|
||||
|
||||
# Configure engine with dummy KV connector
|
||||
# Pass the status file path so the connector can write to it
|
||||
kv_transfer_config = KVTransferConfig(
|
||||
kv_connector="DummyKVConnector",
|
||||
kv_role="kv_both",
|
||||
kv_connector_extra_config={"status_file": str(status_file)},
|
||||
)
|
||||
engine_args = AsyncEngineArgs(
|
||||
model="meta-llama/Llama-3.2-1B-Instruct",
|
||||
enforce_eager=True,
|
||||
async_scheduling=async_scheduling,
|
||||
kv_transfer_config=kv_transfer_config,
|
||||
)
|
||||
|
||||
with set_default_torch_num_threads(1):
|
||||
engine = AsyncLLM.from_engine_args(engine_args)
|
||||
|
||||
try:
|
||||
# Create a request that will complete after just 1 token
|
||||
sampling_params = SamplingParams(
|
||||
max_tokens=1,
|
||||
ignore_eos=True,
|
||||
output_kind=RequestOutputKind.DELTA,
|
||||
)
|
||||
|
||||
# Start generation in a task
|
||||
outputs = []
|
||||
|
||||
async def generate():
|
||||
async for output in engine.generate(
|
||||
request_id=request_id,
|
||||
prompt=TEXT_PROMPT,
|
||||
sampling_params=sampling_params,
|
||||
):
|
||||
outputs.append(output)
|
||||
|
||||
gen_task = asyncio.create_task(generate())
|
||||
|
||||
# Wait for execute_model to signal it has started (with timeout)
|
||||
timeout = 5.0 # 5 second timeout
|
||||
start_time = time.time()
|
||||
while ready_file.exists():
|
||||
if time.time() - start_time > timeout:
|
||||
raise TimeoutError(
|
||||
"Timeout waiting for execute_model to start. "
|
||||
"The monkeypatch may not be working correctly, "
|
||||
"for example if spawn was used instead of fork."
|
||||
)
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
# Abort the request while execute_model is blocked
|
||||
await engine.abort(request_id)
|
||||
|
||||
# Now unblock execute_model by deleting the file
|
||||
# The abort should be processed before the model output
|
||||
block_file.unlink()
|
||||
|
||||
# Wait for generation to complete
|
||||
await gen_task
|
||||
|
||||
# Poll for the KV connector to record the finish status
|
||||
timeout = 5.0
|
||||
start = time.time()
|
||||
captured_statuses = []
|
||||
while time.time() - start < timeout:
|
||||
with open(status_file) as f4:
|
||||
status_lines = f4.read().strip().split("\n")
|
||||
captured_statuses = [
|
||||
line
|
||||
for line in status_lines
|
||||
if line and line.startswith("FINISHED_")
|
||||
]
|
||||
if captured_statuses:
|
||||
break
|
||||
await asyncio.sleep(0.05)
|
||||
else:
|
||||
raise TimeoutError(
|
||||
"Timeout waiting for KV connector to record finish status."
|
||||
)
|
||||
|
||||
# Verify we got output
|
||||
assert len(outputs) > 0, "Should have received at least one output"
|
||||
|
||||
# The final output should have finish_reason="abort"
|
||||
final_output = outputs[-1]
|
||||
assert final_output.finished, (
|
||||
"Final output should be marked as finished"
|
||||
)
|
||||
assert final_output.outputs[0].finish_reason == "abort", (
|
||||
f"Expected finish_reason='abort' but got "
|
||||
f"'{final_output.outputs[0].finish_reason}'. "
|
||||
)
|
||||
|
||||
assert len(captured_statuses) >= 1, (
|
||||
f"Expected at least 1 captured finish status, got "
|
||||
f"{len(captured_statuses)}. File content: {status_lines}"
|
||||
)
|
||||
|
||||
assert "FINISHED_ABORTED" in captured_statuses, (
|
||||
f"KV connector should see FINISHED_ABORTED but got "
|
||||
f"{captured_statuses}. "
|
||||
)
|
||||
|
||||
# Verify cleanup
|
||||
assert not engine.output_processor.has_unfinished_requests()
|
||||
|
||||
finally:
|
||||
# Shutdown the engine
|
||||
engine.shutdown()
|
||||
|
||||
finally:
|
||||
# Clean up temporary files if they still exist
|
||||
if ready_file.exists():
|
||||
ready_file.unlink()
|
||||
if block_file.exists():
|
||||
block_file.unlink()
|
||||
if status_file.exists():
|
||||
status_file.unlink()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,370 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
import socket
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import ray
|
||||
import zmq
|
||||
|
||||
from vllm.utils.network_utils import make_zmq_socket, split_zmq_path
|
||||
from vllm.v1.engine.core import EngineCoreActorMixin
|
||||
from vllm.v1.engine.core_client import BackgroundResources
|
||||
from vllm.v1.engine.utils import (
|
||||
CoreEngineActorManager,
|
||||
EngineZmqAddresses,
|
||||
get_engine_zmq_addresses,
|
||||
launch_core_engines,
|
||||
)
|
||||
from vllm.v1.utils import APIServerProcessManager
|
||||
|
||||
|
||||
class _StubEngineCoreActor(EngineCoreActorMixin):
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: Any,
|
||||
local_client: bool,
|
||||
addresses: EngineZmqAddresses,
|
||||
executor_class: type[Any],
|
||||
log_stats: bool,
|
||||
dp_rank: int = 0,
|
||||
local_dp_rank: int = 0,
|
||||
):
|
||||
# Exercise the production Ray actor mixin without loading a model.
|
||||
EngineCoreActorMixin.__init__(
|
||||
self, vllm_config, addresses, dp_rank, local_dp_rank
|
||||
)
|
||||
|
||||
def _set_visible_devices(self, vllm_config: Any, local_dp_rank: int) -> None:
|
||||
pass
|
||||
|
||||
def wait_for_init(self) -> None:
|
||||
pass
|
||||
|
||||
def run(self) -> None:
|
||||
pass
|
||||
|
||||
def get_nixl_side_channel_host(self) -> str | None:
|
||||
return os.environ.get("VLLM_NIXL_SIDE_CHANNEL_HOST")
|
||||
|
||||
def get_addresses(self) -> tuple[list[str], list[str]]:
|
||||
"""Return the addresses snapshot the actor was constructed with.
|
||||
|
||||
Used by the Ray-DP regression test to assert that no ``tcp://host:0``
|
||||
placeholders were pickled into the actor at ``.remote()`` time.
|
||||
"""
|
||||
return list(self.addresses.inputs), list(self.addresses.outputs)
|
||||
|
||||
|
||||
# Module-level stub worker for the Ray-DP regression test. Must be importable
|
||||
# by ``multiprocessing.spawn`` (no closures, no nesting). Mirrors the worker
|
||||
# in ``tests/entrypoints/test_api_server_process_manager.py``.
|
||||
def _bind_and_report_worker(listen_address, sock, args, client_config):
|
||||
"""Bind ROUTER/PULL with a kernel-assigned port, report the actual
|
||||
endpoints back via ``actual_address_pipe``, then exit."""
|
||||
ctx = zmq.Context()
|
||||
try:
|
||||
in_sock = make_zmq_socket(
|
||||
ctx, client_config["input_address"], zmq.ROUTER, bind=True
|
||||
)
|
||||
out_sock = make_zmq_socket(
|
||||
ctx, client_config["output_address"], zmq.PULL, bind=True
|
||||
)
|
||||
try:
|
||||
pipe = client_config["actual_address_pipe"]
|
||||
try:
|
||||
pipe.send(
|
||||
{
|
||||
"input_address": in_sock.getsockopt(zmq.LAST_ENDPOINT).decode(),
|
||||
"output_address": out_sock.getsockopt(
|
||||
zmq.LAST_ENDPOINT
|
||||
).decode(),
|
||||
}
|
||||
)
|
||||
finally:
|
||||
pipe.close()
|
||||
finally:
|
||||
in_sock.close(linger=0)
|
||||
out_sock.close(linger=0)
|
||||
finally:
|
||||
ctx.term()
|
||||
|
||||
|
||||
class _DummyExecutor:
|
||||
pass
|
||||
|
||||
|
||||
def test_background_resources_passes_worker_shutdown_timeout(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
timeout = 7
|
||||
monkeypatch.setenv("VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS", str(timeout))
|
||||
engine_manager = Mock()
|
||||
resources = BackgroundResources(ctx=None, engine_manager=engine_manager)
|
||||
resources()
|
||||
engine_manager.shutdown.assert_called_once_with(timeout=timeout)
|
||||
|
||||
|
||||
def _make_vllm_config() -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
parallel_config=SimpleNamespace(
|
||||
data_parallel_size=1,
|
||||
data_parallel_size_local=1,
|
||||
enable_elastic_ep=False,
|
||||
world_size=1,
|
||||
),
|
||||
model_config=SimpleNamespace(is_moe=False),
|
||||
kv_transfer_config=None,
|
||||
)
|
||||
|
||||
|
||||
def _make_addresses() -> EngineZmqAddresses:
|
||||
return EngineZmqAddresses(
|
||||
inputs=["tcp://127.0.0.1:12345"],
|
||||
outputs=["tcp://127.0.0.1:12346"],
|
||||
)
|
||||
|
||||
|
||||
def _make_cpu_placement_group():
|
||||
pg = ray.util.placement_group(
|
||||
[{"CPU": 0.001}, {"CPU": 1.0}],
|
||||
strategy="PACK",
|
||||
)
|
||||
ray.get(pg.ready())
|
||||
return pg
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_context():
|
||||
started_ray = False
|
||||
if not ray.is_initialized():
|
||||
project_root = str(Path(__file__).resolve().parents[3])
|
||||
ray.init(
|
||||
num_cpus=2,
|
||||
runtime_env={"env_vars": {"PYTHONPATH": project_root}},
|
||||
log_to_driver=False,
|
||||
)
|
||||
started_ray = True
|
||||
|
||||
yield
|
||||
|
||||
if started_ray:
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("ray_context")
|
||||
def test_driver_nixl_side_channel_host_does_not_leak_to_engine_core_actor(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
driver_marker = f"driver-only-nixl-host-{uuid.uuid4()}"
|
||||
created_placement_groups: list[Any] = []
|
||||
manager: CoreEngineActorManager | None = None
|
||||
|
||||
def create_dp_placement_groups(vllm_config: Any):
|
||||
pg = _make_cpu_placement_group()
|
||||
created_placement_groups.append(pg)
|
||||
return [pg], [0]
|
||||
|
||||
monkeypatch.setenv("VLLM_NIXL_SIDE_CHANNEL_HOST", driver_marker)
|
||||
monkeypatch.setattr("vllm.v1.engine.core.EngineCoreActor", _StubEngineCoreActor)
|
||||
monkeypatch.setattr(
|
||||
CoreEngineActorManager,
|
||||
"create_dp_placement_groups",
|
||||
staticmethod(create_dp_placement_groups),
|
||||
)
|
||||
|
||||
try:
|
||||
manager = CoreEngineActorManager(
|
||||
vllm_config=_make_vllm_config(),
|
||||
addresses=_make_addresses(),
|
||||
executor_class=_DummyExecutor,
|
||||
log_stats=False,
|
||||
)
|
||||
actor = manager.local_engine_actors[0]
|
||||
actor_host = ray.get(actor.get_nixl_side_channel_host.remote())
|
||||
node_host = ray.util.get_node_ip_address()
|
||||
|
||||
assert actor_host != driver_marker
|
||||
assert actor_host == node_host
|
||||
finally:
|
||||
if manager is not None:
|
||||
manager.shutdown()
|
||||
else:
|
||||
for pg in created_placement_groups:
|
||||
ray.util.remove_placement_group(pg)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_context_dp2():
|
||||
"""Ray context sized for two stub actors (each PG needs ~1 CPU)."""
|
||||
started_ray = False
|
||||
if not ray.is_initialized():
|
||||
project_root = str(Path(__file__).resolve().parents[3])
|
||||
ray.init(
|
||||
num_cpus=4,
|
||||
runtime_env={"env_vars": {"PYTHONPATH": project_root}},
|
||||
log_to_driver=False,
|
||||
)
|
||||
started_ray = True
|
||||
|
||||
yield
|
||||
|
||||
if started_ray:
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
def _make_vllm_config_ray_dp_multinode() -> SimpleNamespace:
|
||||
"""Minimal vllm_config that drives the Ray-DP multi-API-server path:
|
||||
``data_parallel_size != data_parallel_size_local`` forces TCP placeholders
|
||||
(multi-node fan-out), and ``data_parallel_backend="ray"`` routes
|
||||
``launch_core_engines`` through the Ray branch.
|
||||
"""
|
||||
return SimpleNamespace(
|
||||
parallel_config=SimpleNamespace(
|
||||
data_parallel_size=2,
|
||||
data_parallel_size_local=1,
|
||||
data_parallel_rank=0,
|
||||
data_parallel_rank_local=None,
|
||||
data_parallel_master_ip="127.0.0.1",
|
||||
data_parallel_backend="ray",
|
||||
data_parallel_rpc_port=29550,
|
||||
local_engines_only=False,
|
||||
enable_elastic_ep=False,
|
||||
world_size=1,
|
||||
),
|
||||
model_config=SimpleNamespace(multimodal_config=None, is_moe=False),
|
||||
cache_config=SimpleNamespace(),
|
||||
needs_dp_coordinator=False,
|
||||
kv_transfer_config=None,
|
||||
# ``_apply_dp_identity_suffix`` reads and rewrites this.
|
||||
instance_id="vllm-ray-dp-regression-test",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(120)
|
||||
@pytest.mark.usefixtures("ray_context_dp2")
|
||||
def test_ray_dp_addresses_resolved_before_actor_creation(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Regression guard for the Ray-DP + multi-API-server hang from PR #42585.
|
||||
|
||||
``launch_core_engines`` Ray branch pickles ``addresses`` into each engine
|
||||
actor at ``.remote()`` time, and ``EngineCoreActorMixin._perform_handshakes``
|
||||
is a no-op, so the actor uses that pickled snapshot for the rest of its
|
||||
life. If ``run_multi_api_server`` allocates ``addresses`` as
|
||||
``tcp://host:0`` placeholders (its default), the actors hold placeholders
|
||||
forever and DEALER-connect to port 0 — ZMQ ``connect`` is async and does
|
||||
not raise, so the failure mode is a deterministic hang.
|
||||
|
||||
The Ray-DP carve-out in ``run_multi_api_server`` forces
|
||||
``defer_api_server_ports=False`` when ``data_parallel_backend == "ray"``
|
||||
so addresses are pre-allocated in the driver and Ray pickles real ports
|
||||
into each actor. This test mirrors that call-site logic and asserts the
|
||||
actors hold real (non-placeholder) endpoints. If the carve-out is
|
||||
removed without an alternative fix, the test fails.
|
||||
"""
|
||||
created_placement_groups: list[Any] = []
|
||||
|
||||
def create_dp_placement_groups(vllm_config: Any):
|
||||
pg1 = _make_cpu_placement_group()
|
||||
pg2 = _make_cpu_placement_group()
|
||||
created_placement_groups.extend([pg1, pg2])
|
||||
return [pg1, pg2], [0, 0]
|
||||
|
||||
monkeypatch.setattr("vllm.v1.engine.core.EngineCoreActor", _StubEngineCoreActor)
|
||||
monkeypatch.setattr(
|
||||
CoreEngineActorManager,
|
||||
"create_dp_placement_groups",
|
||||
staticmethod(create_dp_placement_groups),
|
||||
)
|
||||
|
||||
vllm_config = _make_vllm_config_ray_dp_multinode()
|
||||
|
||||
# Mirror run_multi_api_server's address-allocation logic. The Ray DP
|
||||
# carve-out forces pre-allocation so the addresses pickled into engine
|
||||
# actors at .remote() time are real, not ``tcp://host:0``.
|
||||
is_ray_dp = vllm_config.parallel_config.data_parallel_backend == "ray"
|
||||
addresses = get_engine_zmq_addresses(
|
||||
vllm_config,
|
||||
num_api_servers=2,
|
||||
defer_api_server_ports=not is_ray_dp,
|
||||
)
|
||||
|
||||
sock = socket.socket()
|
||||
engine_manager: CoreEngineActorManager | None = None
|
||||
actor_snapshots: list[tuple[list[str], list[str]]] = []
|
||||
api_server_manager: APIServerProcessManager | None = None
|
||||
try:
|
||||
# Ray actors are spawned here, pickling ``addresses`` into each one.
|
||||
with launch_core_engines(
|
||||
vllm_config,
|
||||
executor_class=_DummyExecutor,
|
||||
log_stats=False,
|
||||
addresses=addresses,
|
||||
num_api_servers=2,
|
||||
) as (
|
||||
engine_manager,
|
||||
_coordinator,
|
||||
_addresses_out,
|
||||
_tensor_queue,
|
||||
):
|
||||
assert isinstance(engine_manager, CoreEngineActorManager)
|
||||
|
||||
# API-server children bind to the pre-allocated ports.
|
||||
api_server_manager = APIServerProcessManager(
|
||||
listen_address="tcp://127.0.0.1:0",
|
||||
sock=sock,
|
||||
args="test_args",
|
||||
num_servers=2,
|
||||
input_addresses=addresses.inputs,
|
||||
output_addresses=addresses.outputs,
|
||||
target_server_fn=_bind_and_report_worker,
|
||||
)
|
||||
|
||||
# run_multi_api_server skips ``gather_actual_addresses`` for
|
||||
# Ray DP (addresses are already real). Mirror that.
|
||||
if not is_ray_dp:
|
||||
actual_inputs, actual_outputs = (
|
||||
api_server_manager.gather_actual_addresses(timeout=15.0)
|
||||
)
|
||||
addresses.inputs = actual_inputs
|
||||
addresses.outputs = actual_outputs
|
||||
|
||||
# Snapshot what each Ray actor actually holds.
|
||||
actors = (
|
||||
engine_manager.local_engine_actors + engine_manager.remote_engine_actors
|
||||
)
|
||||
actor_snapshots = ray.get(
|
||||
[actor.get_addresses.remote() for actor in actors]
|
||||
)
|
||||
finally:
|
||||
if api_server_manager is not None:
|
||||
api_server_manager.shutdown()
|
||||
time.sleep(0.2)
|
||||
sock.close()
|
||||
if engine_manager is not None:
|
||||
engine_manager.shutdown()
|
||||
else:
|
||||
for pg in created_placement_groups:
|
||||
ray.util.remove_placement_group(pg)
|
||||
|
||||
# Every Ray actor must hold real, non-placeholder addresses.
|
||||
assert actor_snapshots, "expected at least one Ray actor to be created"
|
||||
for actor_inputs, actor_outputs in actor_snapshots:
|
||||
for url in actor_inputs + actor_outputs:
|
||||
scheme, _host, port = split_zmq_path(url)
|
||||
assert scheme == "tcp", url
|
||||
assert port and int(port) > 0, (
|
||||
f"Ray actor was pickled with placeholder address {url!r}; "
|
||||
"``run_multi_api_server`` must pre-allocate ports for the "
|
||||
"Ray DP backend so the actors hold real endpoints by the "
|
||||
"time they DEALER-connect. See PR #42585 / Ray-DP "
|
||||
"multi-API-server regression."
|
||||
)
|
||||
@@ -0,0 +1,119 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for VLLM_RAY_DP_PLACEMENT_NODE_IPS."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
import vllm.v1.engine.utils as utils
|
||||
from vllm.v1.engine.utils import CoreEngineActorManager
|
||||
|
||||
|
||||
def _vllm_config(
|
||||
*, dp_size, dp_local, master_ip, world_size=1, all2all_backend="naive"
|
||||
):
|
||||
parallel = SimpleNamespace(
|
||||
data_parallel_master_ip=master_ip,
|
||||
data_parallel_size=dp_size,
|
||||
data_parallel_size_local=dp_local,
|
||||
world_size=world_size,
|
||||
all2all_backend=all2all_backend,
|
||||
)
|
||||
return SimpleNamespace(parallel_config=parallel)
|
||||
|
||||
|
||||
def _resources(node_gpus: dict[str, int]):
|
||||
# node_gpus: {ip: gpu_count}; plus a CPU-only head node.
|
||||
res = {
|
||||
f"id-{ip}": {"GPU": float(g), f"node:{ip}": 1.0} for ip, g in node_gpus.items()
|
||||
}
|
||||
res["id-head"] = {
|
||||
"CPU": 8.0,
|
||||
"node:__internal_head__": 1.0,
|
||||
"node:10.9.9.9": 1.0,
|
||||
}
|
||||
return res
|
||||
|
||||
|
||||
def _run(cfg, resources):
|
||||
created = []
|
||||
|
||||
def fake_pg(name, strategy, bundles):
|
||||
created.append({"name": name, "strategy": strategy, "bundles": bundles})
|
||||
return object()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"ray._private.state.available_resources_per_node",
|
||||
return_value=resources,
|
||||
),
|
||||
patch.object(utils, "current_platform", SimpleNamespace(ray_device_key="GPU")),
|
||||
patch("ray.util.placement_group", side_effect=fake_pg),
|
||||
):
|
||||
pgs, local_ranks = CoreEngineActorManager.create_dp_placement_groups(cfg)
|
||||
return pgs, local_ranks, created
|
||||
|
||||
|
||||
def _pinned_ips(created):
|
||||
return {
|
||||
key.split(":", 1)[1]
|
||||
for pg in created
|
||||
for bundle in pg["bundles"]
|
||||
for key in bundle
|
||||
if key.startswith("node:")
|
||||
}
|
||||
|
||||
|
||||
def test_allowlist_confines_dp_to_listed_nodes(monkeypatch):
|
||||
monkeypatch.setenv("VLLM_RAY_DP_PLACEMENT_NODE_IPS", "10.0.0.1,10.0.0.3")
|
||||
resources = _resources({"10.0.0.1": 8, "10.0.0.2": 8, "10.0.0.3": 8, "10.0.0.4": 8})
|
||||
cfg = _vllm_config(dp_size=16, dp_local=8, master_ip="10.0.0.1")
|
||||
|
||||
pgs, _, created = _run(cfg, resources)
|
||||
|
||||
assert len(pgs) == 16 # 8 on .1 (master) + 8 on .3
|
||||
assert _pinned_ips(created) <= {"10.0.0.1", "10.0.0.3"}
|
||||
|
||||
|
||||
def test_empty_allowlist_is_noop(monkeypatch):
|
||||
monkeypatch.delenv("VLLM_RAY_DP_PLACEMENT_NODE_IPS", raising=False)
|
||||
resources = _resources({"10.0.0.1": 8, "10.0.0.2": 8})
|
||||
cfg = _vllm_config(dp_size=16, dp_local=8, master_ip="10.0.0.1")
|
||||
|
||||
pgs, _, created = _run(cfg, resources)
|
||||
|
||||
assert len(pgs) == 16
|
||||
assert _pinned_ips(created) == {"10.0.0.1", "10.0.0.2"} # all nodes used
|
||||
|
||||
|
||||
def test_master_auto_added_with_warning(monkeypatch):
|
||||
# Allowlist omits the master; vLLM must still keep it and warn.
|
||||
monkeypatch.setenv("VLLM_RAY_DP_PLACEMENT_NODE_IPS", "10.0.0.3")
|
||||
resources = _resources({"10.0.0.1": 8, "10.0.0.3": 8})
|
||||
cfg = _vllm_config(dp_size=16, dp_local=8, master_ip="10.0.0.1")
|
||||
_, _, created = _run(cfg, resources)
|
||||
|
||||
assert _pinned_ips(created) == {"10.0.0.1", "10.0.0.3"}
|
||||
|
||||
|
||||
def test_allowlist_isolates_two_engines(monkeypatch):
|
||||
# Engine B is confined to .2/.4, so it can never touch engine A's master .1.
|
||||
monkeypatch.setenv("VLLM_RAY_DP_PLACEMENT_NODE_IPS", "10.0.0.2,10.0.0.4")
|
||||
resources = _resources({"10.0.0.1": 8, "10.0.0.2": 8, "10.0.0.3": 8, "10.0.0.4": 8})
|
||||
cfg = _vllm_config(dp_size=16, dp_local=8, master_ip="10.0.0.2")
|
||||
|
||||
_, _, created = _run(cfg, resources)
|
||||
|
||||
assert _pinned_ips(created) <= {"10.0.0.2", "10.0.0.4"}
|
||||
|
||||
|
||||
def test_allowlist_too_small_raises(monkeypatch):
|
||||
# Master alone can't hold all ranks and no other node is allowed.
|
||||
monkeypatch.setenv("VLLM_RAY_DP_PLACEMENT_NODE_IPS", "10.0.0.1")
|
||||
resources = _resources({"10.0.0.1": 8, "10.0.0.2": 8})
|
||||
cfg = _vllm_config(dp_size=16, dp_local=8, master_ip="10.0.0.1")
|
||||
|
||||
with pytest.raises(ValueError): # not enough placement groups created
|
||||
_run(cfg, resources)
|
||||
@@ -0,0 +1,130 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from argparse import ArgumentError
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.usage.usage_lib import UsageContext
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.hashing import _xxhash
|
||||
|
||||
|
||||
def test_prefix_caching_from_cli():
|
||||
parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
|
||||
args = parser.parse_args([])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert vllm_config.cache_config.enable_prefix_caching, (
|
||||
"V1 turns on prefix caching by default."
|
||||
)
|
||||
|
||||
# Turn it off possible with flag.
|
||||
args = parser.parse_args(["--no-enable-prefix-caching"])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert not vllm_config.cache_config.enable_prefix_caching
|
||||
|
||||
# Turn it on with flag.
|
||||
args = parser.parse_args(["--enable-prefix-caching"])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert vllm_config.cache_config.enable_prefix_caching
|
||||
|
||||
# default hash algorithm is "builtin"
|
||||
assert vllm_config.cache_config.prefix_caching_hash_algo == "sha256"
|
||||
|
||||
# set hash algorithm to sha256_cbor
|
||||
args = parser.parse_args(["--prefix-caching-hash-algo", "sha256_cbor"])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert vllm_config.cache_config.prefix_caching_hash_algo == "sha256_cbor"
|
||||
|
||||
# set hash algorithm to sha256
|
||||
args = parser.parse_args(["--prefix-caching-hash-algo", "sha256"])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert vllm_config.cache_config.prefix_caching_hash_algo == "sha256"
|
||||
|
||||
# an invalid hash algorithm raises an error
|
||||
parser.exit_on_error = False
|
||||
with pytest.raises(ArgumentError):
|
||||
args = parser.parse_args(["--prefix-caching-hash-algo", "invalid"])
|
||||
|
||||
|
||||
@pytest.mark.skipif(_xxhash is None, reason="xxhash not installed")
|
||||
def test_prefix_caching_xxhash_from_cli():
|
||||
parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
|
||||
|
||||
# set hash algorithm to xxhash (pickle)
|
||||
args = parser.parse_args(["--prefix-caching-hash-algo", "xxhash"])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert vllm_config.cache_config.prefix_caching_hash_algo == "xxhash"
|
||||
|
||||
# set hash algorithm to xxhash_cbor
|
||||
args = parser.parse_args(["--prefix-caching-hash-algo", "xxhash_cbor"])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert vllm_config.cache_config.prefix_caching_hash_algo == "xxhash_cbor"
|
||||
|
||||
|
||||
def test_defaults_with_usage_context():
|
||||
engine_args = EngineArgs(model="facebook/opt-125m")
|
||||
vllm_config: VllmConfig = engine_args.create_engine_config(UsageContext.LLM_CLASS)
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.mem_constants import GiB_bytes
|
||||
|
||||
device_memory = current_platform.get_device_total_memory()
|
||||
device_name = current_platform.get_device_name().lower()
|
||||
if device_memory >= 70 * GiB_bytes and "a100" not in device_name:
|
||||
# For GPUs like H100, H200, and MI300x with >= 70GB memory
|
||||
default_llm_tokens = 16384
|
||||
default_server_tokens = 8192
|
||||
default_max_num_seqs = 1024
|
||||
else:
|
||||
default_llm_tokens = 8192
|
||||
default_server_tokens = 2048
|
||||
default_max_num_seqs = 256
|
||||
|
||||
assert vllm_config.scheduler_config.max_num_seqs == default_max_num_seqs
|
||||
assert vllm_config.scheduler_config.max_num_batched_tokens == default_llm_tokens # noqa: E501
|
||||
|
||||
engine_args = EngineArgs(model="facebook/opt-125m")
|
||||
vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
|
||||
assert vllm_config.scheduler_config.max_num_seqs == default_max_num_seqs
|
||||
assert vllm_config.scheduler_config.max_num_batched_tokens == default_server_tokens # noqa: E501
|
||||
|
||||
|
||||
def test_mm_prefix_lm_raises_batched_tokens_floor():
|
||||
"""Verify that prefix-LM multimodal models auto-raise
|
||||
max_num_batched_tokens to fit at least one multimodal item.
|
||||
|
||||
Regression test for https://github.com/vllm-project/vllm/issues/42687
|
||||
"""
|
||||
from unittest.mock import patch
|
||||
|
||||
# Simulate a prefix-LM multimodal model whose largest modality
|
||||
# (video) requires 2496 tokens — more than the 2048 default.
|
||||
fake_mm_min = (2496, "video")
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model="facebook/opt-125m",
|
||||
max_model_len=2048,
|
||||
enforce_eager=True,
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
type(engine_args),
|
||||
"_get_min_mm_batched_tokens",
|
||||
staticmethod(lambda _mc: fake_mm_min),
|
||||
),
|
||||
patch(
|
||||
"vllm.config.ModelConfig.is_multimodal_model",
|
||||
new_callable=lambda: property(lambda self: True),
|
||||
),
|
||||
patch(
|
||||
"vllm.config.ModelConfig.is_mm_prefix_lm",
|
||||
new_callable=lambda: property(lambda self: True),
|
||||
),
|
||||
):
|
||||
vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
|
||||
|
||||
assert vllm_config.scheduler_config.max_num_batched_tokens >= 2496
|
||||
@@ -0,0 +1,610 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import copy
|
||||
import time
|
||||
import uuid
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
from unittest.mock import PropertyMock, patch
|
||||
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm import SamplingParams
|
||||
from vllm.config import (
|
||||
CacheConfig,
|
||||
ECTransferConfig,
|
||||
KVTransferConfig,
|
||||
ModelConfig,
|
||||
SchedulerConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_default_torch_num_threads
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.engine.core import EngineCore
|
||||
from vllm.v1.executor.abstract import Executor
|
||||
from vllm.v1.executor.uniproc_executor import UniProcExecutor
|
||||
from vllm.v1.kv_cache_interface import KVCacheConfig
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
|
||||
from ...utils import create_new_process_for_each_test, multi_gpu_test
|
||||
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip(reason="V1 currently only supported on CUDA.", allow_module_level=True)
|
||||
|
||||
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
|
||||
TOKENIZER = AutoTokenizer.from_pretrained(MODEL_NAME)
|
||||
# test_engine_core_concurrent_batches assumes exactly 12 tokens per prompt.
|
||||
# Adjust prompt if changing model to maintain 12-token length.
|
||||
PROMPT = "I am Gyoubu Masataka Oniwa"
|
||||
PROMPT_TOKENS = TOKENIZER(PROMPT).input_ids
|
||||
|
||||
_REQUEST_COUNTER = 0
|
||||
|
||||
|
||||
def make_request() -> EngineCoreRequest:
|
||||
global _REQUEST_COUNTER
|
||||
_REQUEST_COUNTER += 1
|
||||
request_id = f"request-{_REQUEST_COUNTER}"
|
||||
return EngineCoreRequest(
|
||||
request_id=request_id,
|
||||
external_req_id=f"{request_id}-{uuid.uuid4()}",
|
||||
prompt_token_ids=PROMPT_TOKENS,
|
||||
mm_features=None,
|
||||
sampling_params=SamplingParams(),
|
||||
pooling_params=None,
|
||||
arrival_time=time.time(),
|
||||
lora_request=None,
|
||||
cache_salt=None,
|
||||
data_parallel_rank=None,
|
||||
)
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_engine_core():
|
||||
"""Setup the EngineCore."""
|
||||
engine_args = EngineArgs(model=MODEL_NAME)
|
||||
vllm_config = engine_args.create_engine_config()
|
||||
executor_class = Executor.get_class(vllm_config)
|
||||
|
||||
with set_default_torch_num_threads(1):
|
||||
engine_core = EngineCore(
|
||||
vllm_config=vllm_config, executor_class=executor_class, log_stats=True
|
||||
)
|
||||
"""Test basic request lifecycle."""
|
||||
|
||||
# First request.
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(make_request()))
|
||||
assert len(engine_core.scheduler.waiting) == 1
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 1
|
||||
|
||||
# Second request.
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(make_request()))
|
||||
assert len(engine_core.scheduler.waiting) == 1
|
||||
assert len(engine_core.scheduler.running) == 1
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 2
|
||||
|
||||
# Add two requests in a row.
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(make_request()))
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(make_request()))
|
||||
assert len(engine_core.scheduler.waiting) == 2
|
||||
assert len(engine_core.scheduler.running) == 2
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 4
|
||||
|
||||
# Loop through until they are all done.
|
||||
while (outs := engine_core.step_fn()[0].get(0)) and outs.outputs:
|
||||
pass
|
||||
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
"""Test abort cycle."""
|
||||
|
||||
# Basic abort.
|
||||
req = make_request()
|
||||
request_id = req.request_id
|
||||
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req))
|
||||
assert len(engine_core.scheduler.waiting) == 1
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
assert engine_core.scheduler.has_unfinished_requests()
|
||||
assert not engine_core.scheduler.has_finished_requests()
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 1
|
||||
assert engine_core.scheduler.has_unfinished_requests()
|
||||
assert not engine_core.scheduler.has_finished_requests()
|
||||
|
||||
engine_core.abort_requests([request_id])
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
assert not engine_core.scheduler.has_unfinished_requests()
|
||||
assert engine_core.scheduler.has_finished_requests()
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert not engine_core.scheduler.has_unfinished_requests()
|
||||
assert not engine_core.scheduler.has_finished_requests()
|
||||
|
||||
# Add, step, abort 1 of the 3.
|
||||
req0 = make_request()
|
||||
req1 = make_request()
|
||||
req2 = make_request()
|
||||
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req0))
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req1))
|
||||
assert len(engine_core.scheduler.waiting) == 2
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 2
|
||||
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req2))
|
||||
assert len(engine_core.scheduler.waiting) == 1
|
||||
assert len(engine_core.scheduler.running) == 2
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 3
|
||||
|
||||
# Abort just one.
|
||||
engine_core.abort_requests([req1.request_id])
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 2
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 2
|
||||
|
||||
# Abort the other requests at the same time.
|
||||
engine_core.abort_requests([req2.request_id, req0.request_id])
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
|
||||
# Sending duplicate requests with same request_id
|
||||
req0 = make_request()
|
||||
req1 = make_request()
|
||||
req0.request_id = req1.request_id = "test"
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req0))
|
||||
|
||||
while engine_core.scheduler.has_requests():
|
||||
engine_core.step_fn()
|
||||
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req1))
|
||||
while engine_core.scheduler.has_requests():
|
||||
engine_core.step_fn()
|
||||
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_engine_core_advanced_sampling():
|
||||
"""
|
||||
A basic end-to-end test to verify that the engine functions correctly
|
||||
when additional sampling parameters, such as top_p, min_tokens, and
|
||||
presence_penalty, are set.
|
||||
"""
|
||||
"""Setup the EngineCore."""
|
||||
engine_args = EngineArgs(model=MODEL_NAME)
|
||||
vllm_config = engine_args.create_engine_config()
|
||||
executor_class = Executor.get_class(vllm_config)
|
||||
|
||||
with set_default_torch_num_threads(1):
|
||||
engine_core = EngineCore(
|
||||
vllm_config=vllm_config, executor_class=executor_class, log_stats=True
|
||||
)
|
||||
"""Test basic request lifecycle."""
|
||||
# First request.
|
||||
request: EngineCoreRequest = make_request()
|
||||
request.sampling_params = SamplingParams(
|
||||
min_tokens=4,
|
||||
presence_penalty=1.0,
|
||||
frequency_penalty=1.0,
|
||||
repetition_penalty=0.1,
|
||||
stop_token_ids=[1001, 1002],
|
||||
)
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(request))
|
||||
|
||||
def _check_engine_state():
|
||||
assert len(engine_core.scheduler.waiting) == 1
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
# Loop through until they are all done.
|
||||
while engine_core.scheduler.has_requests():
|
||||
engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
|
||||
_check_engine_state()
|
||||
|
||||
# Second request.
|
||||
request2 = make_request()
|
||||
request2.sampling_params = SamplingParams(
|
||||
top_p=0.99,
|
||||
top_k=50,
|
||||
)
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(request2))
|
||||
_check_engine_state()
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_engine_core_concurrent_batches():
|
||||
"""
|
||||
Test that the engine can handle multiple concurrent batches.
|
||||
"""
|
||||
|
||||
def make_request_with_max_tokens(req_id: str, max_tokens: int) -> EngineCoreRequest:
|
||||
request = make_request()
|
||||
request.request_id = req_id
|
||||
request.sampling_params.max_tokens = max_tokens
|
||||
return request
|
||||
|
||||
class DummyExecutor(UniProcExecutor):
|
||||
def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None:
|
||||
super().initialize_from_config(kv_cache_configs)
|
||||
|
||||
# Create a thread pool with a single worker
|
||||
self.thread_pool = ThreadPoolExecutor(max_workers=1)
|
||||
|
||||
def execute_model(
|
||||
self,
|
||||
scheduler_output,
|
||||
non_block=False,
|
||||
) -> Future[ModelRunnerOutput | None]:
|
||||
"""Make execute_model non-blocking."""
|
||||
|
||||
# DummyExecutor used only for testing async case.
|
||||
assert non_block
|
||||
|
||||
def _execute():
|
||||
output = self.collective_rpc("execute_model", args=(scheduler_output,))
|
||||
# Make a copy because output[0] may be reused
|
||||
# by the next batch.
|
||||
return copy.deepcopy(output[0])
|
||||
|
||||
# Use the thread pool instead of creating a new thread
|
||||
return self.thread_pool.submit(_execute)
|
||||
|
||||
def sample_tokens(
|
||||
self, grammar_output, non_block=False
|
||||
) -> Future[ModelRunnerOutput]:
|
||||
"""Make sample_tokens non-blocking."""
|
||||
|
||||
# DummyExecutor used only for testing async case.
|
||||
assert non_block
|
||||
|
||||
def _execute():
|
||||
output = self.collective_rpc("sample_tokens", args=(grammar_output,))
|
||||
# Make a copy because output[0] may be reused
|
||||
# by the next batch.
|
||||
return copy.deepcopy(output[0])
|
||||
|
||||
# Use the thread pool instead of creating a new thread
|
||||
return self.thread_pool.submit(_execute)
|
||||
|
||||
def shutdown(self):
|
||||
if hasattr(self, "thread_pool"):
|
||||
self.thread_pool.shutdown(wait=False)
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=MODEL_NAME,
|
||||
# To test concurrent batches.
|
||||
max_num_seqs=2,
|
||||
# Avoid all requests being scheduled once.
|
||||
enable_prefix_caching=False,
|
||||
max_num_batched_tokens=10,
|
||||
# Reduce startup time.
|
||||
enforce_eager=True,
|
||||
# Test concurrent batch behaviour independently of async scheduling.
|
||||
async_scheduling=False,
|
||||
)
|
||||
vllm_config = engine_args.create_engine_config()
|
||||
# Force two concurrent batches to exercise the batch queue independently
|
||||
# of async scheduling (which is disabled above).
|
||||
with (
|
||||
set_default_torch_num_threads(1),
|
||||
patch.object(
|
||||
VllmConfig,
|
||||
"max_concurrent_batches",
|
||||
new_callable=PropertyMock,
|
||||
return_value=2,
|
||||
),
|
||||
):
|
||||
engine_core = EngineCore(
|
||||
vllm_config=vllm_config, log_stats=False, executor_class=DummyExecutor
|
||||
)
|
||||
assert engine_core.batch_queue is not None
|
||||
|
||||
# Add two requests in a row. Each request have 12 prompt tokens.
|
||||
req0 = make_request_with_max_tokens("0", 5)
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req0))
|
||||
req1 = make_request_with_max_tokens("1", 5)
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req1))
|
||||
|
||||
# Schedule Batch 1: (10, req0)
|
||||
assert engine_core.step_with_batch_queue()[0] is None
|
||||
assert len(engine_core.batch_queue) == 1
|
||||
scheduler_output = engine_core.batch_queue[-1][1]
|
||||
assert scheduler_output.num_scheduled_tokens["0"] == 10
|
||||
# num_computed_tokens should have been updated immediately.
|
||||
assert engine_core.scheduler.requests[req0.request_id].num_computed_tokens == 10
|
||||
|
||||
# Schedule Batch 2: (2, req0), (8, req1)
|
||||
assert engine_core.step_with_batch_queue()[0] == {}
|
||||
assert len(engine_core.batch_queue) == 1
|
||||
scheduler_output = engine_core.batch_queue[-1][1]
|
||||
assert scheduler_output.num_scheduled_tokens["0"] == 2
|
||||
assert scheduler_output.num_scheduled_tokens["1"] == 8
|
||||
# num_computed_tokens should have been updated immediately.
|
||||
assert engine_core.scheduler.requests["0"].num_computed_tokens == 12
|
||||
assert engine_core.scheduler.requests["1"].num_computed_tokens == 8
|
||||
|
||||
assert engine_core.scheduler.get_num_unfinished_requests() == 2
|
||||
|
||||
# Finish Batch 1 and schedule Batch 3: (4, req1).
|
||||
# Note that req0 cannot be scheduled
|
||||
# because it is in the decoding stage now.
|
||||
engine_core.step_with_batch_queue()
|
||||
assert len(engine_core.batch_queue) == 1
|
||||
scheduler_output = engine_core.batch_queue[-1][1]
|
||||
assert scheduler_output.num_scheduled_tokens["1"] == 4
|
||||
|
||||
# Finish Batch 2. Get first token of req0.
|
||||
# Schedule Batch 4: (1, req0).
|
||||
output = engine_core.step_with_batch_queue()[0].get(0)
|
||||
assert output is not None
|
||||
assert len(output.outputs) == 1
|
||||
assert engine_core.scheduler.requests[req0.request_id].num_tokens == 13
|
||||
scheduler_output = engine_core.batch_queue[-1][1]
|
||||
assert scheduler_output.num_scheduled_tokens["0"] == 1
|
||||
|
||||
# Finish Batch 3. Get first token of req1. Schedule Batch 5: (1, req1).
|
||||
output = engine_core.step_with_batch_queue()[0].get(0)
|
||||
assert output is not None
|
||||
assert len(output.outputs) == 1
|
||||
assert engine_core.scheduler.requests[req1.request_id].num_tokens == 13
|
||||
scheduler_output = engine_core.batch_queue[-1][1]
|
||||
assert scheduler_output.num_scheduled_tokens["1"] == 1
|
||||
|
||||
# Loop until req0 is finished.
|
||||
req_id = 0
|
||||
expected_num_tokens = [
|
||||
engine_core.scheduler.requests["0"].num_tokens + 1,
|
||||
engine_core.scheduler.requests["1"].num_tokens + 1,
|
||||
]
|
||||
while engine_core.scheduler.get_num_unfinished_requests() == 2:
|
||||
output = engine_core.step_with_batch_queue()[0]
|
||||
# Every step consumes an output.
|
||||
assert output is not None
|
||||
assert len(output[0].outputs) == 1
|
||||
if req_id in engine_core.scheduler.requests:
|
||||
assert (
|
||||
engine_core.scheduler.requests[req_id].num_tokens
|
||||
== expected_num_tokens[req_id]
|
||||
)
|
||||
expected_num_tokens[req_id] += 1
|
||||
req_id = (req_id + 1) % 2
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
def test_engine_core_tp():
|
||||
"""
|
||||
Test engine can initialize worker in tp properly
|
||||
"""
|
||||
|
||||
"""Setup the EngineCore."""
|
||||
engine_args = EngineArgs(
|
||||
model=MODEL_NAME,
|
||||
tensor_parallel_size=2,
|
||||
# Reduce startup time.
|
||||
enforce_eager=True,
|
||||
)
|
||||
vllm_config = engine_args.create_engine_config()
|
||||
executor_class = Executor.get_class(vllm_config)
|
||||
|
||||
with set_default_torch_num_threads(1):
|
||||
engine_core = EngineCore(
|
||||
vllm_config=vllm_config, executor_class=executor_class, log_stats=True
|
||||
)
|
||||
|
||||
def get_worker_cache_config_field(worker, key: str):
|
||||
return getattr(worker.cache_config, key)
|
||||
|
||||
num_gpu_blocks = engine_core.collective_rpc(
|
||||
get_worker_cache_config_field, args=("num_gpu_blocks",)
|
||||
)
|
||||
num_cpu_blocks = engine_core.collective_rpc(
|
||||
get_worker_cache_config_field, args=("num_cpu_blocks",)
|
||||
)
|
||||
assert all(x is not None for x in num_gpu_blocks)
|
||||
assert all(x is not None for x in num_cpu_blocks)
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_engine_core_invalid_request_id_type():
|
||||
"""Test that engine raises TypeError for non-string request_id."""
|
||||
engine_args = EngineArgs(model=MODEL_NAME)
|
||||
vllm_config = engine_args.create_engine_config()
|
||||
executor_class = Executor.get_class(vllm_config)
|
||||
|
||||
with set_default_torch_num_threads(1):
|
||||
engine_core = EngineCore(
|
||||
vllm_config=vllm_config, executor_class=executor_class, log_stats=True
|
||||
)
|
||||
|
||||
# Test with UUID object (common mistake)
|
||||
uuid_request = make_request()
|
||||
uuid_request.request_id = uuid.uuid4() # UUID object instead of string
|
||||
|
||||
with pytest.raises(TypeError, match="request_id must be a string, got.*UUID"):
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(uuid_request))
|
||||
|
||||
# Test with integer
|
||||
int_request = make_request()
|
||||
int_request.request_id = 12345
|
||||
|
||||
with pytest.raises(TypeError, match="request_id must be a string, got.*int"):
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(int_request))
|
||||
|
||||
# Test with None
|
||||
none_request = make_request()
|
||||
none_request.request_id = None
|
||||
|
||||
with pytest.raises(TypeError, match="request_id must be a string, got.*NoneType"):
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(none_request))
|
||||
|
||||
# Verify engine is still functional after errors
|
||||
valid_request = make_request()
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(valid_request))
|
||||
assert len(engine_core.scheduler.waiting) == 1
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
@pytest.mark.parametrize(
|
||||
("ec_role", "gpu_memory_utilization", "enable_prefix_caching"),
|
||||
[
|
||||
("ec_producer", 0.01, False),
|
||||
# NOTE: ec_producer never allows prefix caching
|
||||
("ec_consumer", 0.7, True),
|
||||
("ec_consumer", 0.7, False),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("use_kv_connector", [False, True])
|
||||
def test_encoder_instance_zero_kv_cache(
|
||||
ec_role: str,
|
||||
gpu_memory_utilization: float,
|
||||
enable_prefix_caching: bool,
|
||||
use_kv_connector: bool,
|
||||
):
|
||||
"""EPD (Encoder-Prefill-Decode) Encoder-cache-specific tests
|
||||
|
||||
This test verifies encoder-only instance initializes with 0 KV cache blocks.
|
||||
Under EPD disagg mode, Encoder instances (EC producer role) only execute
|
||||
vision encoder, so they don't need KV cache for text generation.
|
||||
"""
|
||||
# Form vllm config
|
||||
model_config = ModelConfig(
|
||||
model="llava-hf/llava-1.5-7b-hf", # Multimodal model
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
dtype="float16",
|
||||
seed=42,
|
||||
)
|
||||
scheduler_config = SchedulerConfig(
|
||||
max_num_seqs=10,
|
||||
max_num_batched_tokens=512,
|
||||
max_model_len=512,
|
||||
disable_hybrid_kv_cache_manager=True,
|
||||
is_encoder_decoder=model_config.is_encoder_decoder,
|
||||
)
|
||||
cache_config = CacheConfig(
|
||||
block_size=16,
|
||||
gpu_memory_utilization=gpu_memory_utilization,
|
||||
cache_dtype="auto",
|
||||
enable_prefix_caching=enable_prefix_caching,
|
||||
)
|
||||
kv_transfer_config = (
|
||||
KVTransferConfig(
|
||||
kv_connector="ExampleConnector",
|
||||
kv_role="kv_both",
|
||||
kv_connector_extra_config={"shared_storage_path": "local_storage"},
|
||||
)
|
||||
if use_kv_connector
|
||||
else None
|
||||
)
|
||||
ec_transfer_config = ECTransferConfig(
|
||||
ec_connector="ECExampleConnector",
|
||||
ec_role=ec_role,
|
||||
ec_connector_extra_config={"shared_storage_path": "/tmp/ec_test_encoder"},
|
||||
)
|
||||
|
||||
vllm_config = VllmConfig(
|
||||
model_config=model_config,
|
||||
cache_config=cache_config,
|
||||
scheduler_config=scheduler_config,
|
||||
kv_transfer_config=kv_transfer_config,
|
||||
ec_transfer_config=ec_transfer_config,
|
||||
)
|
||||
|
||||
executor_class = Executor.get_class(vllm_config)
|
||||
print(f"executor_class: {executor_class}")
|
||||
|
||||
with set_default_torch_num_threads(1):
|
||||
engine_core = EngineCore(
|
||||
vllm_config=vllm_config, executor_class=executor_class, log_stats=True
|
||||
)
|
||||
|
||||
# Check encoder cache manager exists
|
||||
assert engine_core.scheduler.encoder_cache_manager is not None, (
|
||||
"encoder_cache_manager should exist"
|
||||
)
|
||||
|
||||
if ec_role == "ec_producer":
|
||||
# Check 1: num_blocks should be 0
|
||||
# NOTE: num_blocks=1 as BlockPool always needs a null_block.
|
||||
kv_cache_config = engine_core.scheduler.kv_cache_manager.kv_cache_config
|
||||
print(f"kv_cache_config: {kv_cache_config}")
|
||||
assert kv_cache_config.num_blocks == 1, (
|
||||
f"ec_producer should only have 1 KV blocks, "
|
||||
f"got {kv_cache_config.num_blocks}"
|
||||
)
|
||||
|
||||
# Check 2: kv_cache_groups should be empty
|
||||
assert len(kv_cache_config.kv_cache_groups) == 0, (
|
||||
f"ec_producer should have 0 KV cache groups, "
|
||||
f"got {len(kv_cache_config.kv_cache_groups)}"
|
||||
)
|
||||
|
||||
# Check 3: kv_cache_tensors should be empty
|
||||
assert len(kv_cache_config.kv_cache_tensors) == 0, (
|
||||
f"Encoder instance should have 0 KV cache tensors, "
|
||||
f"got {len(kv_cache_config.kv_cache_tensors)}"
|
||||
)
|
||||
|
||||
# Check 4: Verify EC connector is initialized and is producer
|
||||
assert engine_core.scheduler.ec_connector is not None, (
|
||||
"Encoder instance should have EC connector"
|
||||
)
|
||||
assert engine_core.scheduler.ec_connector.is_producer, (
|
||||
"Encoder instance EC connector should be producer"
|
||||
)
|
||||
|
||||
# Check 5: Verify chunked prefill is disabled
|
||||
assert not vllm_config.scheduler_config.enable_chunked_prefill, (
|
||||
"Encoder instance should disable chunked prefill (no KV cache)"
|
||||
)
|
||||
|
||||
elif ec_role == "ec_consumer":
|
||||
# Check 1: num_blocks should be > 1
|
||||
kv_cache_config = engine_core.scheduler.kv_cache_manager.kv_cache_config
|
||||
print(f"kv_cache_config: {kv_cache_config}")
|
||||
assert kv_cache_config.num_blocks > 1, (
|
||||
f"ec_consumer should have >1 KV blocks, got {kv_cache_config.num_blocks}"
|
||||
)
|
||||
|
||||
# Check 2: kv_cache_groups should NOT be empty
|
||||
assert len(kv_cache_config.kv_cache_groups) > 0, (
|
||||
f"ec_consumer should have KV cache groups, "
|
||||
f"got {len(kv_cache_config.kv_cache_groups)}"
|
||||
)
|
||||
|
||||
# Check 3: Verify EC connector is consumer
|
||||
assert engine_core.scheduler.ec_connector is not None, (
|
||||
"Consumer instance should have EC connector"
|
||||
)
|
||||
assert not engine_core.scheduler.ec_connector.is_producer, (
|
||||
"Consumer instance EC connector should be consumer"
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,196 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.engine.detokenizer import IncrementalDetokenizer
|
||||
|
||||
# ruff: noqa: E501
|
||||
|
||||
|
||||
def test_fast_inc_detok_invalid_utf8_err_case():
|
||||
"""
|
||||
Test edge case where tokenizer can produce non-monotonic,
|
||||
invalid UTF-8 output, which breaks the internal state of
|
||||
tokenizers' DecodeStream.
|
||||
See https://github.com/vllm-project/vllm/issues/17448.
|
||||
|
||||
Thanks to reproducer from @fpaupier:
|
||||
https://gist.github.com/fpaupier/0ed1375bd7633c5be6c894b1c7ac1be3.
|
||||
"""
|
||||
tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-1b-it")
|
||||
|
||||
# Create a test request
|
||||
prompt_token_ids = [107, 4606, 236787, 107]
|
||||
params = SamplingParams(skip_special_tokens=True)
|
||||
request = EngineCoreRequest(
|
||||
request_id="test",
|
||||
external_req_id="test-ext",
|
||||
prompt_token_ids=prompt_token_ids,
|
||||
mm_features=None,
|
||||
sampling_params=params,
|
||||
pooling_params=None,
|
||||
arrival_time=0.0,
|
||||
lora_request=None,
|
||||
cache_salt=None,
|
||||
data_parallel_rank=None,
|
||||
)
|
||||
|
||||
detokenizer = IncrementalDetokenizer.from_new_request(tokenizer, request)
|
||||
|
||||
assert detokenizer.__class__.__name__ == "FastIncrementalDetokenizer", (
|
||||
"Should use FastIncrementalDetokenizer by default"
|
||||
)
|
||||
|
||||
# Process tokens incrementally
|
||||
test_tokens = [
|
||||
236840,
|
||||
107,
|
||||
138,
|
||||
236782,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
6265,
|
||||
1083,
|
||||
623,
|
||||
121908,
|
||||
147418,
|
||||
827,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
6265,
|
||||
236779,
|
||||
2084,
|
||||
1083,
|
||||
623,
|
||||
203292,
|
||||
827,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
6265,
|
||||
236779,
|
||||
7777,
|
||||
1083,
|
||||
623,
|
||||
121908,
|
||||
147418,
|
||||
569,
|
||||
537,
|
||||
236789,
|
||||
65880,
|
||||
569,
|
||||
537,
|
||||
236789,
|
||||
62580,
|
||||
853,
|
||||
115693,
|
||||
210118,
|
||||
35178,
|
||||
16055,
|
||||
1270,
|
||||
759,
|
||||
215817,
|
||||
4758,
|
||||
1925,
|
||||
1117,
|
||||
827,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
5654,
|
||||
1083,
|
||||
623,
|
||||
110733,
|
||||
46291,
|
||||
827,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
5654,
|
||||
236779,
|
||||
2084,
|
||||
1083,
|
||||
623,
|
||||
136955,
|
||||
56731,
|
||||
827,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
5654,
|
||||
236779,
|
||||
7777,
|
||||
1083,
|
||||
623,
|
||||
194776,
|
||||
2947,
|
||||
496,
|
||||
109811,
|
||||
1608,
|
||||
890,
|
||||
215817,
|
||||
4758,
|
||||
1925,
|
||||
1117,
|
||||
2789,
|
||||
432,
|
||||
398,
|
||||
602,
|
||||
31118,
|
||||
569,
|
||||
124866,
|
||||
134772,
|
||||
509,
|
||||
19478,
|
||||
1640,
|
||||
33779,
|
||||
236743,
|
||||
236770,
|
||||
236819,
|
||||
236825,
|
||||
236771,
|
||||
432,
|
||||
398,
|
||||
432,
|
||||
237167,
|
||||
827,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
77984,
|
||||
1083,
|
||||
623,
|
||||
2709,
|
||||
236745,
|
||||
2555,
|
||||
513,
|
||||
236789,
|
||||
602,
|
||||
31118,
|
||||
569,
|
||||
]
|
||||
|
||||
output = ""
|
||||
for i, token_id in enumerate(test_tokens):
|
||||
detokenizer.update([token_id], False)
|
||||
|
||||
finished = i == len(test_tokens) - 1
|
||||
output += detokenizer.get_next_output_text(finished, delta=True)
|
||||
|
||||
assert (
|
||||
output
|
||||
== r"""[
|
||||
{
|
||||
"source": "Résultats",
|
||||
"source_type": "CONCEPT",
|
||||
"source_description": "Résultats de l'analyse de l'impact des opérations israéliennes sur la frontière libanaise",
|
||||
"target": "Israël",
|
||||
"target_type": "ORGANIZATION",
|
||||
"target_description": "Pays qui a obtenu à sa frontière libanaise « un niveau de calme inédit depuis les années 1960 »",
|
||||
"relationship": "Obtention d'un niveau de"""
|
||||
)
|
||||
@@ -0,0 +1,54 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.v1.core.kv_cache_utils import check_enough_kv_cache_memory
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec
|
||||
|
||||
|
||||
def test_kv_cache_oom_no_memory():
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
config = MagicMock()
|
||||
config.model_config.max_model_len = 2048
|
||||
|
||||
spec = {
|
||||
"layer_0": FullAttentionSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=8,
|
||||
head_size=128,
|
||||
dtype="float16",
|
||||
)
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
check_enough_kv_cache_memory(config, spec, 0)
|
||||
|
||||
|
||||
def test_kv_cache_oom_insufficient_memory(monkeypatch):
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
config = MagicMock()
|
||||
config.model_config.max_model_len = 2048
|
||||
config.cache_config.block_size = 16
|
||||
config.parallel_config.tensor_parallel_size = 1
|
||||
config.parallel_config.pipeline_parallel_size = 1
|
||||
config.parallel_config.decode_context_parallel_size = 1
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm.v1.core.kv_cache_utils.max_memory_usage_bytes",
|
||||
lambda c, s: 100 * 1024**3, # 100 GiB
|
||||
)
|
||||
|
||||
spec = {
|
||||
"layer_0": FullAttentionSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=8,
|
||||
head_size=128,
|
||||
dtype="float16",
|
||||
)
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
check_enough_kv_cache_memory(config, spec, 1024**3) # 1 GiB
|
||||
@@ -0,0 +1,238 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import random
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import LLM
|
||||
from vllm.sampling_params import SamplingParams, StructuredOutputsParams
|
||||
from vllm.v1.metrics.reader import Counter, Gauge, Histogram, Metric, Vector
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tests.conftest import VllmRunner
|
||||
else:
|
||||
VllmRunner = object
|
||||
|
||||
MODEL = "facebook/opt-125m"
|
||||
DTYPE = "half"
|
||||
|
||||
|
||||
def _vllm_model(
|
||||
apc: bool,
|
||||
vllm_runner: type[VllmRunner],
|
||||
*,
|
||||
skip_tokenizer_init: bool = False,
|
||||
):
|
||||
"""Set up VllmRunner instance."""
|
||||
return vllm_runner(
|
||||
MODEL,
|
||||
dtype=DTYPE,
|
||||
max_model_len=128,
|
||||
enforce_eager=True,
|
||||
enable_prefix_caching=apc,
|
||||
gpu_memory_utilization=0.5,
|
||||
skip_tokenizer_init=skip_tokenizer_init,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
# Function scope decouples tests & allows
|
||||
# env var adjustment via monkeypatch
|
||||
scope="function",
|
||||
# Prefix caching
|
||||
params=[False, True],
|
||||
)
|
||||
def vllm_model(vllm_runner, request):
|
||||
"""VllmRunner test fixture parameterized by APC True/False."""
|
||||
with _vllm_model(request.param, vllm_runner) as vllm_model:
|
||||
yield vllm_model
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def vllm_model_apc(vllm_runner):
|
||||
"""VllmRunner test fixture with APC."""
|
||||
with _vllm_model(True, vllm_runner) as vllm_model:
|
||||
yield vllm_model
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
# Function scope decouples tests & allows
|
||||
# env var adjustment via monkeypatch
|
||||
scope="function",
|
||||
# Prefix caching
|
||||
params=[False, True],
|
||||
)
|
||||
def vllm_model_skip_tokenizer_init(vllm_runner, request):
|
||||
"""VllmRunner test fixture with APC."""
|
||||
with _vllm_model(
|
||||
request.param,
|
||||
vllm_runner,
|
||||
skip_tokenizer_init=True,
|
||||
) as vllm_model:
|
||||
yield vllm_model
|
||||
|
||||
|
||||
def _get_test_sampling_params(
|
||||
prompt_list: list[str],
|
||||
seed: int | None = 42,
|
||||
structured_outputs: bool = False,
|
||||
) -> tuple[list[SamplingParams], list[int]]:
|
||||
"""Generate random sampling params for a batch."""
|
||||
rng = random.Random(seed)
|
||||
|
||||
def get_mostly_n_gt1() -> int:
|
||||
r"""Mostly n \in [2,20], ~1/3 n=1"""
|
||||
x = rng.randint(0, 28)
|
||||
if x < 10:
|
||||
return 1
|
||||
else:
|
||||
return x - 8
|
||||
|
||||
n_list = [get_mostly_n_gt1() for _ in range(len(prompt_list))]
|
||||
# High temperature to maximize the chance of unique completions
|
||||
return [
|
||||
SamplingParams(
|
||||
temperature=0.95,
|
||||
top_p=0.95,
|
||||
n=n,
|
||||
seed=seed,
|
||||
structured_outputs=StructuredOutputsParams(regex="[0-9]+")
|
||||
if structured_outputs
|
||||
else None,
|
||||
)
|
||||
for n in n_list
|
||||
], n_list
|
||||
|
||||
|
||||
def test_compatibility_with_skip_tokenizer_init(
|
||||
vllm_model_skip_tokenizer_init: VllmRunner,
|
||||
example_prompts: list[str],
|
||||
):
|
||||
# Case 1: Structured output request should raise an error.
|
||||
sampling_params_list, _ = _get_test_sampling_params(
|
||||
example_prompts,
|
||||
structured_outputs=True,
|
||||
)
|
||||
llm: LLM = vllm_model_skip_tokenizer_init.llm
|
||||
with pytest.raises(ValueError):
|
||||
_ = llm.generate(example_prompts, sampling_params_list)
|
||||
|
||||
|
||||
def test_parallel_sampling(vllm_model, example_prompts) -> None:
|
||||
"""Test passes if parallel sampling `n>1` yields `n` unique completions.
|
||||
|
||||
Args:
|
||||
vllm_model: VllmRunner instance under test.
|
||||
example_prompt: test fixture providing prompts for testing.
|
||||
"""
|
||||
sampling_params_list, n_list = _get_test_sampling_params(example_prompts)
|
||||
llm: LLM = vllm_model.llm
|
||||
outputs = llm.generate(example_prompts, sampling_params_list)
|
||||
|
||||
# Validate each request response
|
||||
for out, n in zip(outputs, n_list):
|
||||
completion_counts: dict[str, int] = {}
|
||||
# Assert correct number of completions
|
||||
assert len(out.outputs) == n, f"{len(out.outputs)} completions; {n} expected."
|
||||
for idx in range(n):
|
||||
comp = out.outputs[idx]
|
||||
# Assert correct completion indices
|
||||
assert comp.index == idx, f"Index {comp.index}; expected {idx}."
|
||||
text = comp.text
|
||||
completion_counts[text] = completion_counts.get(text, 0) + 1
|
||||
# Assert unique completions
|
||||
if len(completion_counts) != n:
|
||||
repeats = {txt: num for (txt, num) in completion_counts.items() if num > 1}
|
||||
raise AssertionError(
|
||||
f"{len(completion_counts)} unique completions; expected"
|
||||
f" {n}. Repeats: {repeats}"
|
||||
)
|
||||
|
||||
|
||||
def test_engine_metrics(vllm_runner, example_prompts):
|
||||
max_tokens = 100
|
||||
# Use spec decoding to test num_accepted_tokens_per_pos
|
||||
speculative_config = {
|
||||
"method": "ngram",
|
||||
"prompt_lookup_max": 5,
|
||||
"prompt_lookup_min": 3,
|
||||
"num_speculative_tokens": 5,
|
||||
}
|
||||
|
||||
with vllm_runner(
|
||||
MODEL,
|
||||
speculative_config=speculative_config,
|
||||
disable_log_stats=False,
|
||||
) as vllm_model:
|
||||
llm: LLM = vllm_model.llm
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
|
||||
outputs = llm.generate(example_prompts, sampling_params)
|
||||
|
||||
n_prompts = len(example_prompts)
|
||||
assert len(outputs) == n_prompts
|
||||
|
||||
total_tokens = 0
|
||||
for out in outputs:
|
||||
assert len(out.outputs) == 1
|
||||
total_tokens += len(out.outputs[0].token_ids)
|
||||
assert total_tokens == max_tokens * n_prompts
|
||||
|
||||
metrics = llm.get_metrics()
|
||||
|
||||
def find_metric(name) -> list[Metric]:
|
||||
found = []
|
||||
for metric in metrics:
|
||||
if metric.name == name:
|
||||
found.append(metric)
|
||||
return found
|
||||
|
||||
num_requests_running = find_metric("vllm:num_requests_running")
|
||||
assert len(num_requests_running) == 1
|
||||
assert isinstance(num_requests_running[0], Gauge)
|
||||
assert num_requests_running[0].value == 0.0
|
||||
|
||||
generation_tokens = find_metric("vllm:generation_tokens")
|
||||
assert len(generation_tokens) == 1
|
||||
assert isinstance(generation_tokens[0], Counter)
|
||||
assert generation_tokens[0].value == total_tokens
|
||||
|
||||
request_generation_tokens = find_metric("vllm:request_generation_tokens")
|
||||
assert len(request_generation_tokens) == 1
|
||||
assert isinstance(request_generation_tokens[0], Histogram)
|
||||
assert "+Inf" in request_generation_tokens[0].buckets
|
||||
assert request_generation_tokens[0].buckets["+Inf"] == n_prompts
|
||||
assert request_generation_tokens[0].count == n_prompts
|
||||
assert request_generation_tokens[0].sum == total_tokens
|
||||
|
||||
num_accepted_tokens_per_pos = find_metric(
|
||||
"vllm:spec_decode_num_accepted_tokens_per_pos"
|
||||
)
|
||||
assert len(num_accepted_tokens_per_pos) == 1
|
||||
assert isinstance(num_accepted_tokens_per_pos[0], Vector)
|
||||
assert len(num_accepted_tokens_per_pos[0].values) == 5
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["meta-llama/Llama-3.2-1B-Instruct"])
|
||||
def test_skip_tokenizer_initialization(model: str):
|
||||
# This test checks if the flag skip_tokenizer_init skips the initialization
|
||||
# of tokenizer and detokenizer. The generated output is expected to contain
|
||||
# token ids.
|
||||
llm = LLM(
|
||||
model=model,
|
||||
skip_tokenizer_init=True,
|
||||
enforce_eager=True,
|
||||
)
|
||||
sampling_params = SamplingParams(prompt_logprobs=True, detokenize=True)
|
||||
|
||||
with pytest.raises(ValueError, match="`skip_tokenizer_init=True`"):
|
||||
llm.generate("abc", sampling_params)
|
||||
|
||||
outputs = llm.generate(
|
||||
{"prompt_token_ids": [1, 2, 3]}, sampling_params=sampling_params
|
||||
)
|
||||
assert len(outputs) > 0
|
||||
completions = outputs[0].outputs
|
||||
assert len(completions) > 0
|
||||
assert completions[0].text == ""
|
||||
assert completions[0].token_ids
|
||||
@@ -0,0 +1,66 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for LogprobsProcessor.
|
||||
|
||||
These tests exercise the truncation invariant that the MRV2 sampler relies
|
||||
on: when the sampler returns a row wider than a request's own
|
||||
`num_logprobs + 1` (because another request in the batch needed a wider
|
||||
row), the trailing positions are populated with sentinel values
|
||||
(`token_id=0`, `logprob=-inf`). LogprobsProcessor must read only the first
|
||||
`num_logprobs + 1` entries so those sentinels never reach the user.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
from vllm.logprobs import create_sample_logprobs
|
||||
from vllm.v1.engine.logprobs import LogprobsProcessor
|
||||
from vllm.v1.outputs import LogprobsLists
|
||||
|
||||
|
||||
def _make_processor(num_logprobs: int) -> LogprobsProcessor:
|
||||
return LogprobsProcessor(
|
||||
tokenizer=None,
|
||||
logprobs=create_sample_logprobs(flat_logprobs=False),
|
||||
prompt_logprobs=None,
|
||||
cumulative_logprob=0.0,
|
||||
num_logprobs=num_logprobs,
|
||||
num_prompt_logprobs=None,
|
||||
)
|
||||
|
||||
|
||||
def test_drops_trailing_sentinel_columns():
|
||||
"""A request that asked for 3 custom token logprobs but ended up in a
|
||||
batch padded to width 5 must not surface the trailing -inf entries."""
|
||||
processor = _make_processor(num_logprobs=3)
|
||||
|
||||
sampled = 42
|
||||
# Layout: [sampled, custom_1, custom_2, custom_3, SENTINEL, SENTINEL]
|
||||
# Use float32-exact values so cumulative_logprob compares cleanly.
|
||||
token_ids = np.array([[sampled, 100, 200, 300, 0, 0]], dtype=np.int32)
|
||||
logprobs = np.array([[-0.5, -1.0, -2.0, -3.0, -np.inf, -np.inf]], dtype=np.float32)
|
||||
ranks = np.array([1], dtype=np.int32)
|
||||
|
||||
processor._update_sample_logprobs(LogprobsLists(token_ids, logprobs, ranks))
|
||||
|
||||
assert len(processor.logprobs) == 1
|
||||
pos = processor.logprobs[0]
|
||||
# Exactly sampled + 3 requested tokens; trailing sentinels dropped.
|
||||
assert set(pos.keys()) == {sampled, 100, 200, 300}
|
||||
assert 0 not in pos
|
||||
assert all(np.isfinite(lp.logprob) for lp in pos.values())
|
||||
# cumulative_logprob comes from the sampled token's logprob only.
|
||||
assert processor.cumulative_logprob == -0.5
|
||||
|
||||
|
||||
def test_accepts_exactly_sized_row():
|
||||
"""When the row is exactly num_logprobs+1, no truncation needed."""
|
||||
processor = _make_processor(num_logprobs=2)
|
||||
|
||||
token_ids = np.array([[7, 11, 13]], dtype=np.int32)
|
||||
logprobs = np.array([[-0.5, -1.5, -2.5]], dtype=np.float32)
|
||||
ranks = np.array([1], dtype=np.int32)
|
||||
|
||||
processor._update_sample_logprobs(LogprobsLists(token_ids, logprobs, ranks))
|
||||
|
||||
pos = processor.logprobs[0]
|
||||
assert set(pos.keys()) == {7, 11, 13}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,83 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from vllm import SamplingParams
|
||||
from vllm.outputs import CompletionOutput
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.engine.parallel_sampling import ParentRequest
|
||||
|
||||
|
||||
def test_parent_request_to_output_stream() -> None:
|
||||
parent_request = ParentRequest(make_request(SamplingParams(n=2)))
|
||||
parent_request.child_requests = {"child_id_0", "child_id_1"}
|
||||
output_0 = CompletionOutput(
|
||||
index=0, text="child 0", token_ids=[], cumulative_logprob=None, logprobs=None
|
||||
)
|
||||
output_1 = CompletionOutput(
|
||||
index=1, text="child 1", token_ids=[], cumulative_logprob=None, logprobs=None
|
||||
)
|
||||
# Request not finished
|
||||
assert ([output_0], False) == parent_request.get_outputs("child_id_0", output_0)
|
||||
assert ([output_1], False) == parent_request.get_outputs("child_id_1", output_1)
|
||||
assert ([output_0], False) == parent_request.get_outputs("child_id_0", output_0)
|
||||
assert ([output_1], False) == parent_request.get_outputs("child_id_1", output_1)
|
||||
|
||||
# output_1 finished
|
||||
output_1.finish_reason = "ended"
|
||||
assert ([output_0], False) == parent_request.get_outputs("child_id_0", output_0)
|
||||
assert ([output_1], False) == parent_request.get_outputs("child_id_1", output_1)
|
||||
# Finished output_1 had already returned, DO NOT returned again
|
||||
assert ([output_0], False) == parent_request.get_outputs("child_id_0", output_0)
|
||||
assert parent_request.get_outputs("child_id_1", output_1) == ([], False)
|
||||
|
||||
# output_0 finished
|
||||
output_0.finish_reason = "ended"
|
||||
assert ([output_0], True) == parent_request.get_outputs("child_id_0", output_0)
|
||||
assert parent_request.get_outputs("child_id_1", output_1) == ([], True)
|
||||
# Finished output_0 had already returned, DO NOT returned again
|
||||
assert parent_request.get_outputs("child_id_0", output_0) == ([], True)
|
||||
assert parent_request.get_outputs("child_id_1", output_1) == ([], True)
|
||||
|
||||
|
||||
def test_parent_request_to_output_final_only() -> None:
|
||||
parent_request = ParentRequest(
|
||||
make_request(SamplingParams(n=2, output_kind=RequestOutputKind.FINAL_ONLY))
|
||||
)
|
||||
parent_request.child_requests = {"child_id_0", "child_id_1"}
|
||||
output_0 = CompletionOutput(
|
||||
index=0, text="child 0", token_ids=[], cumulative_logprob=None, logprobs=None
|
||||
)
|
||||
output_1 = CompletionOutput(
|
||||
index=1, text="child 1", token_ids=[], cumulative_logprob=None, logprobs=None
|
||||
)
|
||||
# Request not finished, return nothing
|
||||
assert parent_request.get_outputs("child_id_0", output_0) == ([], False)
|
||||
assert parent_request.get_outputs("child_id_1", output_1) == ([], False)
|
||||
# output_1 finished, but outputs won't be returned until all child requests finished
|
||||
output_1.finish_reason = "ended"
|
||||
assert parent_request.get_outputs("child_id_0", output_0) == ([], False)
|
||||
assert parent_request.get_outputs("child_id_1", output_1) == ([], False)
|
||||
# output_0 finished, as all child requests finished, the output would be returned
|
||||
output_0.finish_reason = "ended"
|
||||
assert ([output_0, output_1], True) == parent_request.get_outputs(
|
||||
"child_id_0", output_0
|
||||
)
|
||||
assert ([output_0, output_1], True) == parent_request.get_outputs(
|
||||
"child_id_1", output_1
|
||||
)
|
||||
|
||||
|
||||
def make_request(sampling_params: SamplingParams) -> EngineCoreRequest:
|
||||
return EngineCoreRequest(
|
||||
request_id="parent_id",
|
||||
external_req_id="ext_parent_id",
|
||||
prompt_token_ids=None,
|
||||
mm_features=None,
|
||||
sampling_params=sampling_params,
|
||||
pooling_params=None,
|
||||
arrival_time=0.0,
|
||||
lora_request=None,
|
||||
cache_salt=None,
|
||||
data_parallel_rank=None,
|
||||
)
|
||||
@@ -0,0 +1,63 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch.cuda
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.engine.core import EngineCore
|
||||
|
||||
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
|
||||
|
||||
|
||||
def test_preprocess_error_handling(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that preprocessing errors are handled gracefully."""
|
||||
|
||||
if current_platform.is_rocm() or current_platform.is_xpu():
|
||||
pytest.skip(
|
||||
"Skipped on ROCm/XPU: this test only works with 'fork', "
|
||||
"but ROCm/XPU uses 'spawn'."
|
||||
)
|
||||
|
||||
assert not torch.cuda.is_initialized(), (
|
||||
"fork needs to be used for the engine "
|
||||
"core process and this isn't possible if cuda is already initialized"
|
||||
)
|
||||
|
||||
# Store original method to call for non-failing requests
|
||||
original_preprocess = EngineCore.preprocess_add_request
|
||||
|
||||
# Monkeypatch to make preprocess_add_request raise an exception
|
||||
# only for requests with "FAIL" in the first token
|
||||
def conditional_failing_preprocess(self, request: EngineCoreRequest):
|
||||
# Fail if the first token id is 333
|
||||
if request.prompt_token_ids and request.prompt_token_ids[0] == 333:
|
||||
raise ValueError("Simulated preprocessing error!")
|
||||
return original_preprocess(self, request)
|
||||
|
||||
monkeypatch.setattr(
|
||||
EngineCore, "preprocess_add_request", conditional_failing_preprocess
|
||||
)
|
||||
|
||||
llm = LLM(model=MODEL_NAME)
|
||||
|
||||
# Create a failing request by crafting a request with an invalid token
|
||||
# We need to use a direct approach since LLM.generate tokenizes for us
|
||||
from vllm.inputs import TokensPrompt
|
||||
|
||||
# This should raise an exception due to the preprocessing failure
|
||||
# Special token id to trigger the failure
|
||||
failing_prompt = TokensPrompt(prompt_token_ids=[333])
|
||||
outputs = llm.generate(failing_prompt, SamplingParams(max_tokens=10)) # type: ignore
|
||||
assert len(outputs) == 1
|
||||
assert len(outputs[0].outputs[0].token_ids) == 0
|
||||
assert outputs[0].finished
|
||||
assert outputs[0].outputs[0].finish_reason == "error"
|
||||
|
||||
# Verify the engine is still functional with a normal request
|
||||
outputs = llm.generate("Hello, my name is", SamplingParams(max_tokens=10))
|
||||
assert len(outputs) == 1
|
||||
assert len(outputs[0].outputs[0].token_ids) > 0
|
||||
assert outputs[0].outputs[0].finish_reason in ("stop", "length")
|
||||
@@ -0,0 +1,432 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from typing import TypeAlias
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from transformers import PythonBackend, TokenizersBackend
|
||||
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.v1.engine import EngineCoreOutput, FinishReason
|
||||
from vllm.v1.metrics.stats import PrefillStats
|
||||
from vllm.v1.outputs import LogprobsLists, LogprobsTensors
|
||||
|
||||
GeneralTokenizerType: TypeAlias = PythonBackend | TokenizersBackend
|
||||
|
||||
# Number of sample logprobs to request when testing sample logprobs
|
||||
NUM_SAMPLE_LOGPROBS_UNDER_TEST = 5
|
||||
# Number of prompt logprobs to request when testing prompt logprobs
|
||||
NUM_PROMPT_LOGPROBS_UNDER_TEST = 7
|
||||
|
||||
TOKENIZER_NAME = "meta-llama/Llama-3.2-1B"
|
||||
|
||||
FULL_STRINGS = [
|
||||
"My name is Robert from Neural Magic and I love working on vLLM so much!",
|
||||
"Red Hat is the best open source company by far across Linux, K8s, and AI.",
|
||||
"Nick is the name of my brother in addition to my colleague from Red Hat.",
|
||||
]
|
||||
STOP_STRINGS = ["I love working on", "company by far", "brother in"]
|
||||
PROMPT_LEN = 5
|
||||
|
||||
random.seed(42)
|
||||
|
||||
|
||||
def _create_random_top_logprob_test_vector(
|
||||
num_logprobs: int,
|
||||
lower: float,
|
||||
upper: float,
|
||||
) -> torch.Tensor:
|
||||
"""Create a random vector of top logprob float values.
|
||||
|
||||
Use to create fake sample logprobs for testing.
|
||||
|
||||
Note that a real production scenario would require
|
||||
logprobs to be sorted in descending order, something
|
||||
which is omitted in this function.
|
||||
|
||||
Args:
|
||||
num_logprobs: number of top logprobs
|
||||
lower: lower range of logprob float values
|
||||
upper: upper range of logprob float values
|
||||
|
||||
Returns:
|
||||
1D length-`num_logprobs` torch Tensor of float logprob values
|
||||
"""
|
||||
return torch.rand(num_logprobs) * (upper - lower) + lower
|
||||
|
||||
|
||||
def _create_random_top_logprob_test_matrix(
|
||||
shape: tuple,
|
||||
lower: float,
|
||||
upper: float,
|
||||
) -> torch.Tensor:
|
||||
"""Create a random matrix of top logprob float values.
|
||||
|
||||
Use to create fake prompt logprobs for testing.
|
||||
|
||||
Note that a real production scenario would require
|
||||
logprobs to be sorted in descending order along rows,
|
||||
something which is omitted in this function.
|
||||
|
||||
Args:
|
||||
shape: (num_tokens,num_logprobs) tuple representing
|
||||
matrix shape
|
||||
lower: lower range of logprob float values
|
||||
upper: upper range of logprob float values
|
||||
|
||||
Returns:
|
||||
2D num_tokens x num_logprobs torch Tensor of float logprob values
|
||||
"""
|
||||
return torch.rand(*shape) * (upper - lower) + lower
|
||||
|
||||
|
||||
def _create_random_top_token_test_vector(
|
||||
num_logprobs: int,
|
||||
lower: int,
|
||||
upper: int,
|
||||
sampled_token_id: int,
|
||||
adjust_num_logprobs: bool = True,
|
||||
) -> tuple[torch.Tensor, int]:
|
||||
"""Create a random vector of top logprob token indices
|
||||
|
||||
Use to create fake sample logprobs for testing. The sampled token
|
||||
ID must always be one of the top logprobs, which this dummy test
|
||||
vector generator enforces. OpenAI API
|
||||
compatible engines must be able to return an additional sample
|
||||
logprob for the sampled token if the sampled token was not
|
||||
among the top sample logprobs; `adjust_num_logprobs` emulates
|
||||
this behavior by increasing the vector length by 1 if
|
||||
`adjust_num_logprobs` is set.
|
||||
|
||||
Args:
|
||||
num_logprobs: number of top logprobs
|
||||
lower: lower range of token ids
|
||||
upper: upper range of token ids
|
||||
sampled_token_id: the token actually sampled
|
||||
adjust_num_logprobs: if True, emulate situation where sampled
|
||||
token logprob must be injected into top
|
||||
logprobs
|
||||
|
||||
Returns:
|
||||
1D length-x torch Tensor of token ids where x is
|
||||
`num_logprobs+1` if `adjust_num_logprobs` and
|
||||
`num_logprobs` otherwise
|
||||
sampled_token_rank: the rank of sampled_token_id in the vocab
|
||||
vector when sorted in descending order by
|
||||
logprob
|
||||
"""
|
||||
|
||||
# Calculate the final number of logprobs required
|
||||
total_logprobs = num_logprobs + 1 if adjust_num_logprobs else num_logprobs
|
||||
|
||||
# Generate random indices using torch
|
||||
choice_tensor = torch.randperm(upper - lower)[:total_logprobs] + lower
|
||||
|
||||
# Ensure the sampled token ID is included in the tensor
|
||||
choice_tensor[0] = sampled_token_id
|
||||
|
||||
# Check if the sampled_token_id occurs in choice_tensor[1:]
|
||||
if sampled_token_id in choice_tensor[1:]:
|
||||
sampled_token_rank = (
|
||||
(choice_tensor[1:] == sampled_token_id).nonzero(as_tuple=True)[0].item()
|
||||
)
|
||||
else:
|
||||
# If not found, assign a random int between num_logprobs and 50700
|
||||
sampled_token_rank = random.randint(num_logprobs, 50700)
|
||||
|
||||
return choice_tensor, sampled_token_rank
|
||||
|
||||
|
||||
def _create_random_top_token_test_matrix(
|
||||
shape: tuple[int, int],
|
||||
lower: int,
|
||||
upper: int,
|
||||
tokens_list: list[int],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Create a random matrix of top logprob token indices
|
||||
|
||||
Use to create fake prompt logprobs for testing.
|
||||
|
||||
Token ids are generated randomly and sampled without
|
||||
replacement.
|
||||
|
||||
Args:
|
||||
shape: (num_tokens, num_logprobs) tuple representing
|
||||
matrix shape
|
||||
lower: lower range of token ids
|
||||
upper: upper range of token ids
|
||||
|
||||
Returns:
|
||||
tuple containing:
|
||||
- 2D num_tokens x num_logprobs+1 torch Tensor of token ids
|
||||
- 1D tensor of ranks of prompt tokens in their respective
|
||||
rows, or random values
|
||||
"""
|
||||
num_elements = shape[0] * shape[1]
|
||||
choice_tensor = torch.randperm(upper - lower)[:num_elements] + lower
|
||||
matrix = torch.cat(
|
||||
(
|
||||
torch.tensor(tokens_list, dtype=torch.int).unsqueeze(-1),
|
||||
choice_tensor.view(shape),
|
||||
),
|
||||
dim=1,
|
||||
)
|
||||
|
||||
# Initialize the tensor for storing the ranks
|
||||
prompt_token_ranks = torch.empty(shape[0], dtype=torch.int)
|
||||
|
||||
# Iterate over each row to check presence of
|
||||
# tokens_list[rdx] and determine its index
|
||||
for rdx in range(shape[0]):
|
||||
row = matrix[rdx, 1:] # Skip the first column as it contains the token list
|
||||
token_index = (row == tokens_list[rdx]).nonzero(as_tuple=True)[0]
|
||||
if token_index.numel() > 0:
|
||||
prompt_token_ranks[rdx] = token_index.item()
|
||||
else:
|
||||
prompt_token_ranks[rdx] = random.randint(shape[1], 50700)
|
||||
|
||||
return matrix, prompt_token_ranks
|
||||
|
||||
|
||||
def decode_token(
|
||||
tok_id: int,
|
||||
tokenizer: PythonBackend,
|
||||
) -> str:
|
||||
"""Reproduce the process of detokenizing a token for testing purposes.
|
||||
|
||||
Args:
|
||||
tok_id: token id to detokenize
|
||||
tokenizer: tokenizer to use for detokenization
|
||||
|
||||
Returns:
|
||||
string representation of token
|
||||
"""
|
||||
return tokenizer.convert_ids_to_tokens(tok_id)
|
||||
|
||||
|
||||
def generate_dummy_sample_logprobs(
|
||||
sampled_tokens_list: list,
|
||||
num_logprobs: int,
|
||||
tokenizer: PythonBackend,
|
||||
) -> list[tuple[list[int], list[float], int]]:
|
||||
"""Generate dummy sample logprobs
|
||||
|
||||
Generate a test data structure which imitates the list of sample logprobs
|
||||
which would be assembled in the engine core during decode phase.
|
||||
|
||||
Args:
|
||||
sampled_tokens_list: list of sampled tokens
|
||||
num_logprobs: return `num_logprobs` or `num_logprobs+1` logprobs per token
|
||||
tokenizer: model tokenizer to use for detokenization
|
||||
|
||||
Returns
|
||||
list of (top token ids vector, logprobs vector, sampled token rank)
|
||||
Python lists tuples; in each tuple the logprobs and top token ids
|
||||
vectors have the same length which is either `num_logprobs` or
|
||||
`num_logprobs+1`. Sampled token rank is the rank (index+1) of the
|
||||
sampled token within the vocab vector when sorted by logprob in
|
||||
descending order.
|
||||
"""
|
||||
res = []
|
||||
for sampled_token_id in sampled_tokens_list:
|
||||
(
|
||||
token_vector,
|
||||
sampled_token_rank,
|
||||
) = _create_random_top_token_test_vector(
|
||||
num_logprobs, 0, len(tokenizer.vocab) - 1, sampled_token_id
|
||||
)
|
||||
|
||||
res.append(
|
||||
(
|
||||
token_vector,
|
||||
_create_random_top_logprob_test_vector(num_logprobs + 1, -100, 0),
|
||||
sampled_token_rank,
|
||||
)
|
||||
)
|
||||
|
||||
# Convert tensors in the list tuples to Python lists
|
||||
res_list_format = [
|
||||
(log_probs_tensor.tolist(), token_ids_tensor.tolist(), sampled_token_rank)
|
||||
for log_probs_tensor, token_ids_tensor, sampled_token_rank in res
|
||||
]
|
||||
|
||||
return res_list_format
|
||||
|
||||
|
||||
def generate_dummy_prompt_logprobs_tensors(
|
||||
prompt_tokens_list: list,
|
||||
num_logprobs: int,
|
||||
tokenizer: PythonBackend,
|
||||
) -> LogprobsTensors:
|
||||
"""Generate dummy prompt logprobs tensors
|
||||
|
||||
Generate a test data structure which imitates the torch Tensors of prompt
|
||||
logprobs which would be assembled in the engine core during chunked
|
||||
prefill.
|
||||
|
||||
Args:
|
||||
prompt_tokens_list: list of prompt tokens
|
||||
num_logprobs: return `num_logprobs` logprobs per token
|
||||
tokenizer: model tokenizer to use for detokenization
|
||||
|
||||
Returns
|
||||
Single tuple of (logprobs matrix, top token ids matrix) torch Tensor,
|
||||
where both matrices have dimensions
|
||||
num_prompt_tokens x num_logprobs
|
||||
"""
|
||||
# For now, assume the whole prompt is processed in one chunk; thus,
|
||||
# the number of non-`None` prompt logprobs is `len(prompt_tokens_list)-1`.
|
||||
# Prior to injecting `None` at the beginning of prompt logprobs (which
|
||||
# happens later in the detokenizer, not here), the prompt logprobs in
|
||||
# the ith position are predicting the probability distribution of the
|
||||
# prompt token in (i+1)st position. Thus, we concat
|
||||
# `prompt_tokens_list[1:]` to the dummy token ids, just as the engine
|
||||
# would.
|
||||
num_prompt_logprobs = len(prompt_tokens_list) - 1
|
||||
(
|
||||
token_vector,
|
||||
prompt_token_ranks,
|
||||
) = _create_random_top_token_test_matrix(
|
||||
(num_prompt_logprobs, num_logprobs),
|
||||
0,
|
||||
len(tokenizer.vocab) - 1,
|
||||
prompt_tokens_list[1:],
|
||||
)
|
||||
return LogprobsTensors(
|
||||
token_vector,
|
||||
_create_random_top_logprob_test_matrix(
|
||||
(num_prompt_logprobs, num_logprobs + 1), -100, 0
|
||||
),
|
||||
prompt_token_ranks,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DummyOutputProcessorTestVectors:
|
||||
"""Dummy test vectors for output processor tests"""
|
||||
|
||||
tokenizer: GeneralTokenizerType
|
||||
vllm_config: EngineArgs
|
||||
full_tokens: list[list[int]] # Prompt + generated tokens
|
||||
prompt_tokens: list[list[int]]
|
||||
generation_tokens: list[list[int]]
|
||||
# Each request is associated with a tuple of
|
||||
# (top tokens, top logprobs, ranks) prompt logprobs tensors
|
||||
prompt_logprobs: list[LogprobsTensors]
|
||||
# Each request is associated with a sample logprobs; a request's
|
||||
# sample logprobs are a list of (top tokens, top logprobs, ranks)
|
||||
# sample logprobs tensors at each sequence position
|
||||
generation_logprobs: list[list[tuple[list[int], list[float], int]]]
|
||||
prompt_strings: list[str]
|
||||
prompt_strings_len: list[int]
|
||||
generation_strings: list[str]
|
||||
|
||||
|
||||
class MockEngineCore:
|
||||
"""Mock engine core outputs form premade tokens lists."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokens_list: list[list[int]],
|
||||
prompts_list: list[list[int]],
|
||||
# For each request, for each sampled token offset,
|
||||
# a tuple of
|
||||
# (list of topk token ids, list of sample logprob vals, rank)
|
||||
generated_logprobs_raw: list[list[tuple[list[int], list[float], int]]]
|
||||
| None = None,
|
||||
# For each request, a tuple of
|
||||
# (prompt logprob val matrix, prompt logprob tok id matrix);
|
||||
# each matrix has dimensions
|
||||
# (num prompt toks) x (num prompt logprobs+1)
|
||||
prompt_logprobs_raw: list[LogprobsTensors] | None = None,
|
||||
eos_token_id: int | None = None,
|
||||
stop_token_ids: list[int] | None = None,
|
||||
request_ids: list[str] | None = None,
|
||||
) -> None:
|
||||
self.num_requests = len(tokens_list)
|
||||
self.tokens_list = tokens_list
|
||||
self.prompts_list = prompts_list
|
||||
self.generated_logprobs_raw = generated_logprobs_raw
|
||||
self.do_logprobs = generated_logprobs_raw is not None
|
||||
self.prompt_logprobs_raw = prompt_logprobs_raw
|
||||
self.do_prompt_logprobs = prompt_logprobs_raw is not None
|
||||
self.request_finished = [False for _ in range(self.num_requests)]
|
||||
self.request_token_idx = [0 for _ in range(self.num_requests)]
|
||||
self.eos_token_id = eos_token_id
|
||||
self.stop_token_ids = stop_token_ids
|
||||
self.request_ids = (
|
||||
request_ids
|
||||
if request_ids is not None
|
||||
else [f"request-{i}" for i in range(self.num_requests)]
|
||||
)
|
||||
|
||||
def get_outputs(self, num_active: int = -1) -> list[EngineCoreOutput]:
|
||||
do_logprobs = self.do_logprobs
|
||||
do_prompt_logprobs = self.do_prompt_logprobs
|
||||
|
||||
outputs = []
|
||||
for req_idx, (token_ids, prompt_token_ids) in enumerate(
|
||||
zip(self.tokens_list, self.prompts_list)
|
||||
):
|
||||
if num_active != -1 and req_idx >= num_active:
|
||||
break
|
||||
if not self.request_finished[req_idx]:
|
||||
token_idx = self.request_token_idx[req_idx]
|
||||
if do_logprobs:
|
||||
assert self.generated_logprobs_raw is not None
|
||||
(logprobs_token_ids_, logprobs_, sampled_token_ranks_) = (
|
||||
self.generated_logprobs_raw[req_idx][token_idx]
|
||||
)
|
||||
logprobs = LogprobsLists(
|
||||
np.array([logprobs_token_ids_]),
|
||||
np.array([logprobs_]),
|
||||
np.array([sampled_token_ranks_]),
|
||||
)
|
||||
else:
|
||||
logprobs = None
|
||||
if do_prompt_logprobs:
|
||||
if token_idx == 0:
|
||||
assert self.prompt_logprobs_raw is not None
|
||||
prompt_logprobs = self.prompt_logprobs_raw[req_idx]
|
||||
else:
|
||||
prompt_logprobs = None
|
||||
else:
|
||||
prompt_logprobs = None
|
||||
|
||||
# Add prefill_stats on first output (prefill) for this request
|
||||
if token_idx == 0:
|
||||
prefill_stats = PrefillStats()
|
||||
prefill_stats.set(
|
||||
num_prompt_tokens=len(prompt_token_ids),
|
||||
num_local_cached_tokens=0,
|
||||
num_external_cached_tokens=0,
|
||||
)
|
||||
else:
|
||||
prefill_stats = None
|
||||
|
||||
new_token_id = token_ids[token_idx]
|
||||
output = EngineCoreOutput(
|
||||
request_id=self.request_ids[req_idx],
|
||||
new_token_ids=[new_token_id],
|
||||
new_logprobs=logprobs,
|
||||
new_prompt_logprobs_tensors=prompt_logprobs,
|
||||
prefill_stats=prefill_stats,
|
||||
)
|
||||
if token_idx == len(token_ids) - 1:
|
||||
output.finish_reason = FinishReason.LENGTH
|
||||
self.request_finished[req_idx] = True
|
||||
if new_token_id == self.eos_token_id:
|
||||
output.finish_reason = FinishReason.STOP
|
||||
self.request_finished[req_idx] = True
|
||||
if new_token_id in (self.stop_token_ids or ()):
|
||||
output.finish_reason = FinishReason.STOP
|
||||
output.stop_reason = new_token_id
|
||||
self.request_finished[req_idx] = True
|
||||
outputs.append(output)
|
||||
|
||||
self.request_token_idx[req_idx] += 1
|
||||
|
||||
return outputs
|
||||
Reference in New Issue
Block a user