Spaces:
Runtime error
Runtime error
File size: 2,364 Bytes
6c46b3f | 1 2 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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | use std::collections::BinaryHeap;
use std::sync::{Arc, Mutex};
use std::time::Instant;
use llama_cpp_2::token::LlamaToken;
#[derive(Clone)]
pub struct ChatRequest {
pub task_id: String,
pub prompt: String,
pub device_name: String,
pub tokens: Vec<LlamaToken>,
}
pub struct QueueItem {
pub request: ChatRequest,
pub usage_count: u64,
pub timestamp: Instant,
pub response_tx: tokio::sync::mpsc::Sender<Result<String, String>>,
}
// We want a Min-Heap based on usage_count, and then FIFO for timestamp.
// BinaryHeap is a Max-Heap, so we reverse the ordering.
impl PartialEq for QueueItem {
fn eq(&self, other: &Self) -> bool {
self.usage_count == other.usage_count && self.timestamp == other.timestamp
}
}
impl Eq for QueueItem {}
impl PartialOrd for QueueItem {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for QueueItem {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
// Lower usage_count = higher priority (greater in Max-Heap)
match other.usage_count.cmp(&self.usage_count) {
std::cmp::Ordering::Equal => {
// Older timestamp = higher priority
other.timestamp.cmp(&self.timestamp)
}
ord => ord,
}
}
}
pub struct QueueManager {
heap: Arc<Mutex<BinaryHeap<QueueItem>>>,
notifier: Arc<tokio::sync::Notify>,
}
impl QueueManager {
pub fn new() -> Self {
Self {
heap: Arc::new(Mutex::new(BinaryHeap::new())),
notifier: Arc::new(tokio::sync::Notify::new()),
}
}
pub fn enqueue(&self, item: QueueItem) {
self.heap.lock().unwrap().push(item);
self.notifier.notify_one();
}
pub fn get_position_by_device(&self, device_name: &str) -> Option<usize> {
let heap = self.heap.lock().unwrap();
let mut items: Vec<_> = heap.iter().collect();
items.sort();
let position = items.iter().rev().position(|item| item.request.device_name == device_name);
position.map(|p| p + 1)
}
pub async fn wait_for_next(&self) -> QueueItem {
loop {
if let Some(item) = self.heap.lock().unwrap().pop() {
return item;
}
self.notifier.notified().await;
}
}
}
|