diff --git a/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py b/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py index a15a624c90..7047ba1c17 100644 --- a/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py +++ b/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py @@ -9,13 +9,14 @@ import pathlib import subprocess import sys from typing import Any +from unittest.mock import Mock import pytest import torch import vllm.model_executor.model_loader.tensorizer from tests.utils import VLLM_PATH, RemoteOpenAIServer -from vllm import LLM, SamplingParams +from vllm import SamplingParams from vllm.engine.arg_utils import EngineArgs from vllm.model_executor.model_loader.tensorizer import ( TensorizerConfig, @@ -325,7 +326,7 @@ def test_load_with_just_model_tensors(just_serialize_model_tensors, model_ref): pass -def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path): +def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path, vllm_runner): serialization_params = { "limit_cpu_concurrency": 2, } @@ -334,9 +335,6 @@ def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path): config = TensorizerConfig( tensorizer_uri=str(model_path), serialization_kwargs=serialization_params ) - llm = LLM( - model=model_ref, - ) def serialization_test(self, *args, **kwargs): # This is performed in the ephemeral worker process, so monkey-patching @@ -365,7 +363,66 @@ def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path): kwargs = {"tensorizer_config": config.to_serializable()} - assert assert_from_collective_rpc(llm, serialization_test, kwargs) + with vllm_runner(model_ref) as runner: + assert assert_from_collective_rpc(runner.get_llm(), serialization_test, kwargs) + + +@pytest.mark.parametrize( + ( + "serialization_error", + "renderer_shutdown_error", + "engine_shutdown_error", + "expected_error", + ), + [ + (None, None, None, None), + (RuntimeError("serialization failed"), None, None, "serialization failed"), + ( + RuntimeError("serialization failed"), + ValueError("renderer shutdown failed"), + ValueError("engine shutdown failed"), + "serialization failed", + ), + ( + None, + ValueError("renderer shutdown failed"), + None, + "renderer shutdown failed", + ), + (None, None, ValueError("engine shutdown failed"), "engine shutdown failed"), + ], +) +def test_tensorize_vllm_model_shuts_down_engine( + monkeypatch, + serialization_error, + renderer_shutdown_error, + engine_shutdown_error, + expected_error, +): + from vllm.v1.engine.llm_engine import LLMEngine + + engine_args = Mock() + config = Mock(encryption_keyfile=None) + engine = Mock() + engine.collective_rpc.side_effect = serialization_error + engine.renderer.shutdown.side_effect = renderer_shutdown_error + engine.engine_core.shutdown.side_effect = engine_shutdown_error + monkeypatch.setattr(LLMEngine, "from_vllm_config", Mock(return_value=engine)) + monkeypatch.setenv("VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS", "7") + + if expected_error is None: + tensorize_vllm_model(engine_args, config) + else: + with pytest.raises(Exception, match=expected_error) as exc_info: + tensorize_vllm_model(engine_args, config) + notes = getattr(exc_info.value, "__notes__", ()) + if serialization_error is not None and renderer_shutdown_error is not None: + assert any("renderer shutdown failed" in note for note in notes) + if serialization_error is not None and engine_shutdown_error is not None: + assert any("engine shutdown failed" in note for note in notes) + + engine.engine_core.shutdown.assert_called_once_with(timeout=17.0) + engine.renderer.shutdown.assert_called_once_with() def test_assert_deserialization_kwargs_passed_to_tensor_deserializer(tmp_path, capfd): diff --git a/vllm/model_executor/model_loader/tensorizer.py b/vllm/model_executor/model_loader/tensorizer.py index 008abb6fdf..befe560929 100644 --- a/vllm/model_executor/model_loader/tensorizer.py +++ b/vllm/model_executor/model_loader/tensorizer.py @@ -66,6 +66,7 @@ __all__ = [ ] logger = init_logger(__name__) +_TENSORIZER_ENGINE_CLEANUP_GRACE_S = 10.0 def is_valid_deserialization_uri(uri: str | None) -> bool: @@ -730,10 +731,38 @@ def tensorize_vllm_model( from vllm.v1.engine.llm_engine import LLMEngine engine = LLMEngine.from_vllm_config(engine_config) - engine.collective_rpc( - "save_tensorized_model", - kwargs={"tensorizer_config": tensorizer_config.to_serializable()}, - ) + error: BaseException | None = None + try: + engine.collective_rpc( + "save_tensorized_model", + kwargs={"tensorizer_config": tensorizer_config.to_serializable()}, + ) + except BaseException as operation_error: + error = operation_error + + def shutdown_engine_core() -> None: + engine.engine_core.shutdown( + timeout=( + envs.VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS + + _TENSORIZER_ENGINE_CLEANUP_GRACE_S + ) + ) + + for name, callback in ( + ("renderer", engine.renderer.shutdown), + ("engine core", shutdown_engine_core), + ): + try: + callback() + except BaseException as shutdown_error: + logger.exception("Failed to shut down tensorization %s", name) + if error is None: + error = shutdown_error + elif hasattr(error, "add_note"): + error.add_note(f"{name} shutdown also failed: {shutdown_error!r}") + + if error is not None: + raise error def tensorize_lora_adapter(lora_path: str, tensorizer_config: TensorizerConfig):