From df7e12715fa5cac36069d3a58d869d9b939027a9 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 13 Jan 2026 01:14:30 -0600 Subject: [PATCH] [ROCm][CI] Fix engine core client tests for ROCm spawn multiprocessing (#32061) Signed-off-by: Andreas Karatzas --- tests/v1/engine/test_engine_core_client.py | 251 +++++++++++++++------ 1 file changed, 181 insertions(+), 70 deletions(-) diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 4b9c0e746..d5a904d18 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -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" )