| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Audit vLLM compiled libraries for PyTorch stable ABI compliance.""" |
| 5 | |
| 6 | import fnmatch |
| 7 | import sys |
| 8 | from pathlib import Path |
| 9 | |
| 10 | from torch_abi_audit import inspect_package |
| 11 | from torch_abi_audit.report import ExtensionReport, PackageReport |
| 12 | |
| 13 | # Temporary allowlist of extensions not yet on the stable ABI. |
| 14 | # Shrink and remove over time. |
| 15 | ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = ( |
| 16 | # ROCm-only legacy extensions. Most ROCm ops have moved to |
| 17 | # _C_stable_libtorch, but these two modules still bind libtorch directly. |
| 18 | "_C.abi3.so", |
| 19 | "_rocm_C.abi3.so", |
| 20 | "third_party/deep_gemm/_C*.so", |
| 21 | ) |
| 22 | |
| 23 | |
| 24 | def _relative_path(lib: ExtensionReport, package_root: Path) -> str: |
| 25 | try: |
| 26 | return lib.path.relative_to(package_root).as_posix() |
| 27 | except ValueError: |
| 28 | return lib.path.name |
| 29 | |
| 30 | |
| 31 | def _is_torch_unstable(lib: ExtensionReport) -> bool: |
| 32 | return lib.error is None and lib.torch.uses_torch and not lib.torch.stable |
| 33 | |
| 34 | |
| 35 | def _matches_allowlist(rel_path: str, patterns: tuple[str, ...]) -> bool: |
| 36 | return any(fnmatch.fnmatch(rel_path, pattern) for pattern in patterns) |
| 37 | |
| 38 | |
| 39 | def _iter_libs(report: PackageReport) -> tuple[ExtensionReport, ...]: |
| 40 | return (*report.extensions, *report.bundled_libs) |
| 41 | |
| 42 | |
| 43 | def _collect_unstable(report: PackageReport) -> list[str]: |
| 44 | return sorted( |
| 45 | _relative_path(lib, report.root) |
| 46 | for lib in _iter_libs(report) |
| 47 | if _is_torch_unstable(lib) |
| 48 | ) |
| 49 | |
| 50 | |
| 51 | def _find_stale_allowlist_entries( |
| 52 | report: PackageReport, patterns: tuple[str, ...] |
| 53 | ) -> list[str]: |
| 54 | """Allowlist patterns that match a built library which is no longer unstable.""" |
| 55 | stale: list[str] = [] |
| 56 | for pattern in patterns: |
| 57 | for lib in _iter_libs(report): |
| 58 | if lib.error is not None: |
| 59 | continue |
| 60 | if not fnmatch.fnmatch(_relative_path(lib, report.root), pattern): |
| 61 | continue |
| 62 | if not _is_torch_unstable(lib): |
| 63 | stale.append(pattern) |
| 64 | break |
| 65 | return stale |
| 66 | |
| 67 | |
| 68 | def check_torch_abi( |
| 69 | package: str = "vllm", |
| 70 | patterns: tuple[str, ...] = ALLOWED_UNSTABLE_LIBRARIES, |
| 71 | ) -> int: |
| 72 | report = inspect_package(package) |
| 73 | if report.error: |
| 74 | print(f"error: failed to inspect {package!r}: {report.error}", file=sys.stderr) |
| 75 | return 2 |
| 76 | |
| 77 | unstable = _collect_unstable(report) |
| 78 | unexpected = [ |
| 79 | rel_path for rel_path in unstable if not _matches_allowlist(rel_path, patterns) |
| 80 | ] |
| 81 | stale = _find_stale_allowlist_entries(report, patterns) |
| 82 | |
| 83 | if unexpected or stale: |
| 84 | if unexpected: |
| 85 | print( |
| 86 | "Not allowed: torch-unstable libraries outside " |
| 87 | f"ALLOWED_UNSTABLE_LIBRARIES: {', '.join(unexpected)}", |
| 88 | file=sys.stderr, |
| 89 | ) |
| 90 | if stale: |
| 91 | print( |
| 92 | "Not allowed: stale ALLOWED_UNSTABLE_LIBRARIES entries: " |
| 93 | f"{', '.join(stale)}", |
| 94 | file=sys.stderr, |
| 95 | ) |
| 96 | return 1 |
| 97 | |
| 98 | print("Torch stable ABI check passed.") |
| 99 | return 0 |
| 100 | |
| 101 | |
| 102 | if __name__ == "__main__": |
| 103 | print(">>> Auditing vLLM extension modules for PyTorch stable ABI compliance") |
| 104 | sys.exit(check_torch_abi()) |