Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "891b9d33de7ee7b3ee95b9bd7bb8a9cffae0e08c"
Unverified Commit b37cdce2 authored by Antoni Baum's avatar Antoni Baum Committed by GitHub
Browse files

[Core] Cache some utils (#3474)

parent b30880a7
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 0 deletions
vllm/utils.py
View file @ b37cdce2
...@@ -4,6 +4,7 @@ import socket ...@@ -4,6 +4,7 @@ import socket
import subprocess import subprocess
import uuid import uuid
import gc import gc
from functools import cache
from platform import uname from platform import uname
from typing import List, Tuple, Union from typing import List, Tuple, Union
from packaging.version import parse, Version from packaging.version import parse, Version
...@@ -120,6 +121,7 @@ def is_hip() -> bool: ...@@ -120,6 +121,7 @@ def is_hip() -> bool:
return torch.version.hip is not None return torch.version.hip is not None
@cache
def is_neuron() -> bool: def is_neuron() -> bool:
try: try:
import transformers_neuronx import transformers_neuronx
...@@ -128,6 +130,7 @@ def is_neuron() -> bool: ...@@ -128,6 +130,7 @@ def is_neuron() -> bool:
return transformers_neuronx is not None return transformers_neuronx is not None
@cache
def get_max_shared_memory_bytes(gpu: int = 0) -> int: def get_max_shared_memory_bytes(gpu: int = 0) -> int:
"""Returns the maximum shared memory per thread block in bytes.""" """Returns the maximum shared memory per thread block in bytes."""
# NOTE: This import statement should be executed lazily since # NOTE: This import statement should be executed lazily since
...@@ -151,6 +154,7 @@ def random_uuid() -> str: ...@@ -151,6 +154,7 @@ def random_uuid() -> str:
return str(uuid.uuid4().hex) return str(uuid.uuid4().hex)
@cache
def in_wsl() -> bool: def in_wsl() -> bool:
# Reference: https://github.com/microsoft/WSL/issues/4071 # Reference: https://github.com/microsoft/WSL/issues/4071
return "microsoft" in " ".join(uname()).lower() return "microsoft" in " ".join(uname()).lower()
...@@ -225,6 +229,7 @@ def set_cuda_visible_devices(device_ids: List[int]) -> None: ...@@ -225,6 +229,7 @@ def set_cuda_visible_devices(device_ids: List[int]) -> None:
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, device_ids)) os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, device_ids))
@cache
def get_nvcc_cuda_version() -> Optional[Version]: def get_nvcc_cuda_version() -> Optional[Version]:
cuda_home = os.environ.get('CUDA_HOME') cuda_home = os.environ.get('CUDA_HOME')
if not cuda_home: if not cuda_home:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment