df7e12715f · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -3,6 +3,7 @@

3

4 import asyncio

5 import importlib

6 +import inspect

7 import os

8 import signal

9 import time

@@ -47,6 +48,7 @@ MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"

48 TOKENIZER = AutoTokenizer.from_pretrained(MODEL_NAME)

49 PROMPT = "Hello my name is Robert and I love quantization kernels"

50 PROMPT_TOKENS = TOKENIZER(PROMPT).input_ids

51 +TEST_MODULE = "tests.v1.engine.test_engine_core_client"

52

53 _REQUEST_COUNTER = 0

54

@@ -223,10 +225,170 @@ def echo(self, msg: str, err_msg: str | None = None, sleep: float | None = None)

225 return msg

226

227

228 +@dataclass

229 +class TestMessage:

230 + """Test dataclass for verifying custom type serialization."""

231 +

232 + message: str

233 +

234 +

235 +# Dummy utility function to monkey-patch into engine core.

236 +def echo_dc(

237 + self,

238 + msg: str,

239 + return_list: bool = False,

240 +) -> TestMessage | list[TestMessage]:

241 + print(f"echo dc util function called: {msg}")

242 + val = None if msg is None else TestMessage(msg)

243 + # Return dataclass to verify support for returning custom types

244 + # (for which there is special handling to make it work with msgspec).

245 + return [val for _ in range(3)] if return_list else val

246 +

247 +

248 +# Dummy utility function to test dict serialization with custom types.

249 +def echo_dc_dict(

250 + self,

251 + msg: str,

252 + return_dict: bool = False,

253 +) -> TestMessage | dict[str, TestMessage]:

254 + print(f"echo dc dict util function called: {msg}")

255 + val = None if msg is None else TestMessage(msg)

256 + # Return dict of dataclasses to verify support for returning dicts

257 + # with custom value types.

258 + if return_dict:

259 + return {"key1": val, "key2": val, "key3": val}

260 + else:

261 + return val

262 +

263 +

264 +# Dummy utility function to test nested structures with custom types.

265 +def echo_dc_nested(

266 + self,

267 + msg: str,

268 + structure_type: str = "list_of_dicts",

269 +) -> Any:

270 + print(f"echo dc nested util function called: {msg}, structure: {structure_type}")

271 + val = None if msg is None else TestMessage(msg)

272 +

273 + structures = {

274 + "list_of_dicts": [{"a": val, "b": val}, {"c": val, "d": val}],

275 + "dict_of_lists": {"list1": [val, val], "list2": [val, val]},

276 + "deep_nested": {"outer": [{"inner": [val, val]}, {"inner": [val]}]},

277 + }

278 + return structures.get(structure_type, val)

279 +

280 +

281 +# --- Fixtures for subprocess patching ---

282 +# These create sitecustomize.py files that patch EngineCore in spawned

283 +# subprocesses. This is necessary because ROCm requires 'spawn' multiprocessing

284 +# start method, which creates fresh Python interpreters that don't inherit

285 +# monkey-patches from the parent process.

286 +

287 +

288 +@pytest.fixture

289 +def subprocess_echo_patch(monkeypatch, tmp_path):

290 + """Create sitecustomize.py so spawned subprocesses have echo method.

291 +

292 + This is needed because ROCm uses 'spawn' multiprocessing start method,

293 + which creates a fresh Python interpreter that doesn't inherit monkey-patches.

294 + By using sitecustomize.py, we ensure the patch is applied when Python starts.

295 + """

296 + sc = tmp_path / "sitecustomize.py"

297 + sc.write_text(

298 + "\n".join(

299 + [

300 + "import time",

301 + "from vllm.v1.engine.core import EngineCore",

302 + inspect.getsource(echo),

303 + "EngineCore.echo = echo",

304 + ]

305 + )

306 + )

307 + monkeypatch.setenv(

308 + "PYTHONPATH",

309 + os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),

310 + )

311 +

312 +

313 +@pytest.fixture

314 +def subprocess_echo_dc_patch(monkeypatch, tmp_path):

315 + """Create sitecustomize.py so spawned subprocesses have echo_dc method."""

316 + sc = tmp_path / "sitecustomize.py"

317 + sc.write_text(

318 + "\n".join(

319 + [

320 + "from dataclasses import dataclass",

321 + "",

322 + inspect.getsource(TestMessage),

323 + f"TestMessage.__module__ = '{TEST_MODULE}'",

324 + "",

325 + "from vllm.v1.engine.core import EngineCore",

326 + inspect.getsource(echo_dc),

327 + "EngineCore.echo_dc = echo_dc",

328 + ]

329 + )

330 + )

331 + monkeypatch.setenv(

332 + "PYTHONPATH",

333 + os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),

334 + )

335 +

336 +

337 +@pytest.fixture

338 +def subprocess_echo_dc_dict_patch(monkeypatch, tmp_path):

339 + """Create sitecustomize.py so spawned subprocesses have echo_dc_dict method."""

340 + sc = tmp_path / "sitecustomize.py"

