biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[mypy] Enable following imports for some directories (#6681)

Browse Source
This commit is contained in:
Cyrus Leung
2024-07-31 10:38:03 +08:00
committed by GitHub
parent c32ab8be1a
commit da1f7cc12a
18 changed files with 185 additions and 143 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
vllm/utils.py
View File

@@ -94,8 +94,10 @@ class LRUCache(Generic[T]):
def __len__(self) -> int:
return len(self.cache)
def __getitem__(self, key: Hashable) -> Optional[T]:
return self.get(key)
def __getitem__(self, key: Hashable) -> T:
value = self.cache[key] # Raise KeyError if not exists
self.cache.move_to_end(key)
return value
def __setitem__(self, key: Hashable, value: T) -> None:
self.put(key, value)
@@ -109,8 +111,9 @@ class LRUCache(Generic[T]):
def get(self,
key: Hashable,
default_value: Optional[T] = None) -> Optional[T]:
value: Optional[T]
if key in self.cache:
value: Optional[T] = self.cache[key]
value = self.cache[key]
self.cache.move_to_end(key)
else:
value = default_value
@@ -590,8 +593,8 @@ class CudaMemoryProfiler:
torch.cuda.reset_peak_memory_stats(self.device)
mem = torch.cuda.max_memory_allocated(self.device)
elif is_xpu():
torch.xpu.reset_peak_memory_stats(self.device)
mem = torch.xpu.max_memory_allocated(self.device)
torch.xpu.reset_peak_memory_stats(self.device) # type: ignore
mem = torch.xpu.max_memory_allocated(self.device) # type: ignore
return mem
def __enter__(self):