|
| 1 | +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +"""Tests for the VllmBackend teardown lifecycle.""" |
| 5 | + |
| 6 | +from unittest.mock import MagicMock, patch |
| 7 | + |
| 8 | +import pytest |
| 9 | + |
| 10 | +from nemo_safe_synthesizer.generation import vllm_backend as vllm_backend_mod |
| 11 | +from nemo_safe_synthesizer.generation.vllm_backend import VllmBackend |
| 12 | + |
| 13 | + |
| 14 | +@pytest.fixture |
| 15 | +def _mock_vllm_cleanup(): |
| 16 | + """Patch vLLM distributed cleanup so tests run without a GPU.""" |
| 17 | + with ( |
| 18 | + patch.object(vllm_backend_mod, "cleanup_dist_env_and_memory") as mock_dist, |
| 19 | + patch.object(vllm_backend_mod, "cleanup_memory") as mock_mem, |
| 20 | + ): |
| 21 | + yield mock_dist, mock_mem |
| 22 | + |
| 23 | + |
| 24 | +@pytest.fixture |
| 25 | +def backend(_mock_vllm_cleanup, fixture_session_cache_dir): |
| 26 | + """Create a VllmBackend with mocked dependencies.""" |
| 27 | + mock_metadata = MagicMock() |
| 28 | + mock_metadata.adapter_path = None |
| 29 | + mock_metadata.instruction = "Generate" |
| 30 | + mock_metadata.prompt_config = MagicMock() |
| 31 | + mock_metadata.prompt_config.template = "{instruction} {schema}" |
| 32 | + |
| 33 | + mock_config = MagicMock() |
| 34 | + # Pin branching fields so create_processor() selects TabularDataProcessor deterministically. |
| 35 | + mock_config.time_series.is_timeseries = False |
| 36 | + mock_config.data.group_training_examples_by = None |
| 37 | + # Pin to a valid literal so StructuredOutputsConfig Pydantic validation passes in initialize(). |
| 38 | + mock_config.generation.structured_generation_backend = "xgrammar" |
| 39 | + mock_config.generation.attention_backend = None |
| 40 | + |
| 41 | + mock_workdir = MagicMock() |
| 42 | + mock_workdir.schema_file = fixture_session_cache_dir / "schema.json" |
| 43 | + mock_workdir.schema_file.parent.mkdir(parents=True, exist_ok=True) |
| 44 | + mock_workdir.schema_file.write_text('{"properties": {"col_a": {"type": "string"}}}') |
| 45 | + mock_workdir.adapter_path = None |
| 46 | + |
| 47 | + return VllmBackend(config=mock_config, model_metadata=mock_metadata, workdir=mock_workdir) |
| 48 | + |
| 49 | + |
| 50 | +class TestTeardownIdempotency: |
| 51 | + def test_first_teardown_runs_cleanup(self, backend, _mock_vllm_cleanup): |
| 52 | + mock_dist, mock_mem = _mock_vllm_cleanup |
| 53 | + backend.llm = MagicMock() |
| 54 | + |
| 55 | + backend.teardown() |
| 56 | + |
| 57 | + mock_dist.assert_called_once() |
| 58 | + mock_mem.assert_called_once() |
| 59 | + assert backend.llm is None |
| 60 | + assert backend._torn_down is True |
| 61 | + |
| 62 | + def test_second_teardown_is_noop(self, backend, _mock_vllm_cleanup): |
| 63 | + mock_dist, mock_mem = _mock_vllm_cleanup |
| 64 | + |
| 65 | + backend.teardown() |
| 66 | + mock_dist.reset_mock() |
| 67 | + mock_mem.reset_mock() |
| 68 | + |
| 69 | + backend.teardown() |
| 70 | + |
| 71 | + mock_dist.assert_not_called() |
| 72 | + mock_mem.assert_not_called() |
| 73 | + |
| 74 | + def test_initialize_resets_guard(self, backend, _mock_vllm_cleanup): |
| 75 | + backend.teardown() |
| 76 | + assert backend._torn_down is True |
| 77 | + |
| 78 | + with patch.object(vllm_backend_mod, "vLLM"): |
| 79 | + backend.initialize() |
| 80 | + |
| 81 | + assert backend._torn_down is False |
| 82 | + |
| 83 | + |
| 84 | +class TestTeardownResilience: |
| 85 | + def test_cleanup_memory_runs_even_if_dist_cleanup_fails(self, backend, _mock_vllm_cleanup): |
| 86 | + mock_dist, mock_mem = _mock_vllm_cleanup |
| 87 | + mock_dist.side_effect = RuntimeError("distributed cleanup failed") |
| 88 | + |
| 89 | + backend.teardown() |
| 90 | + |
| 91 | + mock_dist.assert_called_once() |
| 92 | + mock_mem.assert_called_once() |
| 93 | + assert backend.llm is None |
| 94 | + |
| 95 | + def test_llm_cleared_even_if_dist_cleanup_fails(self, backend, _mock_vllm_cleanup): |
| 96 | + mock_dist, _ = _mock_vllm_cleanup |
| 97 | + mock_dist.side_effect = RuntimeError("boom") |
| 98 | + backend.llm = MagicMock() |
| 99 | + |
| 100 | + backend.teardown() |
| 101 | + |
| 102 | + assert backend.llm is None |
| 103 | + |
| 104 | + |
| 105 | +class TestDunderDel: |
| 106 | + def test_del_calls_teardown(self, backend, _mock_vllm_cleanup): |
| 107 | + mock_dist, _ = _mock_vllm_cleanup |
| 108 | + |
| 109 | + # Reset to isolate only this explicit __del__ call. |
| 110 | + mock_dist.reset_mock() |
| 111 | + backend.__del__() |
| 112 | + |
| 113 | + mock_dist.assert_called_once() |
| 114 | + assert backend._torn_down is True |
| 115 | + |
| 116 | + def test_del_suppresses_exceptions(self, backend, _mock_vllm_cleanup): |
| 117 | + mock_dist, _ = _mock_vllm_cleanup |
| 118 | + mock_dist.side_effect = RuntimeError("boom") |
| 119 | + |
| 120 | + backend.__del__() |
| 121 | + |
| 122 | + def test_del_after_teardown_is_noop(self, backend, _mock_vllm_cleanup): |
| 123 | + mock_dist, _ = _mock_vllm_cleanup |
| 124 | + |
| 125 | + backend.teardown() |
| 126 | + mock_dist.reset_mock() |
| 127 | + |
| 128 | + backend.__del__() |
| 129 | + |
| 130 | + mock_dist.assert_not_called() |
0 commit comments