Skip to content

Commit 6b78837

Browse files
authored
Fix setup.py neuron-ls issue (#2671)
1 parent 120157f commit 6b78837

File tree

1 file changed

+24
-19
lines changed

1 file changed

+24
-19
lines changed

setup.py

Lines changed: 24 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import io
33
import os
44
import re
5+
import shutil
56
import subprocess
67
import warnings
78
from pathlib import Path
@@ -38,6 +39,10 @@
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+
4146
def _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.
5960
CXX_FLAGS = ["-g", "-O2", "-std=c++17"]
6061
# TODO(woosuk): Should we use -O3?
@@ -400,7 +401,12 @@ def find_version(filepath: str) -> str:
400401
def 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

435436
def 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

Comments
 (0)