[Bugfix] Shut down private Tensorizer engines (#49840)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Andreas Karatzas committed
Jul 30, 2026 at 14:33 UTC
30e333ca5f0119385123e7e713462ffc7c8d7e45
2 files changed
+96
-10
tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py
+63
-6
@@ -9,13 +9,14 @@ import pathlib
9
import subprocess
10
import sys
11
from typing import Any
12
+from unittest.mock import Mock
13
14
import pytest
15
import torch
16
17
import vllm.model_executor.model_loader.tensorizer
18
from tests.utils import VLLM_PATH, RemoteOpenAIServer
18
-from vllm import LLM, SamplingParams
19
+from vllm import SamplingParams
20
from vllm.engine.arg_utils import EngineArgs
21
from vllm.model_executor.model_loader.tensorizer import (
22
TensorizerConfig,
@@ -325,7 +326,7 @@ def test_load_with_just_model_tensors(just_serialize_model_tensors, model_ref):
326
pass
327
328
328
-def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path):
329
+def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path, vllm_runner):
330
serialization_params = {
331
"limit_cpu_concurrency": 2,
332
}
@@ -334,9 +335,6 @@ def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path):
335
config = TensorizerConfig(
336
tensorizer_uri=str(model_path), serialization_kwargs=serialization_params
337
)
337
- llm = LLM(
338
- model=model_ref,
339
- )
338
339
def serialization_test(self, *args, **kwargs):
340
# 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):
363
364
kwargs = {"tensorizer_config": config.to_serializable()}
365
368
- assert assert_from_collective_rpc(llm, serialization_test, kwargs)
366
+ with vllm_runner(model_ref) as runner:
367
+ assert assert_from_collective_rpc(runner.get_llm(), serialization_test, kwargs)
368
+
369
+
370
+@pytest.mark.parametrize(
371
+ (
372
+ "serialization_error",
373
+ "renderer_shutdown_error",
374
+ "engine_shutdown_error",
375
+ "expected_error",
376
+ ),
377
+ [
378
+ (None, None, None, None),
379
+ (RuntimeError("serialization failed"), None, None, "serialization failed"),
380
+ (
381
+ RuntimeError("serialization failed"),
382
+ ValueError("renderer shutdown failed"),
383
+ ValueError("engine shutdown failed"),
384
+ "serialization failed",
385
+ ),
386
+ (
387
+ None,
388
+ ValueError("renderer shutdown failed"),
389
+ None,
390
+ "renderer shutdown failed",
391
+ ),
392
+ (None, None, ValueError("engine shutdown failed"), "engine shutdown failed"),
393
+ ],
394
+)
395
+def test_tensorize_vllm_model_shuts_down_engine(
396
+ monkeypatch,
397
+ serialization_error,
398
+ renderer_shutdown_error,
399
+ engine_shutdown_error,
400
+ expected_error,
401
+):
402
+ from vllm.v1.engine.llm_engine import LLMEngine
403
+
404
+ engine_args = Mock()
405
+ config = Mock(encryption_keyfile=None)
406
+ engine = Mock()
407
+ engine.collective_rpc.side_effect = serialization_error
408
+ engine.renderer.shutdown.side_effect = renderer_shutdown_error
409
+ engine.engine_core.shutdown.side_effect = engine_shutdown_error
410
+ monkeypatch.setattr(LLMEngine, "from_vllm_config", Mock(return_value=engine))
411
+ monkeypatch.setenv("VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS", "7")
412
+
413
+ if expected_error is None:
414
+ tensorize_vllm_model(engine_args, config)
415
+ else:
416
+ with pytest.raises(Exception, match=expected_error) as exc_info:
417
+ tensorize_vllm_model(engine_args, config)
418
+ notes = getattr(exc_info.value, "__notes__", ())
419
+ if serialization_error is not None and renderer_shutdown_error is not None:
420
+ assert any("renderer shutdown failed" in note for note in notes)
421
+ if serialization_error is not None and engine_shutdown_error is not None:
422
+ assert any("engine shutdown failed" in note for note in notes)
423
+
424
+ engine.engine_core.shutdown.assert_called_once_with(timeout=17.0)
425
+ engine.renderer.shutdown.assert_called_once_with()
426
427
428
def test_assert_deserialization_kwargs_passed_to_tensor_deserializer(tmp_path, capfd):
vllm/model_executor/model_loader/tensorizer.py
+33
-4
@@ -66,6 +66,7 @@ __all__ = [
66
]
67
68
logger = init_logger(__name__)
69
+_TENSORIZER_ENGINE_CLEANUP_GRACE_S = 10.0
70
71
72
def is_valid_deserialization_uri(uri: str | None) -> bool:
@@ -730,10 +731,38 @@ def tensorize_vllm_model(
731
from vllm.v1.engine.llm_engine import LLMEngine
732
733
engine = LLMEngine.from_vllm_config(engine_config)
733
- engine.collective_rpc(
734
- "save_tensorized_model",
735
- kwargs={"tensorizer_config": tensorizer_config.to_serializable()},
736
- )
734
+ error: BaseException | None = None
735
+ try:
736
+ engine.collective_rpc(
737
+ "save_tensorized_model",
738
+ kwargs={"tensorizer_config": tensorizer_config.to_serializable()},
739
+ )
740
+ except BaseException as operation_error:
741
+ error = operation_error
742
+
743
+ def shutdown_engine_core() -> None:
744
+ engine.engine_core.shutdown(
745
+ timeout=(
746
+ envs.VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS
747
+ + _TENSORIZER_ENGINE_CLEANUP_GRACE_S
748
+ )
749
+ )
750
+
751
+ for name, callback in (
752
+ ("renderer", engine.renderer.shutdown),
753
+ ("engine core", shutdown_engine_core),
754
+ ):
755
+ try:
756
+ callback()
757
+ except BaseException as shutdown_error:
758
+ logger.exception("Failed to shut down tensorization %s", name)
759
+ if error is None:
760
+ error = shutdown_error
761
+ elif hasattr(error, "add_note"):
762
+ error.add_note(f"{name} shutdown also failed: {shutdown_error!r}")
763
+
764
+ if error is not None:
765
+ raise error
766
767
768
def tensorize_lora_adapter(lora_path: str, tensorizer_config: TensorizerConfig):