| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Pytest configuration for vLLM pooling MTEB tests.""" |
| 4 | |
| 5 | import gc |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from tests.utils import wait_for_memory_to_settle |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.platforms import current_platform |
| 13 | |
| 14 | logger = init_logger(__name__) |
| 15 | |
| 16 | |
| 17 | @pytest.fixture(autouse=True) |
| 18 | def release_gpu_memory_between_tests(): |
| 19 | """Reclaim cached GPU memory between tests in this shard. |
| 20 | |
| 21 | These tests start many vLLM/HF engines sequentially in one pytest |
| 22 | process; without explicit reclamation between tests, VRAM fragmentation |
| 23 | accumulates and later tests fail at engine startup ("Engine core |
| 24 | initialization failed"). |
| 25 | """ |
| 26 | yield |
| 27 | gc.collect() |
| 28 | if torch.accelerator.is_available(): |
| 29 | torch.accelerator.empty_cache() |
| 30 | if current_platform.is_rocm(): |
| 31 | return |
| 32 | try: |
| 33 | wait_for_memory_to_settle() |
| 34 | except ValueError as e: |
| 35 | # Longer-lived (class/module-scoped) engine fixtures may legitimately |
| 36 | # still hold VRAM; nothing further to reclaim in that case. |
| 37 | logger.info("Failed to clean GPU memory: %s", e) |