341 + sc.write_text(

342 + "\n".join(

343 + [

344 + "from dataclasses import dataclass",

345 + "",

346 + inspect.getsource(TestMessage),

347 + f"TestMessage.__module__ = '{TEST_MODULE}'",

348 + "",

349 + "from vllm.v1.engine.core import EngineCore",

350 + inspect.getsource(echo_dc_dict),

351 + "EngineCore.echo_dc_dict = echo_dc_dict",

352 + ]

353 + )

354 + )

355 + monkeypatch.setenv(

356 + "PYTHONPATH",

357 + os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),

358 + )

359 +

360 +

361 +@pytest.fixture

362 +def subprocess_echo_dc_nested_patch(monkeypatch, tmp_path):

363 + """Create sitecustomize.py so spawned subprocesses have echo_dc_nested method."""

364 + sc = tmp_path / "sitecustomize.py"

365 + sc.write_text(

366 + "\n".join(

367 + [

368 + "from dataclasses import dataclass",

369 + "from typing import Any",

370 + "",

371 + inspect.getsource(TestMessage),

372 + f"TestMessage.__module__ = '{TEST_MODULE}'",

373 + "",

374 + "from vllm.v1.engine.core import EngineCore",

375 + inspect.getsource(echo_dc_nested),

376 + "EngineCore.echo_dc_nested = echo_dc_nested",

377 + ]

378 + )

379 + )

380 + monkeypatch.setenv(

381 + "PYTHONPATH",

382 + os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),

383 + )

384 +

385 +

386 @create_new_process_for_each_test()

387 @pytest.mark.parametrize("multiprocessing_mode", [True, False])

388 def test_engine_core_client(

229 - monkeypatch: pytest.MonkeyPatch, multiprocessing_mode: bool

389 + monkeypatch: pytest.MonkeyPatch,

390 + multiprocessing_mode: bool,

391 + subprocess_echo_patch,

392 ):

393 with monkeypatch.context() as m:

394 # Monkey-patch core engine utility function to test.

