| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test that platform imports do not prematurely initialize CUDA. |
| 4 | |
| 5 | This is critical for Ray-based multi-GPU setups where workers need to |
| 6 | set CUDA_VISIBLE_DEVICES after importing vLLM but before CUDA is initialized. |
| 7 | If CUDA is initialized during import, device_count() gets locked and ignores |
| 8 | subsequent env var changes. |
| 9 | """ |
| 10 | |
| 11 | import subprocess |
| 12 | import sys |
| 13 | from pathlib import Path |
| 14 | |
| 15 | import pytest |
| 16 | |
| 17 | SCRIPTS_DIR = Path(__file__).parent / "scripts" |
| 18 | |
| 19 | |
| 20 | def run_script(script_name: str) -> subprocess.CompletedProcess: |
| 21 | """Run a test script in a subprocess with clean CUDA state.""" |
| 22 | script_path = SCRIPTS_DIR / script_name |
| 23 | return subprocess.run( |
| 24 | [sys.executable, str(script_path)], |
| 25 | capture_output=True, |
| 26 | text=True, |
| 27 | ) |
| 28 | |
| 29 | |
| 30 | def test_platform_import_does_not_init_cuda(): |
| 31 | """Test that importing vllm.platforms does not initialize CUDA.""" |
| 32 | result = run_script("check_platform_no_cuda_init.py") |
| 33 | if result.returncode != 0: |
| 34 | pytest.fail(f"Platform import initialized CUDA:\n{result.stderr}") |
| 35 | |
| 36 | |
| 37 | def test_device_count_respects_env_after_platform_import(): |
| 38 | """Test that device_count respects CUDA_VISIBLE_DEVICES after import.""" |
| 39 | result = run_script("check_device_count_respects_env.py") |
| 40 | if result.returncode != 0: |
| 41 | pytest.fail( |
| 42 | f"device_count does not respect env var after import:\n{result.stderr}" |
| 43 | ) |
| 44 | |
| 45 | |
| 46 | if __name__ == "__main__": |
| 47 | pytest.main([__file__, "-v"]) |