| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import sys |
| 4 | |
| 5 | import regex as re |
| 6 | |
| 7 | # --------------------------------------------------------------------------- # |
| 8 | # Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx` |
| 9 | # --------------------------------------------------------------------------- # |
| 10 | _TORCH_CUDA_PATTERNS = [ |
| 11 | r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|memory_reserved|memory_allocated|max_memory_allocated|max_memory_reserved|reset_peak_memory_stats|memory_stats|mem_get_info|set_device|device\()\b", |
| 12 | r"\btorch\.cuda\.(manual_seed|manual_seed_all)\b", |
| 13 | r"\bwith\storch\.cuda\.device\b", |
| 14 | # Calls torch.cuda.{_is_compiled/_device_count_amdsmi/_device_count_nvml} internally |
| 15 | r"\bcuda_device_count_stateless\(\)\b", |
| 16 | r"\bcurrent_platform\.mem_get_info\(\)\b", |
| 17 | ] |
| 18 | |
| 19 | ALLOWED_FILES = { |
| 20 | "vllm/platforms/", |
| 21 | "vllm/device_allocator/", |
| 22 | "vllm/distributed/weight_transfer/ipc_engine.py", |
| 23 | "vllm/distributed/weight_transfer/sharded_rdt_engine.py", |
| 24 | "tests/distributed/test_packed_tensor.py", |
| 25 | } |
| 26 | |
| 27 | |
| 28 | def scan_file(path: str) -> int: |
| 29 | with open(path, encoding="utf-8") as f: |
| 30 | content = f.read() |
| 31 | for pattern in _TORCH_CUDA_PATTERNS: |
| 32 | for match in re.finditer(pattern, content, re.MULTILINE): |
| 33 | # Calculate line number from match position |
| 34 | line_num = content[: match.start() + 1].count("\n") + 1 |
| 35 | matched_text = match.group(0) |
| 36 | if "manual_seed" in matched_text: |
| 37 | print( |
| 38 | f"{path}:{line_num}: " |
| 39 | "\033[91merror:\033[0m " |
| 40 | f"Found {matched_text} API call. Use set_random_seed instead." |
| 41 | ) |
| 42 | return 1 |
| 43 | print( |
| 44 | f"{path}:{line_num}: " |
| 45 | "\033[91merror:\033[0m " # red color |
| 46 | "Found torch.cuda API call. Please refer RFC " |
| 47 | "https://github.com/vllm-project/vllm/issues/30679, use " |
| 48 | "torch.accelerator API instead." |
| 49 | ) |
| 50 | return 1 |
| 51 | return 0 |
| 52 | |
| 53 | |
| 54 | def main(): |
| 55 | returncode = 0 |
| 56 | for filename in sys.argv[1:]: |
| 57 | if any(filename.startswith(prefix) for prefix in ALLOWED_FILES): |
| 58 | continue |
| 59 | returncode |= scan_file(filename) |
| 60 | return returncode |
| 61 | |
| 62 | |
| 63 | if __name__ == "__main__": |
| 64 | sys.exit(main()) |