22import io
33import os
44import re
5+ import shutil
56import subprocess
67import warnings
78from pathlib import Path
3839# SUPPORTED_ARCHS = NVIDIA_SUPPORTED_ARCHS.union(ROCM_SUPPORTED_ARCHS)
3940
4041
42+ def _is_cuda () -> bool :
43+ return torch .version .cuda is not None
44+
45+
4146def _is_hip () -> bool :
4247 return torch .version .hip is not None
4348
@@ -46,15 +51,11 @@ def _is_neuron() -> bool:
4651 torch_neuronx_installed = True
4752 try :
4853 subprocess .run (["neuron-ls" ], capture_output = True , check = True )
49- except (FileNotFoundError , PermissionError ):
54+ except (FileNotFoundError , PermissionError , subprocess . CalledProcessError ):
5055 torch_neuronx_installed = False
5156 return torch_neuronx_installed
5257
5358
54- def _is_cuda () -> bool :
55- return (torch .version .cuda is not None ) and not _is_neuron ()
56-
57-
5859# Compiler flags.
5960CXX_FLAGS = ["-g" , "-O2" , "-std=c++17" ]
6061# TODO(woosuk): Should we use -O3?
@@ -400,7 +401,12 @@ def find_version(filepath: str) -> str:
400401def get_vllm_version () -> str :
401402 version = find_version (get_path ("vllm" , "__init__.py" ))
402403
403- if _is_hip ():
404+ if _is_cuda ():
405+ cuda_version = str (nvcc_cuda_version )
406+ if cuda_version != MAIN_CUDA_VERSION :
407+ cuda_version_str = cuda_version .replace ("." , "" )[:3 ]
408+ version += f"+cu{ cuda_version_str } "
409+ elif _is_hip ():
404410 # Get the HIP version
405411 hipcc_version = get_hipcc_rocm_version ()
406412 if hipcc_version != MAIN_CUDA_VERSION :
@@ -412,13 +418,8 @@ def get_vllm_version() -> str:
412418 if neuron_version != MAIN_CUDA_VERSION :
413419 neuron_version_str = neuron_version .replace ("." , "" )[:3 ]
414420 version += f"+neuron{ neuron_version_str } "
415- elif _is_cuda ():
416- cuda_version = str (nvcc_cuda_version )
417- if cuda_version != MAIN_CUDA_VERSION :
418- cuda_version_str = cuda_version .replace ("." , "" )[:3 ]
419- version += f"+cu{ cuda_version_str } "
420421 else :
421- raise RuntimeError ("Unknown runtime environment. " )
422+ raise RuntimeError ("Unknown runtime environment" )
422423
423424 return version
424425
@@ -434,13 +435,7 @@ def read_readme() -> str:
434435
435436def get_requirements () -> List [str ]:
436437 """Get Python package dependencies from requirements.txt."""
437- if _is_hip ():
438- with open (get_path ("requirements-rocm.txt" )) as f :
439- requirements = f .read ().strip ().split ("\n " )
440- elif _is_neuron ():
441- with open (get_path ("requirements-neuron.txt" )) as f :
442- requirements = f .read ().strip ().split ("\n " )
443- else :
438+ if _is_cuda ():
444439 with open (get_path ("requirements.txt" )) as f :
445440 requirements = f .read ().strip ().split ("\n " )
446441 if nvcc_cuda_version <= Version ("11.8" ):
@@ -449,6 +444,16 @@ def get_requirements() -> List[str]:
449444 if requirements [i ].startswith ("cupy-cuda12x" ):
450445 requirements [i ] = "cupy-cuda11x"
451446 break
447+ elif _is_hip ():
448+ with open (get_path ("requirements-rocm.txt" )) as f :
449+ requirements = f .read ().strip ().split ("\n " )
450+ elif _is_neuron ():
451+ with open (get_path ("requirements-neuron.txt" )) as f :
452+ requirements = f .read ().strip ().split ("\n " )
453+ else :
454+ raise ValueError (
455+ "Unsupported platform, please use CUDA, ROCM or Neuron." )
456+
452457 return requirements
453458
454459
0 commit comments