pie 0.2.0

Pie: A Programmable LLM Serving System
Documentation 
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
86
87
88
89
90
91
92

use crate::service::{self, Service, ServiceError};
use dashmap::DashMap;
use std::sync::{Arc, OnceLock};
use tokio::sync::oneshot;

static SERVICE_ID_KVS: OnceLock<usize> = OnceLock::new();

/// Dispatches a command to the key-value store service.
pub fn dispatch(command: Command) -> Result<(), ServiceError> {
    let service_id = *SERVICE_ID_KVS.get_or_init(|| service::get_service_id("kvs").unwrap());
    service::dispatch(service_id, command)
}

/// Defines the set of operations available for the key-value store.
#[derive(Debug)]
pub enum Command {
    /// Retrieves a value associated with a key.
    /// The result is sent back as an `Option<String>`.
    Get {
        key: String,
        response: oneshot::Sender<Option<String>>,
    },
    /// Inserts or updates a key-value pair.
    /// The `oneshot::Sender` is used to signal completion.
    Set { key: String, value: String },
    /// Removes a key-value pair.
    /// The `oneshot::Sender` is used to signal completion.
    Delete { key: String },
    /// Checks if a key exists in the store.
    /// The result is sent back as a `bool`.
    Exists {
        key: String,
        response: oneshot::Sender<bool>,
    },
    /// Retrieves a list of all keys currently in the store.
    /// The result is sent back as a `Vec<String>`.
    ListKeys {
        response: oneshot::Sender<Vec<String>>,
    },
}

/// An in-memory key-value store service.
///
/// It uses a `DashMap` for concurrent, lock-free reads and writes,
/// making it suitable for a multi-threaded, asynchronous environment.
#[derive(Debug, Clone)]
pub struct KeyValueStore {
    store: Arc<DashMap<String, String>>,
}

impl KeyValueStore {
    /// Creates a new, empty `KeyValueStore`.
    pub fn new() -> Self {
        KeyValueStore {
            store: Arc::new(DashMap::new()),
        }
    }
}

impl Default for KeyValueStore {
    fn default() -> Self {
        Self::new()
    }
}

impl Service for KeyValueStore {
    type Command = Command;

    async fn handle(&mut self, cmd: Self::Command) {
        match cmd {
            Command::Get { key, response } => {
                let value = self.store.get(&key).map(|v| v.value().clone());
                let _ = response.send(value);
            }
            Command::Set { key, value } => {
                self.store.insert(key, value);
            }
            Command::Delete { key } => {
                self.store.remove(&key);
            }
            Command::Exists { key, response } => {
                let exists = self.store.contains_key(&key);
                let _ = response.send(exists);
            }
            Command::ListKeys { response } => {
                let keys: Vec<String> =
                    self.store.iter().map(|entry| entry.key().clone()).collect();
                let _ = response.send(keys);
            }
        }
    }
}