From ab68783757a7da393d57ab30f1261c5b073c9c5e Mon Sep 17 00:00:00 2001 From: simon-mo Date: Tue, 30 Jan 2024 20:34:13 +0000 Subject: [PATCH 1/2] Fix setup.py neuron-ls issue --- setup.py | 38 +++++++++++++++++--------------------- 1 file changed, 17 insertions(+), 21 deletions(-) diff --git a/setup.py b/setup.py index 3e2127855a75..08ccfe72f0da 100644 --- a/setup.py +++ b/setup.py @@ -2,6 +2,7 @@ import io import os import re +import shutil import subprocess import warnings from pathlib import Path @@ -23,21 +24,16 @@ # SUPPORTED_ARCHS = NVIDIA_SUPPORTED_ARCHS.union(ROCM_SUPPORTED_ARCHS) +def _is_cuda() -> bool: + return torch.version.cuda is not None + + def _is_hip() -> bool: return torch.version.hip is not None def _is_neuron() -> bool: - torch_neuronx_installed = True - try: - subprocess.run(["neuron-ls"], capture_output=True, check=True) - except FileNotFoundError: - torch_neuronx_installed = False - return torch_neuronx_installed - - -def _is_cuda() -> bool: - return (torch.version.cuda is not None) and not _is_neuron() + return shutil.which("neuron-ls") is not None # Compiler flags. @@ -317,7 +313,7 @@ def get_torch_arch_list() -> Set[str]: vllm_extension_sources.append("csrc/quantization/awq/gemm_kernels.cu") vllm_extension_sources.append("csrc/custom_all_reduce.cu") -if not _is_neuron(): +if _is_cuda() or _is_hip(): vllm_extension = CUDAExtension( name="vllm._C", sources=vllm_extension_sources, @@ -350,7 +346,12 @@ def find_version(filepath: str) -> str: def get_vllm_version() -> str: 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 hipcc_version = get_hipcc_rocm_version() if hipcc_version != MAIN_CUDA_VERSION: @@ -362,11 +363,6 @@ def get_vllm_version() -> str: if neuron_version != MAIN_CUDA_VERSION: neuron_version_str = neuron_version.replace(".", "")[:3] version += f"+neuron{neuron_version_str}" - else: - 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}" return version @@ -382,15 +378,15 @@ def read_readme() -> str: def get_requirements() -> List[str]: """Get Python package dependencies from requirements.txt.""" - if _is_hip(): + if _is_cuda(): + with open(get_path("requirements.txt")) as f: + requirements = f.read().strip().split("\n") + 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: - with open(get_path("requirements.txt")) as f: - requirements = f.read().strip().split("\n") return requirements From 9fdb20aa587da7e8006965b6d80c2a9a6e16cd2d Mon Sep 17 00:00:00 2001 From: simon-mo Date: Fri, 15 Mar 2024 22:55:05 -0700 Subject: [PATCH 2/2] address comments --- setup.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index f5a1a2014c54..19d8c28a6adb 100644 --- a/setup.py +++ b/setup.py @@ -38,7 +38,12 @@ def _is_hip() -> bool: def _is_neuron() -> bool: - return shutil.which("neuron-ls") is not None + torch_neuronx_installed = True + try: + subprocess.run(["neuron-ls"], capture_output=True, check=True) + except subprocess.CalledProcessError: + torch_neuronx_installed = False + return torch_neuronx_installed # Compiler flags. @@ -425,6 +430,9 @@ def get_requirements() -> List[str]: 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