perf: waiting_queue 改用 deque,pull_candidates 从 O(n²) 降到 O(1)
- list.pop(0) 每次左移全部元素,改 deque.popleft() 指针操作 - return_to_waiting 从 slice 整体复制改 appendleft 逐个插入 - 热路径 refill 阶段不再卡顿
This commit is contained in:
parent
e3382f6bb5
commit
513f1f7826
|
|
@ -2,8 +2,9 @@ import logging
|
|||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections import deque
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
from typing import Any, Callable, Deque, Dict, List, Optional
|
||||
|
||||
from astrai.tokenize.tokenizer import AutoTokenizer
|
||||
|
||||
|
|
@ -76,7 +77,7 @@ class TaskManager:
|
|||
self.max_seq_len = max_seq_len
|
||||
self.max_prompt_len = max_prompt_len
|
||||
|
||||
self.waiting_queue: List[Task] = []
|
||||
self.waiting_queue: Deque[Task] = deque()
|
||||
self.active_tasks: List[Task] = []
|
||||
|
||||
self._task_event = threading.Event()
|
||||
|
|
@ -129,7 +130,9 @@ class TaskManager:
|
|||
def remove_task(self, task_id: str) -> List[Task]:
|
||||
with self._lock:
|
||||
removed_active = [t for t in self.active_tasks if t.task_id == task_id]
|
||||
self.waiting_queue = [t for t in self.waiting_queue if t.task_id != task_id]
|
||||
self.waiting_queue = deque(
|
||||
t for t in self.waiting_queue if t.task_id != task_id
|
||||
)
|
||||
self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
|
||||
return removed_active
|
||||
|
||||
|
|
@ -166,7 +169,7 @@ class TaskManager:
|
|||
with self._lock:
|
||||
take = min(n, len(self.waiting_queue))
|
||||
for _ in range(take):
|
||||
to_add.append(self.waiting_queue.pop(0))
|
||||
to_add.append(self.waiting_queue.popleft())
|
||||
return to_add
|
||||
|
||||
def activate(self, task: Task) -> None:
|
||||
|
|
@ -176,7 +179,8 @@ class TaskManager:
|
|||
|
||||
def return_to_waiting(self, tasks: List[Task]) -> None:
|
||||
with self._lock:
|
||||
self.waiting_queue[:0] = tasks
|
||||
for task in reversed(tasks):
|
||||
self.waiting_queue.appendleft(task)
|
||||
|
||||
def has_work(self) -> bool:
|
||||
return bool(self.active_tasks or self.waiting_queue)
|
||||
|
|
|
|||
Loading…
Reference in New Issue