2
2
import io
3
3
import os
4
4
import re
5
+ import shutil
5
6
import subprocess
6
7
import warnings
7
8
from pathlib import Path
38
39
# SUPPORTED_ARCHS = NVIDIA_SUPPORTED_ARCHS.union(ROCM_SUPPORTED_ARCHS)
39
40
40
41
42
+ def _is_cuda () -> bool :
43
+ return torch .version .cuda is not None
44
+
45
+
41
46
def _is_hip () -> bool :
42
47
return torch .version .hip is not None
43
48
@@ -46,15 +51,11 @@ def _is_neuron() -> bool:
46
51
torch_neuronx_installed = True
47
52
try :
48
53
subprocess .run (["neuron-ls" ], capture_output = True , check = True )
49
- except (FileNotFoundError , PermissionError ):
54
+ except (FileNotFoundError , PermissionError , subprocess . CalledProcessError ):
50
55
torch_neuronx_installed = False
51
56
return torch_neuronx_installed
52
57
53
58
54
- def _is_cuda () -> bool :
55
- return (torch .version .cuda is not None ) and not _is_neuron ()
56
-
57
-
58
59
# Compiler flags.
59
60
CXX_FLAGS = ["-g" , "-O2" , "-std=c++17" ]
60
61
# TODO(woosuk): Should we use -O3?
@@ -400,7 +401,12 @@ def find_version(filepath: str) -> str:
400
401
def get_vllm_version () -> str :
401
402
version = find_version (get_path ("vllm" , "__init__.py" ))
402
403
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 ():
404
410
# Get the HIP version
405
411
hipcc_version = get_hipcc_rocm_version ()
406
412
if hipcc_version != MAIN_CUDA_VERSION :
@@ -412,13 +418,8 @@ def get_vllm_version() -> str:
412
418
if neuron_version != MAIN_CUDA_VERSION :
413
419
neuron_version_str = neuron_version .replace ("." , "" )[:3 ]
414
420
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 } "
420
421
else :
421
- raise RuntimeError ("Unknown runtime environment. " )
422
+ raise RuntimeError ("Unknown runtime environment" )
422
423
423
424
return version
424
425
@@ -434,13 +435,7 @@ def read_readme() -> str:
434
435
435
436
def get_requirements () -> List [str ]:
436
437
"""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 ():
444
439
with open (get_path ("requirements.txt" )) as f :
445
440
requirements = f .read ().strip ().split ("\n " )
446
441
if nvcc_cuda_version <= Version ("11.8" ):
@@ -449,6 +444,16 @@ def get_requirements() -> List[str]:
449
444
if requirements [i ].startswith ("cupy-cuda12x" ):
450
445
requirements [i ] = "cupy-cuda11x"
451
446
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
+
452
457
return requirements
453
458
454
459
0 commit comments