Skip to content

Commit cf8cac8

Browse files
rkooo567cadedaniel
andauthored
[mypy][6/N] Fix all the core subdirectory typing (#4450)
Co-authored-by: Cade Daniel <[email protected]>
1 parent 5e401bc commit cf8cac8

File tree

10 files changed

+275
-83
lines changed

10 files changed

+275
-83
lines changed

.github/workflows/mypy.yaml

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ jobs:
3333
- name: Mypy
3434
run: |
3535
mypy vllm/attention --config-file pyproject.toml
36+
mypy vllm/core --config-file pyproject.toml
3637
mypy vllm/distributed --config-file pyproject.toml
3738
mypy vllm/entrypoints --config-file pyproject.toml
3839
mypy vllm/executor --config-file pyproject.toml
@@ -42,9 +43,6 @@ jobs:
4243
mypy vllm/engine --config-file pyproject.toml
4344
mypy vllm/worker --config-file pyproject.toml
4445
mypy vllm/spec_decode --config-file pyproject.toml
45-
mypy vllm/lora --config-file pyproject.toml
4646
mypy vllm/model_executor --config-file pyproject.toml
47-
48-
# TODO(sang): Fix nested dir
49-
mypy vllm/core/*.py --follow-imports=skip --config-file pyproject.toml
47+
mypy vllm/lora --config-file pyproject.toml
5048

format.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ echo 'vLLM yapf: Done'
9595
# Run mypy
9696
echo 'vLLM mypy:'
9797
mypy vllm/attention --config-file pyproject.toml
98-
mypy vllm/core/*.py --follow-imports=skip --config-file pyproject.toml
98+
mypy vllm/core --config-file pyproject.toml
9999
mypy vllm/distributed --config-file pyproject.toml
100100
mypy vllm/entrypoints --config-file pyproject.toml
101101
mypy vllm/executor --config-file pyproject.toml

vllm/core/block/block_table.py

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,9 @@ def __init__(
4040
):
4141
self._block_size = block_size
4242
self._allocator = block_allocator
43-
self._blocks: Optional[List[Block]] = _blocks
43+
if _blocks is None:
44+
_blocks = []
45+
self._blocks: List[Block] = _blocks
4446

4547
# Use helper method instead of directly calculating, as blocks
4648
# may not be allocated.
@@ -104,7 +106,7 @@ def append_token_ids(self,
104106
token_ids (List[int]): The sequence of token IDs to be appended.
105107
"""
106108
assert self._is_allocated
107-
assert self._blocks is not None
109+
assert len(self._blocks) > 0
108110

109111
self.ensure_num_empty_slots(num_empty_slots=len(token_ids) +
110112
num_lookahead_slots)
@@ -141,6 +143,7 @@ def ensure_num_empty_slots(self, num_empty_slots: int) -> None:
141143
blocks_to_allocate = cdiv(slots_to_allocate, self._block_size)
142144

143145
for _ in range(blocks_to_allocate):
146+
assert len(self._blocks) > 0
144147
self._blocks.append(
145148
self._allocator.allocate_mutable(prev_block=self._blocks[-1],
146149
device=device))
@@ -159,6 +162,7 @@ def fork(self) -> "BlockTable":
159162
the current instance.
160163
"""
161164
assert self._is_allocated
165+
assert len(self._blocks) > 0
162166
forked_blocks = self._allocator.fork(self._blocks[-1])
163167
return BlockTable(
164168
block_size=self._block_size,
@@ -177,10 +181,10 @@ def free(self) -> None:
177181
assert self._is_allocated
178182
for block in self._blocks:
179183
self._allocator.free(block)
180-
self._blocks = None
184+
self._blocks = []
181185

182186
@property
183-
def physical_block_ids(self) -> List[int]:
187+
def physical_block_ids(self) -> List[Optional[int]]:
184188
"""Returns a list of physical block indices for the blocks in the
185189
BlockTable.
186190
@@ -235,7 +239,7 @@ def _allocate_blocks_for_token_ids(self, prev_block: Optional[Block],
235239

236240
def _get_all_token_ids(self) -> List[int]:
237241
# NOTE: This function is O(seq_len); use sparingly.
238-
token_ids = []
242+
token_ids: List[int] = []
239243

240244
if not self._is_allocated:
241245
return token_ids
@@ -247,7 +251,7 @@ def _get_all_token_ids(self) -> List[int]:
247251

248252
@property
249253
def _is_allocated(self) -> bool:
250-
return self._blocks is not None
254+
return len(self._blocks) > 0
251255

252256
@property
253257
def _num_empty_slots(self) -> int:

vllm/core/block/common.py

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,25 @@
11
from collections import defaultdict
2-
from typing import Dict, Iterable, List, Optional
2+
from typing import Dict, Iterable, List, Optional, Protocol
33

44
from vllm.core.block.interfaces import Block, BlockAllocator
55

66
BlockId = int
77
RefCount = int
88

99

10-
class RefCounter:
10+
class RefCounterProtocol(Protocol):
11+
12+
def incr(self, block_id: BlockId) -> RefCount:
13+
raise NotImplementedError
14+
15+
def decr(self, block_id: BlockId) -> RefCount:
16+
raise NotImplementedError
17+
18+
def get(self, block_id: BlockId) -> RefCount:
19+
raise NotImplementedError
20+
21+
22+
class RefCounter(RefCounterProtocol):
1123
"""A class for managing reference counts for a set of block indices.
1224
1325
The RefCounter class maintains a dictionary that maps block indices to their
@@ -54,7 +66,7 @@ def as_readonly(self) -> "ReadOnlyRefCounter":
5466
return ReadOnlyRefCounter(self)
5567

5668

57-
class ReadOnlyRefCounter:
69+
class ReadOnlyRefCounter(RefCounterProtocol):
5870
"""A read-only view of the RefCounter class.
5971
6072
The ReadOnlyRefCounter class provides a read-only interface to access the
@@ -96,7 +108,7 @@ class CopyOnWriteTracker:
96108

97109
def __init__(
98110
self,
99-
refcounter: RefCounter,
111+
refcounter: RefCounterProtocol,
100112
allocator: BlockAllocator,
101113
):
102114
self._copy_on_writes: Dict[BlockId, List[BlockId]] = defaultdict(list)

vllm/core/block/cpu_gpu_block_allocator.py

Lines changed: 33 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
1-
from typing import Dict, List, Optional
1+
from typing import Dict, FrozenSet, List, Optional
22

3-
from vllm.core.block.interfaces import (Block, BlockAllocator,
3+
from vllm.core.block.interfaces import (Block, BlockAllocator, BlockId,
44
DeviceAwareBlockAllocator)
55
from vllm.core.block.naive_block import NaiveBlock, NaiveBlockAllocator
66
from vllm.core.block.prefix_caching_block import PrefixCachingBlockAllocator
@@ -57,15 +57,15 @@ def create(
5757
cpu_block_ids = block_ids[num_gpu_blocks:]
5858

5959
if allocator_type == "naive":
60-
gpu_allocator = NaiveBlockAllocator(
61-
create_block=NaiveBlock,
60+
gpu_allocator: BlockAllocator = NaiveBlockAllocator(
61+
create_block=NaiveBlock, # type: ignore
6262
num_blocks=num_gpu_blocks,
6363
block_size=block_size,
6464
block_ids=gpu_block_ids,
6565
)
6666

67-
cpu_allocator = NaiveBlockAllocator(
68-
create_block=NaiveBlock,
67+
cpu_allocator: BlockAllocator = NaiveBlockAllocator(
68+
create_block=NaiveBlock, # type: ignore
6969
num_blocks=num_cpu_blocks,
7070
block_size=block_size,
7171
block_ids=cpu_block_ids,
@@ -105,13 +105,14 @@ def __init__(
105105
Device.GPU: gpu_block_allocator,
106106
}
107107

108-
self._block_ids_to_allocator = {}
108+
self._block_ids_to_allocator: Dict[int, BlockAllocator] = {}
109109
for _, allocator in self._allocators.items():
110110
for block_id in allocator.all_block_ids:
111111
self._block_ids_to_allocator[block_id] = allocator
112112

113-
def allocate_mutable(self, prev_block: Optional[Block],
114-
device: Device) -> Block:
113+
def allocate_mutable(self,
114+
prev_block: Optional[Block],
115+
device: Optional[Device] = None) -> Block:
115116
"""Allocates a new mutable block on the specified device.
116117
117118
Args:
@@ -122,10 +123,13 @@ def allocate_mutable(self, prev_block: Optional[Block],
122123
Returns:
123124
Block: The newly allocated mutable block.
124125
"""
126+
assert device is not None
125127
return self._allocators[device].allocate_mutable(prev_block)
126128

127-
def allocate_immutable(self, prev_block: Optional[Block],
128-
token_ids: List[int], device: Device) -> Block:
129+
def allocate_immutable(self,
130+
prev_block: Optional[Block],
131+
token_ids: List[int],
132+
device: Optional[Device] = None) -> Block:
129133
"""Allocates a new immutable block with the provided token IDs on the
130134
specified device.
131135
@@ -140,6 +144,7 @@ def allocate_immutable(self, prev_block: Optional[Block],
140144
Block: The newly allocated immutable block containing the provided
141145
token IDs.
142146
"""
147+
assert device is not None
143148
return self._allocators[device].allocate_immutable(
144149
prev_block, token_ids)
145150

@@ -149,7 +154,9 @@ def free(self, block: Block) -> None:
149154
Args:
150155
block (Block): The block to be freed.
151156
"""
152-
allocator = self._block_ids_to_allocator[block.block_id]
157+
block_id = block.block_id
158+
assert block_id is not None
159+
allocator = self._block_ids_to_allocator[block_id]
153160
return allocator.free(block)
154161

155162
def fork(self, last_block: Block) -> List[Block]:
@@ -163,19 +170,22 @@ def fork(self, last_block: Block) -> List[Block]:
163170
List[Block]: A new list of blocks that shares the same memory as the
164171
original sequence.
165172
"""
166-
allocator = self._block_ids_to_allocator[last_block.block_id]
173+
block_id = last_block.block_id
174+
assert block_id is not None
175+
allocator = self._block_ids_to_allocator[block_id]
167176
return allocator.fork(last_block)
168177

169-
def get_num_free_blocks(self, device: Device) -> int:
178+
def get_num_free_blocks(self, device: Optional[Device] = None) -> int:
170179
"""Returns the number of free blocks available on the specified device.
171180
172181
Args:
173182
device (Device): The device for which to query the number of free
174-
blocks.
183+
blocks. AssertionError is raised if None is passed.
175184
176185
Returns:
177186
int: The number of free blocks available on the specified device.
178187
"""
188+
assert device is not None
179189
return self._allocators[device].get_num_free_blocks()
180190

181191
def clear_copy_on_writes(self) -> Dict[int, List[int]]:
@@ -210,5 +220,12 @@ def get_common_computed_block_ids(
210220
return self._allocators[device].get_common_computed_block_ids(
211221
seq_block_ids)
212222

213-
def all_block_ids(self) -> frozenset[int]:
223+
@property
224+
def all_block_ids(self) -> FrozenSet[int]:
214225
return frozenset(self._block_ids_to_allocator.keys())
226+
227+
def promote_to_immutable_block(self, block: Block) -> BlockId:
228+
raise NotImplementedError
229+
230+
def cow_block_if_not_appendable(self, block: Block) -> Optional[BlockId]:
231+
raise NotImplementedError

0 commit comments

Comments
 (0)