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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use std::any::Any;
use std::collections::HashMap;
use std::fmt::Debug;
use std::sync::{OnceLock, RwLock};
use thiserror::Error;
use tokio::sync::mpsc::{UnboundedSender, unbounded_channel};
use tokio::task;

// Common driver routines

static SERVICE_DISPATCHER: OnceLock<ServiceDispatcher> = OnceLock::new();

pub fn install_service<T>(name: &str, driver: T) -> Option<usize>
where
    T: Service + 'static + Send,
{
    // check if the dispatcher is initialized
    let dispatcher = SERVICE_DISPATCHER.get_or_init(|| ServiceDispatcher {
        maps: RwLock::new(HashMap::new()),
        channels: boxcar::Vec::new(),
    });

    dispatcher.install(name, driver)
}

pub fn dispatch<C>(service_id: usize, cmd: C) -> Result<(), ServiceError>
where
    C: Any + Send + Sync + Debug + 'static,
{
    SERVICE_DISPATCHER
        .get()
        .expect("Service is not initialized")
        .dispatch(service_id, cmd)
}

pub fn get_service_id(name: &str) -> Option<usize> {
    SERVICE_DISPATCHER
        .get()
        .expect("Service not found")
        .get_service_id(name)
}

#[derive(Debug, Error)]
pub enum ServiceError {
    #[error("Driver '{0}' not found")]
    DriverNotFound(String),

    #[error("Invalid driver index: {0}")]
    InvalidDriverIndex(usize),
}
//#[async_trait]
pub trait Service: Send {
    type Command: Send + Sync + 'static;

    fn handle(&mut self, cmd: Self::Command) -> impl Future<Output = ()> + Send;
}

pub type AnyCommand = Box<dyn Any + Send + Sync>;

#[derive(Debug)]
pub struct ServiceDispatcher {
    maps: RwLock<HashMap<String, usize>>,

    // Note: Using `boxcar::Vec` for performance
    // boxcar::Vec<> took:  14.79875ms
    // RwLock<Vec<>> took: 54.038458ms
    channels: boxcar::Vec<UnboundedSender<AnyCommand>>,
}

impl ServiceDispatcher {
    pub fn get_service_id(&self, name: &str) -> Option<usize> {
        self.maps.read().unwrap().get(name).copied()
    }

    pub fn install<T>(&self, name: &str, mut driver: T) -> Option<usize>
    where
        T: Service + 'static + Send,
    {
        let (tx, mut rx) = unbounded_channel();

        // first, make sure the name is not already registered
        if self.get_service_id(name).is_some() {
            return None;
        }

        self.channels.push(tx);
        let service_id = self.channels.count() - 1;

        self.maps
            .write()
            .unwrap()
            .insert(name.to_string(), service_id);

        task::spawn(async move {
            while let Some(cmd) = rx.recv().await {
                driver.handle(*cmd.downcast().unwrap()).await;
            }
        });

        Some(service_id)
    }

    pub fn dispatch<C>(&self, service_id: usize, cmd: C) -> Result<(), ServiceError>
    where
        C: Any + Send + Sync + Debug + 'static,
    {
        let cmd = Box::new(cmd);

        self.channels
            .get(service_id)
            .ok_or(ServiceError::InvalidDriverIndex(service_id))?
            .send(cmd)
            .unwrap();

        Ok(())
    }

    pub fn dispatch_with<C>(&self, name: &str, cmd: C) -> Result<(), ServiceError>
    where
        C: Any + Send + Sync + Debug + 'static,
    {
        let service_id = self
            .get_service_id(name)
            .ok_or(ServiceError::DriverNotFound(name.to_string()))?;

        self.dispatch(service_id, cmd)?;

        Ok(())
    }
}