Skip to content

Commit 5e8400b

Browse files
authored
Merge pull request #44 from powerli2002/fix_cache_eviction
Fix cache eviction and soft deleted using MySQL and Mivuls
2 parents c0a45e6 + dd9093c commit 5e8400b

File tree

3 files changed

+55
-16
lines changed

3 files changed

+55
-16
lines changed

modelcache/manager/eviction_manager.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,10 @@ def check_evict(self):
2929
return True
3030
return False
3131

32-
def delete(self):
32+
def delete(self,model):
3333
mark_ids = self._scalar_storage.get_ids(deleted=True)
3434
self._scalar_storage.clear_deleted_data()
35-
self._vector_base.delete(mark_ids)
35+
self._vector_base.delete(mark_ids,model)
3636
self.delete_count += 1
3737
if self.delete_count >= self.REBUILD_CONDITION:
3838
self.rebuild()

modelcache/manager/scalar_data/sql_storage.py

+51-13
Original file line numberDiff line numberDiff line change
@@ -41,15 +41,15 @@ def _insert(self, data: List):
4141
model = data[3]
4242
answer_type = 0
4343
embedding_data = embedding_data.tobytes()
44+
is_deleted = 0
4445

45-
table_name = "cache_codegpt_answer"
46-
insert_sql = "INSERT INTO {} (question, answer, answer_type, model, embedding_data) VALUES (%s, %s, %s, %s, _binary%s)".format(table_name)
47-
46+
table_name = "modelcache_llm_answer"
47+
insert_sql = "INSERT INTO {} (question, answer, answer_type, model, embedding_data, is_deleted) VALUES (%s, %s, %s, %s, _binary%s, %s)".format(table_name)
4848
conn = self.pool.connection()
4949
try:
5050
with conn.cursor() as cursor:
5151
# 执行插入数据操作
52-
values = (question, answer, answer_type, model, embedding_data)
52+
values = (question, answer, answer_type, model, embedding_data, is_deleted)
5353
cursor.execute(insert_sql, values)
5454
conn.commit()
5555
id = cursor.lastrowid
@@ -91,7 +91,7 @@ def insert_query_resp(self, query_resp, **kwargs):
9191
conn.close()
9292

9393
def get_data_by_id(self, key: int):
94-
table_name = "cache_codegpt_answer"
94+
table_name = "modelcache_llm_answer"
9595
query_sql = "select question, answer, embedding_data, model from {} where id={}".format(table_name, key)
9696
conn_start = time.time()
9797
conn = self.pool.connection()
@@ -112,7 +112,7 @@ def get_data_by_id(self, key: int):
112112
return None
113113

114114
def update_hit_count_by_id(self, primary_id: int):
115-
table_name = "cache_codegpt_answer"
115+
table_name = "modelcache_llm_answer"
116116
update_sql = "UPDATE {} SET hit_count = hit_count+1 WHERE id={}".format(table_name, primary_id)
117117
conn = self.pool.connection()
118118

@@ -127,18 +127,30 @@ def update_hit_count_by_id(self, primary_id: int):
127127
conn.close()
128128

129129
def get_ids(self, deleted=True):
130-
pass
130+
table_name = "modelcache_llm_answer"
131+
state = 1 if deleted else 0
132+
query_sql = "Select id FROM {} WHERE is_deleted = {}".format(table_name, state)
133+
134+
conn = self.pool.connection()
135+
try:
136+
with conn.cursor() as cursor:
137+
cursor.execute(query_sql)
138+
ids = [row[0] for row in cursor.fetchall()]
139+
finally:
140+
conn.close()
141+
142+
return ids
131143

132144
def mark_deleted(self, keys):
133-
table_name = "cache_codegpt_answer"
134-
delete_sql = "Delete from {} WHERE id in ({})".format(table_name, ",".join([str(i) for i in keys]))
145+
table_name = "modelcache_llm_answer"
146+
mark_sql = " update {} set is_deleted=1 WHERE id in ({})".format(table_name, ",".join([str(i) for i in keys]))
135147

136148
# 从连接池中获取连接
137149
conn = self.pool.connection()
138150
try:
139151
with conn.cursor() as cursor:
140152
# 执行删除数据操作
141-
cursor.execute(delete_sql)
153+
cursor.execute(mark_sql)
142154
delete_count = cursor.rowcount
143155
conn.commit()
144156
finally:
@@ -147,7 +159,7 @@ def mark_deleted(self, keys):
147159
return delete_count
148160

149161
def model_deleted(self, model_name):
150-
table_name = "cache_codegpt_answer"
162+
table_name = "modelcache_llm_answer"
151163
delete_sql = "Delete from {} WHERE model='{}'".format(table_name, model_name)
152164

153165
table_log_name = "modelcache_query_log"
@@ -169,10 +181,36 @@ def model_deleted(self, model_name):
169181
return resp
170182

171183
def clear_deleted_data(self):
172-
pass
184+
table_name = "modelcache_llm_answer"
185+
delete_sql = "DELETE FROM {} WHERE is_deleted = 1".format(table_name)
186+
187+
conn = self.pool.connection()
188+
try:
189+
with conn.cursor() as cursor:
190+
cursor.execute(delete_sql)
191+
delete_count = cursor.rowcount
192+
conn.commit()
193+
finally:
194+
conn.close()
195+
196+
return delete_count
173197

174198
def count(self, state: int = 0, is_all: bool = False):
175-
pass
199+
table_name = "modelcache_llm_answer"
200+
if is_all:
201+
count_sql = "SELECT COUNT(*) FROM {}".format(table_name)
202+
else:
203+
count_sql = "SELECT COUNT(*) FROM {} WHERE is_deleted = {}".format(table_name,state)
204+
205+
conn = self.pool.connection()
206+
try:
207+
with conn.cursor() as cursor:
208+
cursor.execute(count_sql)
209+
num = cursor.fetchone()[0]
210+
finally:
211+
conn.close()
212+
213+
return num
176214

177215
def close(self):
178216
pass

reference_doc/create_table.sql

+2-1
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,9 @@ CREATE TABLE `modelcache_llm_answer` (
88
`hit_count` int(11) NOT NULL DEFAULT '0' comment 'hit_count',
99
`model` varchar(1000) NOT NULL comment 'model',
1010
`embedding_data` blob NOT NULL comment 'embedding_data',
11+
`is_deleted` tinyint(1) NOT NULL DEFAULT '0' COMMENT 'delete state(0 Not deleted,-1 deleted)',
1112
PRIMARY KEY(`id`)
12-
) AUTO_INCREMENT = 1 DEFAULT CHARSET = utf8mb4 COMMENT = 'modelcache_llm_answer';
13+
) AUTO_INCREMENT = 1 DEFAULT CHARSET = utf8mb4 COMMENT = 'cache_codegpt_answer';
1314

1415

1516
CREATE TABLE `modelcache_query_log` (

0 commit comments

Comments
 (0)