Fix setup.py neuron-ls issue (#2671)

This commit is contained in:
Simon Mo 2024-03-16 16:00:25 -07:00 committed by GitHub
parent 120157fd2a
commit 6b78837b29
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -2,6 +2,7 @@ import contextlib
import io import io
import os import os
import re import re
import shutil
import subprocess import subprocess
import warnings import warnings
from pathlib import Path from pathlib import Path
@ -38,6 +39,10 @@ ROCM_SUPPORTED_ARCHS = {"gfx908", "gfx90a", "gfx942", "gfx1100"}
# SUPPORTED_ARCHS = NVIDIA_SUPPORTED_ARCHS.union(ROCM_SUPPORTED_ARCHS) # SUPPORTED_ARCHS = NVIDIA_SUPPORTED_ARCHS.union(ROCM_SUPPORTED_ARCHS)
def _is_cuda() -> bool:
return torch.version.cuda is not None
def _is_hip() -> bool: def _is_hip() -> bool:
return torch.version.hip is not None return torch.version.hip is not None
@ -46,15 +51,11 @@ def _is_neuron() -> bool:
torch_neuronx_installed = True torch_neuronx_installed = True
try: try:
subprocess.run(["neuron-ls"], capture_output=True, check=True) subprocess.run(["neuron-ls"], capture_output=True, check=True)
except (FileNotFoundError, PermissionError): except (FileNotFoundError, PermissionError, subprocess.CalledProcessError):
torch_neuronx_installed = False torch_neuronx_installed = False
return torch_neuronx_installed return torch_neuronx_installed
def _is_cuda() -> bool:
return (torch.version.cuda is not None) and not _is_neuron()
# Compiler flags. # Compiler flags.
CXX_FLAGS = ["-g", "-O2", "-std=c++17"] CXX_FLAGS = ["-g", "-O2", "-std=c++17"]
# TODO(woosuk): Should we use -O3? # TODO(woosuk): Should we use -O3?
@ -400,7 +401,12 @@ def find_version(filepath: str) -> str:
def get_vllm_version() -> str: def get_vllm_version() -> str:
version = find_version(get_path("vllm", "__init__.py")) version = find_version(get_path("vllm", "__init__.py"))
if _is_hip(): if _is_cuda():
cuda_version = str(nvcc_cuda_version)
if cuda_version != MAIN_CUDA_VERSION:
cuda_version_str = cuda_version.replace(".", "")[:3]
version += f"+cu{cuda_version_str}"
elif _is_hip():
# Get the HIP version # Get the HIP version
hipcc_version = get_hipcc_rocm_version() hipcc_version = get_hipcc_rocm_version()
if hipcc_version != MAIN_CUDA_VERSION: if hipcc_version != MAIN_CUDA_VERSION:
@ -412,13 +418,8 @@ def get_vllm_version() -> str:
if neuron_version != MAIN_CUDA_VERSION: if neuron_version != MAIN_CUDA_VERSION:
neuron_version_str = neuron_version.replace(".", "")[:3] neuron_version_str = neuron_version.replace(".", "")[:3]
version += f"+neuron{neuron_version_str}" version += f"+neuron{neuron_version_str}"
elif _is_cuda():
cuda_version = str(nvcc_cuda_version)
if cuda_version != MAIN_CUDA_VERSION:
cuda_version_str = cuda_version.replace(".", "")[:3]
version += f"+cu{cuda_version_str}"
else: else:
raise RuntimeError("Unknown runtime environment.") raise RuntimeError("Unknown runtime environment")
return version return version
@ -434,13 +435,7 @@ def read_readme() -> str:
def get_requirements() -> List[str]: def get_requirements() -> List[str]:
"""Get Python package dependencies from requirements.txt.""" """Get Python package dependencies from requirements.txt."""
if _is_hip(): if _is_cuda():
with open(get_path("requirements-rocm.txt")) as f:
requirements = f.read().strip().split("\n")
elif _is_neuron():
with open(get_path("requirements-neuron.txt")) as f:
requirements = f.read().strip().split("\n")
else:
with open(get_path("requirements.txt")) as f: with open(get_path("requirements.txt")) as f:
requirements = f.read().strip().split("\n") requirements = f.read().strip().split("\n")
if nvcc_cuda_version <= Version("11.8"): if nvcc_cuda_version <= Version("11.8"):
@ -449,6 +444,16 @@ def get_requirements() -> List[str]:
if requirements[i].startswith("cupy-cuda12x"): if requirements[i].startswith("cupy-cuda12x"):
requirements[i] = "cupy-cuda11x" requirements[i] = "cupy-cuda11x"
break break
elif _is_hip():
with open(get_path("requirements-rocm.txt")) as f:
requirements = f.read().strip().split("\n")
elif _is_neuron():
with open(get_path("requirements-neuron.txt")) as f:
requirements = f.read().strip().split("\n")
else:
raise ValueError(
"Unsupported platform, please use CUDA, ROCM or Neuron.")
return requirements return requirements