biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[ROCm][CI] Fix engine core client tests for ROCm spawn multiprocessing (#32061)

Browse Source 
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
Andreas Karatzas
2026-01-13 01:14:30 -06:00
committed by GitHub
parent 44c34f22d9
commit df7e12715f
1 changed files with 181 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files

251
tests/v1/engine/test_engine_core_client.py
View File

@@ -3,6 +3,7 @@
import asyncio
import importlib
import inspect
import os
import signal
import time
@@ -47,6 +48,7 @@ MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
TOKENIZER = AutoTokenizer.from_pretrained(MODEL_NAME)
PROMPT = "Hello my name is Robert and I love quantization kernels"
PROMPT_TOKENS = TOKENIZER(PROMPT).input_ids
TEST_MODULE = "tests.v1.engine.test_engine_core_client"
_REQUEST_COUNTER = 0
@@ -223,10 +225,170 @@ def echo(self, msg: str, err_msg: str | None = None, sleep: float | None = None)
return msg
@dataclass
class TestMessage:
"""Test dataclass for verifying custom type serialization."""
message: str
# Dummy utility function to monkey-patch into engine core.
def echo_dc(
self,
msg: str,
return_list: bool = False,
) -> TestMessage | list[TestMessage]:
print(f"echo dc util function called: {msg}")
val = None if msg is None else TestMessage(msg)
# Return dataclass to verify support for returning custom types
# (for which there is special handling to make it work with msgspec).
return [val for _ in range(3)] if return_list else val
# Dummy utility function to test dict serialization with custom types.
def echo_dc_dict(
self,
msg: str,
return_dict: bool = False,
) -> TestMessage | dict[str, TestMessage]:
print(f"echo dc dict util function called: {msg}")
val = None if msg is None else TestMessage(msg)
# Return dict of dataclasses to verify support for returning dicts
# with custom value types.
if return_dict:
return {"key1": val, "key2": val, "key3": val}
else:
return val
# Dummy utility function to test nested structures with custom types.
def echo_dc_nested(
self,
msg: str,
structure_type: str = "list_of_dicts",
) -> Any:
print(f"echo dc nested util function called: {msg}, structure: {structure_type}")
val = None if msg is None else TestMessage(msg)
structures = {
"list_of_dicts": [{"a": val, "b": val}, {"c": val, "d": val}],
"dict_of_lists": {"list1": [val, val], "list2": [val, val]},
"deep_nested": {"outer": [{"inner": [val, val]}, {"inner": [val]}]},
}
return structures.get(structure_type, val)
# --- Fixtures for subprocess patching ---
# These create sitecustomize.py files that patch EngineCore in spawned
# subprocesses. This is necessary because ROCm requires 'spawn' multiprocessing
# start method, which creates fresh Python interpreters that don't inherit
# monkey-patches from the parent process.
@pytest.fixture
def subprocess_echo_patch(monkeypatch, tmp_path):
"""Create sitecustomize.py so spawned subprocesses have echo method.
This is needed because ROCm uses 'spawn' multiprocessing start method,
which creates a fresh Python interpreter that doesn't inherit monkey-patches.
By using sitecustomize.py, we ensure the patch is applied when Python starts.
"""
sc = tmp_path / "sitecustomize.py"
sc.write_text(
"\n".join(
[
"import time",
"from vllm.v1.engine.core import EngineCore",
inspect.getsource(echo),
"EngineCore.echo = echo",
]
)
)
monkeypatch.setenv(
"PYTHONPATH",
os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),
)
@pytest.fixture
def subprocess_echo_dc_patch(monkeypatch, tmp_path):
"""Create sitecustomize.py so spawned subprocesses have echo_dc method."""
sc = tmp_path / "sitecustomize.py"
sc.write_text(
"\n".join(
[
"from dataclasses import dataclass",
"",
inspect.getsource(TestMessage),
f"TestMessage.__module__ = '{TEST_MODULE}'",
"",
"from vllm.v1.engine.core import EngineCore",
inspect.getsource(echo_dc),
"EngineCore.echo_dc = echo_dc",
]
)
)
monkeypatch.setenv(
"PYTHONPATH",
os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),
)
@pytest.fixture
def subprocess_echo_dc_dict_patch(monkeypatch, tmp_path):
"""Create sitecustomize.py so spawned subprocesses have echo_dc_dict method."""
sc = tmp_path / "sitecustomize.py"
sc.write_text(
"\n".join(
[
"from dataclasses import dataclass",
"",
inspect.getsource(TestMessage),
f"TestMessage.__module__ = '{TEST_MODULE}'",
"",
"from vllm.v1.engine.core import EngineCore",
inspect.getsource(echo_dc_dict),
"EngineCore.echo_dc_dict = echo_dc_dict",
]
)
)
monkeypatch.setenv(
"PYTHONPATH",
os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),
)
@pytest.fixture
def subprocess_echo_dc_nested_patch(monkeypatch, tmp_path):
"""Create sitecustomize.py so spawned subprocesses have echo_dc_nested method."""
sc = tmp_path / "sitecustomize.py"
sc.write_text(
"\n".join(
[
"from dataclasses import dataclass",
"from typing import Any",
"",
inspect.getsource(TestMessage),
f"TestMessage.__module__ = '{TEST_MODULE}'",
"",
"from vllm.v1.engine.core import EngineCore",
inspect.getsource(echo_dc_nested),
"EngineCore.echo_dc_nested = echo_dc_nested",
]
)
)
monkeypatch.setenv(
"PYTHONPATH",
os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),
)
@create_new_process_for_each_test()
@pytest.mark.parametrize("multiprocessing_mode", [True, False])
def test_engine_core_client(
monkeypatch: pytest.MonkeyPatch, multiprocessing_mode: bool
monkeypatch: pytest.MonkeyPatch,
multiprocessing_mode: bool,
subprocess_echo_patch,
):
with monkeypatch.context() as m:
# Monkey-patch core engine utility function to test.
@@ -313,7 +475,10 @@ def test_engine_core_client(
@pytest.mark.asyncio(loop_scope="function")
async def test_engine_core_client_asyncio(monkeypatch: pytest.MonkeyPatch):
async def test_engine_core_client_asyncio(
monkeypatch: pytest.MonkeyPatch,
subprocess_echo_patch,
):
with monkeypatch.context() as m:
# Monkey-patch core engine utility function to test.
m.setattr(EngineCore, "echo", echo, raising=False)
@@ -406,66 +571,10 @@ async def test_engine_core_client_asyncio(monkeypatch: pytest.MonkeyPatch):
client.shutdown()
@dataclass
class MyDataclass:
message: str
# Dummy utility function to monkey-patch into engine core.
def echo_dc(
self,
msg: str,
return_list: bool = False,
) -> MyDataclass | list[MyDataclass]:
print(f"echo dc util function called: {msg}")
val = None if msg is None else MyDataclass(msg)
# Return dataclass to verify support for returning custom types
# (for which there is special handling to make it work with msgspec).
return [val for _ in range(3)] if return_list else val
# Dummy utility function to test dict serialization with custom types.
def echo_dc_dict(
self,
msg: str,
return_dict: bool = False,
) -> MyDataclass | dict[str, MyDataclass]:
print(f"echo dc dict util function called: {msg}")
val = None if msg is None else MyDataclass(msg)
# Return dict of dataclasses to verify support for returning dicts
# with custom value types.
if return_dict:
return {"key1": val, "key2": val, "key3": val}
else:
return val
# Dummy utility function to test nested structures with custom types.
def echo_dc_nested(
self,
msg: str,
structure_type: str = "list_of_dicts",
) -> Any:
print(f"echo dc nested util function called: {msg}, structure: {structure_type}")
val = None if msg is None else MyDataclass(msg)
if structure_type == "list_of_dicts": # noqa
# Return list of dicts: [{"a": val, "b": val}, {"c": val, "d": val}]
return [{"a": val, "b": val}, {"c": val, "d": val}]
elif structure_type == "dict_of_lists":
# Return dict of lists: {"list1": [val, val], "list2": [val, val]}
return {"list1": [val, val], "list2": [val, val]}
elif structure_type == "deep_nested":
# Return deeply nested: {"outer": [{"inner": [val, val]},
# {"inner": [val]}]}
return {"outer": [{"inner": [val, val]}, {"inner": [val]}]}
else:
return val
@pytest.mark.asyncio(loop_scope="function")
async def test_engine_core_client_util_method_custom_return(
monkeypatch: pytest.MonkeyPatch,
subprocess_echo_dc_patch,
):
with monkeypatch.context() as m:
# Must set insecure serialization to allow returning custom types.
@@ -494,10 +603,10 @@ async def test_engine_core_client_util_method_custom_return(
core_client: AsyncMPClient = client
result = await core_client.call_utility_async("echo_dc", "testarg2", False)
assert isinstance(result, MyDataclass) and result.message == "testarg2"
assert isinstance(result, TestMessage) and result.message == "testarg2"
result = await core_client.call_utility_async("echo_dc", "testarg2", True)
assert isinstance(result, list) and all(
isinstance(r, MyDataclass) and r.message == "testarg2" for r in result
isinstance(r, TestMessage) and r.message == "testarg2" for r in result
)
# Test returning None and list of Nones
@@ -513,6 +622,7 @@ async def test_engine_core_client_util_method_custom_return(
@pytest.mark.asyncio(loop_scope="function")
async def test_engine_core_client_util_method_custom_dict_return(
monkeypatch: pytest.MonkeyPatch,
subprocess_echo_dc_dict_patch,
):
with monkeypatch.context() as m:
# Must set insecure serialization to allow returning custom types.
@@ -544,7 +654,7 @@ async def test_engine_core_client_util_method_custom_dict_return(
result = await core_client.call_utility_async(
"echo_dc_dict", "testarg3", False
)
assert isinstance(result, MyDataclass) and result.message == "testarg3"
assert isinstance(result, TestMessage) and result.message == "testarg3"
# Test dict return with custom value types
result = await core_client.call_utility_async(
@@ -553,7 +663,7 @@ async def test_engine_core_client_util_method_custom_dict_return(
assert isinstance(result, dict) and len(result) == 3
for key, val in result.items():
assert key in ["key1", "key2", "key3"]
assert isinstance(val, MyDataclass) and val.message == "testarg3"
assert isinstance(val, TestMessage) and val.message == "testarg3"
# Test returning dict with None values
result = await core_client.call_utility_async("echo_dc_dict", None, True)
@@ -569,6 +679,7 @@ async def test_engine_core_client_util_method_custom_dict_return(
@pytest.mark.asyncio(loop_scope="function")
async def test_engine_core_client_util_method_nested_structures(
monkeypatch: pytest.MonkeyPatch,
subprocess_echo_dc_nested_patch,
):
with monkeypatch.context() as m:
# Must set insecure serialization to allow returning custom types.
@@ -605,21 +716,21 @@ async def test_engine_core_client_util_method_nested_structures(
if i == 0:
assert "a" in item and "b" in item
assert (
isinstance(item["a"], MyDataclass)
isinstance(item["a"], TestMessage)
and item["a"].message == "nested1"
)
assert (
isinstance(item["b"], MyDataclass)
isinstance(item["b"], TestMessage)
and item["b"].message == "nested1"
)
else:
assert "c" in item and "d" in item
assert (
isinstance(item["c"], MyDataclass)
isinstance(item["c"], TestMessage)
and item["c"].message == "nested1"
)
assert (
isinstance(item["d"], MyDataclass)
isinstance(item["d"], TestMessage)
and item["d"].message == "nested1"
)
@@ -632,7 +743,7 @@ async def test_engine_core_client_util_method_nested_structures(
for key, lst in result.items():
assert isinstance(lst, list) and len(lst) == 2
for item in lst:
assert isinstance(item, MyDataclass) and item.message == "nested2"
assert isinstance(item, TestMessage) and item.message == "nested2"
# Test deeply nested: {"outer": [{"inner": [val, val]},
# {"inner": [val]}]}
@@ -649,7 +760,7 @@ async def test_engine_core_client_util_method_nested_structures(
inner_list1 = inner_dict1["inner"]
assert isinstance(inner_list1, list) and len(inner_list1) == 2
for item in inner_list1:
assert isinstance(item, MyDataclass) and item.message == "nested3"
assert isinstance(item, TestMessage) and item.message == "nested3"
# Second dict in outer list should have "inner" with 1 item
inner_dict2 = outer_list[1]
@@ -657,7 +768,7 @@ async def test_engine_core_client_util_method_nested_structures(
inner_list2 = inner_dict2["inner"]
assert isinstance(inner_list2, list) and len(inner_list2) == 1
assert (
isinstance(inner_list2[0], MyDataclass)
isinstance(inner_list2[0], TestMessage)
and inner_list2[0].message == "nested3"
)