| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Pytest configuration for vLLM language generation 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 | def pytest_sessionstart(session): |
| 18 | """Configure ROCm-specific settings before test session starts.""" |
| 19 | if current_platform.is_rocm(): |
| 20 | torch.set_float32_matmul_precision("high") |
| 21 | |
| 22 | |
| 23 | @pytest.fixture(autouse=True) |
| 24 | def release_gpu_memory_between_tests(): |
| 25 | """Reclaim cached GPU memory between tests in this shard. |
| 26 | |
| 27 | These tests start many vLLM/HF engines sequentially in one pytest |
| 28 | process; without explicit reclamation between tests, VRAM fragmentation |
| 29 | accumulates and later tests fail at engine startup ("Engine core |
| 30 | initialization failed"). |
| 31 | """ |
| 32 | yield |
| 33 | gc.collect() |
| 34 | if torch.accelerator.is_available(): |
| 35 | torch.accelerator.empty_cache() |
| 36 | if current_platform.is_rocm(): |
| 37 | return |
| 38 | try: |
| 39 | wait_for_memory_to_settle() |
| 40 | except ValueError as e: |
| 41 | # Longer-lived (class/module-scoped) engine fixtures may legitimately |
| 42 | # still hold VRAM; nothing further to reclaim in that case. |
| 43 | logger.info("Failed to clean GPU memory: %s", e) |