39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Verify the pinned vLLM and compiled CUDA imports on a GPU CI worker."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
import re
|
|
from pathlib import Path
|
|
|
|
|
|
def main() -> None:
|
|
dockerfile = Path(__file__).resolve().parents[3] / "docker" / "Dockerfile.ci"
|
|
match = re.search(
|
|
r"^ARG VLLM_VERSION=([0-9]+\.[0-9]+\.[0-9]+)$",
|
|
dockerfile.read_text(),
|
|
re.MULTILINE,
|
|
)
|
|
if match is None:
|
|
raise RuntimeError(f"Missing vLLM release version in {dockerfile}")
|
|
expected_version = match.group(1)
|
|
|
|
import torch
|
|
import vllm
|
|
|
|
assert vllm.__version__ == expected_version, vllm.__version__
|
|
assert torch.cuda.is_available(), "vLLM runtime check requires a GPU worker"
|
|
for module in ("flashinfer", "vllm._C_stable_libtorch"):
|
|
importlib.import_module(module)
|
|
print(
|
|
f"Verified vLLM CUDA runtime: vllm={vllm.__version__} "
|
|
f"expected={expected_version} torch={torch.__version__} "
|
|
f"cuda={torch.version.cuda} device={torch.cuda.get_device_name()}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|