Skip to content

Commit

Permalink
[CI/Build] build on empty device for better dev experience (#4773)
Browse files Browse the repository at this point in the history
  • Loading branch information
tomeras91 authored Aug 11, 2024
1 parent c08e2b3 commit 3860879
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 7 deletions.
4 changes: 2 additions & 2 deletions requirements-cuda.txt
Original file line number Diff line number Diff line change
Expand Up @@ -7,5 +7,5 @@ nvidia-ml-py # for pynvml package
torch == 2.4.0
# These must be updated alongside torch
torchvision == 0.19 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
xformers == 0.0.27.post2 # Requires PyTorch 2.4.0
vllm-flash-attn == 2.6.1 # Requires PyTorch 2.4.0
xformers == 0.0.27.post2; platform_system == 'Linux' and platform_machine == 'x86_64' # Requires PyTorch 2.4.0
vllm-flash-attn == 2.6.1; platform_system == 'Linux' and platform_machine == 'x86_64' # Requires PyTorch 2.4.0
24 changes: 19 additions & 5 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,9 +61,12 @@ def embed_commit_hash():

VLLM_TARGET_DEVICE = envs.VLLM_TARGET_DEVICE

# vLLM only supports Linux platform
assert sys.platform.startswith(
"linux"), "vLLM only supports Linux platform (including WSL)."
if not sys.platform.startswith("linux"):
logger.warning(
"vLLM only supports Linux platform (including WSL). "
"Building on %s, "
"so vLLM may not be able to run correctly", sys.platform)
VLLM_TARGET_DEVICE = "empty"

MAIN_CUDA_VERSION = "12.1"

Expand Down Expand Up @@ -231,6 +234,10 @@ def build_extensions(self) -> None:
subprocess.check_call(["cmake", *build_args], cwd=self.build_temp)


def _no_device() -> bool:
return VLLM_TARGET_DEVICE == "empty"


def _is_cuda() -> bool:
has_cuda = torch.version.cuda is not None
return (VLLM_TARGET_DEVICE == "cuda" and has_cuda
Expand Down Expand Up @@ -350,7 +357,9 @@ def find_version(filepath: str) -> str:
def get_vllm_version() -> str:
version = find_version(get_path("vllm", "version.py"))

if _is_cuda():
if _no_device():
version += "+empty"
elif _is_cuda():
cuda_version = str(get_nvcc_cuda_version())
if cuda_version != MAIN_CUDA_VERSION:
cuda_version_str = cuda_version.replace(".", "")[:3]
Expand Down Expand Up @@ -404,7 +413,9 @@ def _read_requirements(filename: str) -> List[str]:
resolved_requirements.append(line)
return resolved_requirements

if _is_cuda():
if _no_device():
requirements = _read_requirements("requirements-cuda.txt")
elif _is_cuda():
requirements = _read_requirements("requirements-cuda.txt")
cuda_major, cuda_minor = torch.version.cuda.split(".")
modified_requirements = []
Expand Down Expand Up @@ -453,6 +464,9 @@ def _read_requirements(filename: str) -> List[str]:
ext_modules = []
package_data["vllm"].append("*.so")

if _no_device():
ext_modules = []

setup(
name="vllm",
version=get_vllm_version(),
Expand Down

0 comments on commit 3860879

Please sign in to comment.