evictor.py 4.42 KB
Newer Older
1
import enum
2
from abc import ABC, abstractmethod
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
from typing import OrderedDict, Tuple


class EvictionPolicy(enum.Enum):
    """Enum for eviction policy used by make_evictor to instantiate the correct
       Evictor subclass.
    """
    LRU = enum.auto()


class Evictor(ABC):
    """The Evictor subclasses should be used by the BlockAllocator class to
    handle eviction of freed PhysicalTokenBlocks.
    """

    @abstractmethod
    def __init__(self):
        pass

    @abstractmethod
    def __contains__(self, block_id: int) -> bool:
        pass

    @abstractmethod
    def evict(self) -> Tuple[int, int]:
        """Runs the eviction algorithm and returns the evicted block's
        content hash along with physical block id along with physical block id
        """
        pass

    @abstractmethod
    def add(self, block_id: int, content_hash: int, num_hashed_tokens: int,
35
            last_accessed: float):
36
37
38
39
        """Adds block to the evictor, making it a candidate for eviction"""
        pass

    @abstractmethod
40
    def update(self, block_id: int, last_accessed: float):
41
42
43
        """Update corresponding block's access time in metadata"""
        pass

44
45
46
47
48
    @abstractmethod
    def remove(self, block_id: int):
        """Remove a given block id from the cache."""
        pass

49
50
    @property
    @abstractmethod
51
52
53
54
    def num_blocks(self) -> int:
        pass


55
class BlockMetaData:
56
57
58
59
60
61
62
63
    """Data structure for storing key data describe cached block, so that
    evitor could use to make its decision which one to choose for eviction

    Here we use physical block id as the dict key, as there maybe several
    blocks with the same content hash, but their physical id is unique.
    """

    def __init__(self, content_hash: int, num_hashed_tokens: int,
64
                 last_accessed: float):
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
        self.content_hash = content_hash
        self.num_hashed_tokens = num_hashed_tokens
        self.last_accessed = last_accessed


class LRUEvictor(Evictor):
    """Evicts in a least-recently-used order using the last_accessed timestamp
    that's recorded in the PhysicalTokenBlock. If there are multiple blocks with
    the same last_accessed time, then the one with the largest num_hashed_tokens
    will be evicted. If two blocks each have the lowest last_accessed time and
    highest num_hashed_tokens value, then one will be chose arbitrarily
    """

    def __init__(self):
        self.free_table: OrderedDict[int, BlockMetaData] = OrderedDict()

    def __contains__(self, block_id: int) -> bool:
        return block_id in self.free_table

    def evict(self) -> Tuple[int, int]:
        if len(self.free_table) == 0:
            raise ValueError("No usable cache memory left")

88
        evicted_block, evicted_block_id = None, None
89
90
91
92
        # The blocks with the lowest timestamps should be placed consecutively
        # at the start of OrderedDict. Loop through all these blocks to
        # find the one with maximum number of hashed tokens.
        for _id, block in self.free_table.items():
93
94
95
            if evicted_block is None:
                evicted_block, evicted_block_id = block, _id
                continue
96
97
            if evicted_block.last_accessed < block.last_accessed:
                break
98
99
            if evicted_block.num_hashed_tokens < block.num_hashed_tokens:
                evicted_block, evicted_block_id = block, _id
100

101
102
        assert evicted_block is not None
        assert evicted_block_id is not None
103
104
105
106
107
        self.free_table.pop(evicted_block_id)

        return evicted_block_id, evicted_block.content_hash

    def add(self, block_id: int, content_hash: int, num_hashed_tokens: int,
108
            last_accessed: float):
109
110
111
112
        self.free_table[block_id] = BlockMetaData(content_hash,
                                                  num_hashed_tokens,
                                                  last_accessed)

113
    def update(self, block_id: int, last_accessed: float):
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
        self.free_table[block_id].last_accessed = last_accessed

    def remove(self, block_id: int):
        if block_id not in self.free_table:
            raise ValueError(
                "Attempting to remove block that's not in the evictor")
        self.free_table.pop(block_id)

    @property
    def num_blocks(self) -> int:
        return len(self.free_table)


def make_evictor(eviction_policy: EvictionPolicy) -> Evictor:
    if eviction_policy == EvictionPolicy.LRU:
        return LRUEvictor()
    else:
        raise ValueError(f"Unknown cache eviction policy: {eviction_policy}")