@@ -313,7 +475,10 @@ def test_engine_core_client(

475

476

477 @pytest.mark.asyncio(loop_scope="function")

316 -async def test_engine_core_client_asyncio(monkeypatch: pytest.MonkeyPatch):

478 +async def test_engine_core_client_asyncio(

479 + monkeypatch: pytest.MonkeyPatch,

480 + subprocess_echo_patch,

481 +):

482 with monkeypatch.context() as m:

483 # Monkey-patch core engine utility function to test.

484 m.setattr(EngineCore, "echo", echo, raising=False)

@@ -406,66 +571,10 @@ async def test_engine_core_client_asyncio(monkeypatch: pytest.MonkeyPatch):

571 client.shutdown()

572

573

409 -@dataclass

410 -class MyDataclass:

411 - message: str

412 -

413 -

414 -# Dummy utility function to monkey-patch into engine core.

415 -def echo_dc(

416 - self,

417 - msg: str,

418 - return_list: bool = False,

419 -) -> MyDataclass | list[MyDataclass]:

420 - print(f"echo dc util function called: {msg}")

421 - val = None if msg is None else MyDataclass(msg)

422 - # Return dataclass to verify support for returning custom types

423 - # (for which there is special handling to make it work with msgspec).

424 - return [val for _ in range(3)] if return_list else val

425 -

426 -

427 -# Dummy utility function to test dict serialization with custom types.

428 -def echo_dc_dict(

429 - self,

430 - msg: str,

431 - return_dict: bool = False,

432 -) -> MyDataclass | dict[str, MyDataclass]:

433 - print(f"echo dc dict util function called: {msg}")

434 - val = None if msg is None else MyDataclass(msg)

435 - # Return dict of dataclasses to verify support for returning dicts

436 - # with custom value types.

437 - if return_dict:

438 - return {"key1": val, "key2": val, "key3": val}

439 - else:

440 - return val

441 -

442 -

443 -# Dummy utility function to test nested structures with custom types.

444 -def echo_dc_nested(

445 - self,

446 - msg: str,

447 - structure_type: str = "list_of_dicts",

448 -) -> Any:

449 - print(f"echo dc nested util function called: {msg}, structure: {structure_type}")

450 - val = None if msg is None else MyDataclass(msg)

451 -

452 - if structure_type == "list_of_dicts": # noqa

453 - # Return list of dicts: [{"a": val, "b": val}, {"c": val, "d": val}]

454 - return [{"a": val, "b": val}, {"c": val, "d": val}]

455 - elif structure_type == "dict_of_lists":

456 - # Return dict of lists: {"list1": [val, val], "list2": [val, val]}

457 - return {"list1": [val, val], "list2": [val, val]}

458 - elif structure_type == "deep_nested":

459 - # Return deeply nested: {"outer": [{"inner": [val, val]},

460 - # {"inner": [val]}]}

461 - return {"outer": [{"inner": [val, val]}, {"inner": [val]}]}

462 - else:

463 - return val

464 -

465 -

574 @pytest.mark.asyncio(loop_scope="function")

575 async def test_engine_core_client_util_method_custom_return(

576 monkeypatch: pytest.MonkeyPatch,

577 + subprocess_echo_dc_patch,

578 ):

579 with monkeypatch.context() as m:

580 # Must set insecure serialization to allow returning custom types.

@@ -494,10 +603,10 @@ async def test_engine_core_client_util_method_custom_return(

603 core_client: AsyncMPClient = client

604

605 result = await core_client.call_utility_async("echo_dc", "testarg2", False)

497 - assert isinstance(result, MyDataclass) and result.message == "testarg2"

606 + assert isinstance(result, TestMessage) and result.message == "testarg2"

607 result = await core_client.call_utility_async("echo_dc", "testarg2", True)

608 assert isinstance(result, list) and all(

500 - isinstance(r, MyDataclass) and r.message == "testarg2" for r in result

609 + isinstance(r, TestMessage) and r.message == "testarg2" for r in result

610 )

611

612 # Test returning None and list of Nones

@@ -513,6 +622,7 @@ async def test_engine_core_client_util_method_custom_return(

622 @pytest.mark.asyncio(loop_scope="function")

623 async def test_engine_core_client_util_method_custom_dict_return(

624 monkeypatch: pytest.MonkeyPatch,

625 + subprocess_echo_dc_dict_patch,

626 ):

627 with monkeypatch.context() as m:

628 # Must set insecure serialization to allow returning custom types.

@@ -544,7 +654,7 @@ async def test_engine_core_client_util_method_custom_dict_return(

654 result = await core_client.call_utility_async(

655 "echo_dc_dict", "testarg3", False

656 )

547 - assert isinstance(result, MyDataclass) and result.message == "testarg3"

657 + assert isinstance(result, TestMessage) and result.message == "testarg3"

658

659 # Test dict return with custom value types

660 result = await core_client.call_utility_async(

@@ -553,7 +663,7 @@ async def test_engine_core_client_util_method_custom_dict_return(

663 assert isinstance(result, dict) and len(result) == 3

664 for key, val in result.items():

665 assert key in ["key1", "key2", "key3"]

556 - assert isinstance(val, MyDataclass) and val.message == "testarg3"

666 + assert isinstance(val, TestMessage) and val.message == "testarg3"

667

668 # Test returning dict with None values

669 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(

679 @pytest.mark.asyncio(loop_scope="function")

680 async def test_engine_core_client_util_method_nested_structures(

681 monkeypatch: pytest.MonkeyPatch,

682 + subprocess_echo_dc_nested_patch,

683 ):

684 with monkeypatch.context() as m:

685 # Must set insecure serialization to allow returning custom types.

@@ -605,21 +716,21 @@ async def test_engine_core_client_util_method_nested_structures(

716 if i == 0:

717 assert "a" in item and "b" in item

718 assert (

608 - isinstance(item["a"], MyDataclass)

719 + isinstance(item["a"], TestMessage)

720 and item["a"].message == "nested1"

721 )

722 assert (

612 - isinstance(item["b"], MyDataclass)

723 + isinstance(item["b"], TestMessage)

724 and item["b"].message == "nested1"

725 )

726 else:

727 assert "c" in item and "d" in item

728 assert (

618 - isinstance(item["c"], MyDataclass)

729 + isinstance(item["c"], TestMessage)

730 and item["c"].message == "nested1"

731 )

732 assert (

622 - isinstance(item["d"], MyDataclass)

733 + isinstance(item["d"], TestMessage)

734 and item["d"].message == "nested1"

735 )

736

@@ -632,7 +743,7 @@ async def test_engine_core_client_util_method_nested_structures(

743 for key, lst in result.items():

744 assert isinstance(lst, list) and len(lst) == 2

745 for item in lst:

635 - assert isinstance(item, MyDataclass) and item.message == "nested2"

746 + assert isinstance(item, TestMessage) and item.message == "nested2"

747

748 # Test deeply nested: {"outer": [{"inner": [val, val]},

749 # {"inner": [val]}]}

@@ -649,7 +760,7 @@ async def test_engine_core_client_util_method_nested_structures(

760 inner_list1 = inner_dict1["inner"]

761 assert isinstance(inner_list1, list) and len(inner_list1) == 2

762 for item in inner_list1:

652 - assert isinstance(item, MyDataclass) and item.message == "nested3"

763 + assert isinstance(item, TestMessage) and item.message == "nested3"

764

765 # Second dict in outer list should have "inner" with 1 item

766 inner_dict2 = outer_list[1]

@@ -657,7 +768,7 @@ async def test_engine_core_client_util_method_nested_structures(

768 inner_list2 = inner_dict2["inner"]

769 assert isinstance(inner_list2, list) and len(inner_list2) == 1

770 assert (

660 - isinstance(inner_list2[0], MyDataclass)

771 + isinstance(inner_list2[0], TestMessage)

772 and inner_list2[0].message == "nested3"

773 )

774

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