Skip to content

Commit

Permalink
formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
@Bellk17
Bellk17 committed May 10, 2024
1 parent 7f1df86 commit aba83e3
Showing 1 changed file with 29 additions and 34 deletions.
63 changes: 29 additions & 34 deletions setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,7 @@ def load_module_from_path(module_name, path):

# vLLM only supports Linux platform
assert sys.platform.startswith(
"linux"
), "vLLM only supports Linux platform (including WSL)."
"linux"), "vLLM only supports Linux platform (including WSL)."

MAIN_CUDA_VERSION = "12.1"

Expand All @@ -55,7 +54,7 @@ def is_ninja_available() -> bool:

def remove_prefix(text, prefix):
if text.startswith(prefix):
return text[len(prefix) :]
return text[len(prefix):]
return text


Expand Down Expand Up @@ -98,8 +97,8 @@ def compute_num_jobs(self):
if nvcc_threads is not None:
nvcc_threads = int(nvcc_threads)
logger.info(
"Using NVCC_THREADS=%d as the number of nvcc threads.", nvcc_threads
)
"Using NVCC_THREADS=%d as the number of nvcc threads.",
nvcc_threads)
else:
nvcc_threads = 1
num_jobs = max(1, num_jobs // nvcc_threads)
Expand All @@ -124,7 +123,8 @@ def configure(self, ext: CMakeExtension) -> None:

# where .so files will be written, should be the same for all extensions
# that use the same CMakeLists.txt.
outdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))
outdir = os.path.abspath(
os.path.dirname(self.get_ext_fullpath(ext.name)))

cmake_args = [
"-DCMAKE_BUILD_TYPE={}".format(cfg),
Expand Down Expand Up @@ -210,17 +210,13 @@ def build_extensions(self) -> None:


def _is_cuda() -> bool:
return (
VLLM_TARGET_DEVICE == "cuda"
and torch.version.cuda is not None
and not _is_neuron()
)
return (VLLM_TARGET_DEVICE == "cuda" and torch.version.cuda is not None
and not _is_neuron())


def _is_hip() -> bool:
return (
VLLM_TARGET_DEVICE == "cuda" or VLLM_TARGET_DEVICE == "rocm"
) and torch.version.hip is not None
return (VLLM_TARGET_DEVICE == "cuda"
or VLLM_TARGET_DEVICE == "rocm") and torch.version.hip is not None


def _is_neuron() -> bool:
Expand Down Expand Up @@ -268,7 +264,8 @@ def get_neuronxcc_version():
import sysconfig

site_dir = sysconfig.get_paths()["purelib"]
version_file = os.path.join(site_dir, "neuronxcc", "version", "__init__.py")
version_file = os.path.join(site_dir, "neuronxcc", "version",
"__init__.py")

# Check if the command was executed successfully
with open(version_file, "rt") as fp:
Expand All @@ -289,9 +286,8 @@ def get_nvcc_cuda_version() -> Version:
Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py
"""
assert CUDA_HOME is not None, "CUDA_HOME is not set"
nvcc_output = subprocess.check_output(
[CUDA_HOME + "/bin/nvcc", "-V"], universal_newlines=True
)
nvcc_output = subprocess.check_output([CUDA_HOME + "/bin/nvcc", "-V"],
universal_newlines=True)
output = nvcc_output.split()
release_idx = output.index("release") + 1
nvcc_cuda_version = parse(output[release_idx].split(",")[0])
Expand All @@ -308,9 +304,8 @@ def find_version(filepath: str) -> str:
Adapted from https://github.com/ray-project/ray/blob/0b190ee1160eeca9796bc091e07eaebf4c85b511/python/setup.py
"""
with open(filepath) as fp:
version_match = re.search(
r"^__version__ = ['\"]([^'\"]*)['\"]", fp.read(), re.M
)
version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]",
fp.read(), re.M)
if version_match:
return version_match.group(1)
raise RuntimeError("Unable to find version string.")
Expand Down Expand Up @@ -373,10 +368,10 @@ def _read_requirements(filename: str) -> List[str]:
modified_requirements = []
for req in requirements:
if "vllm-nccl-cu12" in req:
req = req.replace("vllm-nccl-cu12", f"vllm-nccl-cu{cuda_major}")
elif "vllm-flash-attn" in req and not (
cuda_major == "12" and cuda_minor == "1"
):
req = req.replace("vllm-nccl-cu12",
f"vllm-nccl-cu{cuda_major}")
elif "vllm-flash-attn" in req and not (cuda_major == "12"
and cuda_minor == "1"):
# vllm-flash-attn is built only for CUDA 12.1.
# Skip for other versions.
continue
Expand All @@ -389,7 +384,8 @@ def _read_requirements(filename: str) -> List[str]:
elif _is_cpu():
requirements = _read_requirements("requirements-cpu.txt")
else:
raise ValueError("Unsupported platform, please use CUDA, ROCm, Neuron, or CPU.")
raise ValueError(
"Unsupported platform, please use CUDA, ROCm, Neuron, or CPU.")
return requirements


Expand All @@ -404,7 +400,9 @@ def _read_requirements(filename: str) -> List[str]:
if _install_punica():
ext_modules.append(CMakeExtension(name="vllm._punica_C"))

package_data = {"vllm": ["py.typed", "model_executor/layers/fused_moe/configs/*.json"]}
package_data = {
"vllm": ["py.typed", "model_executor/layers/fused_moe/configs/*.json"]
}
if envs.VLLM_USE_PRECOMPILED:
ext_modules = []
package_data["vllm"].append("*.so")
Expand All @@ -414,10 +412,8 @@ def _read_requirements(filename: str) -> List[str]:
version=get_vllm_version(),
author="vLLM Team",
license="Apache 2.0",
description=(
"A high-throughput and memory-efficient inference and "
"serving engine for LLMs"
),
description=("A high-throughput and memory-efficient inference and "
"serving engine for LLMs"),
long_description=read_readme(),
long_description_content_type="text/markdown",
url="https://github.com/vllm-project/vllm",
Expand All @@ -433,9 +429,8 @@ def _read_requirements(filename: str) -> List[str]:
"License :: OSI Approved :: Apache Software License",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
],
packages=find_packages(
exclude=("benchmarks", "csrc", "docs", "examples", "tests*")
),
packages=find_packages(exclude=("benchmarks", "csrc", "docs", "examples",
"tests*")),
python_requires=">=3.8",
install_requires=get_requirements(),
ext_modules=ext_modules,
Expand Down

0 comments on commit aba83e3

Please sign in to comment.