server.rs 5.62 KB
Newer Older
1
use crate::router::PolicyConfig;
2
use crate::router::Router;
3
use actix_web::{get, post, web, App, HttpRequest, HttpResponse, HttpServer, Responder};
4
use bytes::Bytes;
5
use env_logger::Builder;
Byron Hsu's avatar
Byron Hsu committed
6
use log::{info, LevelFilter};
7
use std::collections::HashMap;
8
use std::io::Write;
9
10
11

#[derive(Debug)]
pub struct AppState {
12
    router: Router,
13
14
15
    client: reqwest::Client,
}

16
impl AppState {
17
18
19
20
    pub fn new(
        worker_urls: Vec<String>,
        client: reqwest::Client,
        policy_config: PolicyConfig,
21
    ) -> Result<Self, String> {
22
        // Create router based on policy
23
24
        let router = Router::new(worker_urls, policy_config)?;
        Ok(Self { router, client })
25
26
27
    }
}

28
#[get("/health")]
29
30
31
32
async fn health(req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
    data.router
        .route_to_first(&data.client, "/health", &req)
        .await
33
34
35
}

#[get("/health_generate")]
36
async fn health_generate(req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
37
    data.router
38
        .route_to_first(&data.client, "/health_generate", &req)
39
        .await
40
41
}

42
#[get("/get_server_info")]
43
async fn get_server_info(req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
44
    data.router
45
        .route_to_first(&data.client, "/get_server_info", &req)
46
        .await
47
48
}

49
#[get("/v1/models")]
50
51
52
53
async fn v1_models(req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
    data.router
        .route_to_first(&data.client, "/v1/models", &req)
        .await
54
55
}

56
#[get("/get_model_info")]
57
async fn get_model_info(req: HttpRequest, data: web::Data<AppState>) -> impl Responder {
58
    data.router
59
        .route_to_first(&data.client, "/get_model_info", &req)
60
        .await
61
}
62

63
64
#[post("/generate")]
async fn generate(req: HttpRequest, body: Bytes, data: web::Data<AppState>) -> impl Responder {
65
    data.router
66
        .route_generate_request(&data.client, &req, &body, "/generate")
67
68
69
70
71
72
73
74
75
76
        .await
}

#[post("/v1/chat/completions")]
async fn v1_chat_completions(
    req: HttpRequest,
    body: Bytes,
    data: web::Data<AppState>,
) -> impl Responder {
    data.router
77
        .route_generate_request(&data.client, &req, &body, "/v1/chat/completions")
78
79
80
81
82
83
84
85
86
87
        .await
}

#[post("/v1/completions")]
async fn v1_completions(
    req: HttpRequest,
    body: Bytes,
    data: web::Data<AppState>,
) -> impl Responder {
    data.router
88
        .route_generate_request(&data.client, &req, &body, "/v1/completions")
89
        .await
90
91
}

92
93
94
95
96
97
98
99
100
101
102
103
#[post("/add_worker")]
async fn add_worker(
    query: web::Query<HashMap<String, String>>,
    data: web::Data<AppState>,
) -> impl Responder {
    let worker_url = match query.get("url") {
        Some(url) => url.to_string(),
        None => {
            return HttpResponse::BadRequest()
                .body("Worker URL required. Provide 'url' query parameter")
        }
    };
104

105
    match data.router.add_worker(&worker_url).await {
106
107
108
        Ok(message) => HttpResponse::Ok().body(message),
        Err(error) => HttpResponse::BadRequest().body(error),
    }
109
110
}

111
112
113
114
115
116
117
118
119
#[post("/remove_worker")]
async fn remove_worker(
    query: web::Query<HashMap<String, String>>,
    data: web::Data<AppState>,
) -> impl Responder {
    let worker_url = match query.get("url") {
        Some(url) => url.to_string(),
        None => return HttpResponse::BadRequest().finish(),
    };
120
    data.router.remove_worker(&worker_url);
121
    HttpResponse::Ok().body(format!("Successfully removed worker: {}", worker_url))
122
123
}

124
125
126
127
128
129
pub struct ServerConfig {
    pub host: String,
    pub port: u16,
    pub worker_urls: Vec<String>,
    pub policy_config: PolicyConfig,
    pub verbose: bool,
130
    pub max_payload_size: usize,
131
132
133
}

pub async fn startup(config: ServerConfig) -> std::io::Result<()> {
134
    // Initialize logger
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
    Builder::new()
        .format(|buf, record| {
            use chrono::Local;
            writeln!(
                buf,
                "[Router (Rust)] {} - {} - {}",
                Local::now().format("%Y-%m-%d %H:%M:%S"),
                record.level(),
                record.args()
            )
        })
        .filter(
            None,
            if config.verbose {
                LevelFilter::Debug
            } else {
                LevelFilter::Info
            },
        )
        .init();

156
157
158
159
160
161
162
163
    info!("🚧 Initializing router on {}:{}", config.host, config.port);
    info!("🚧 Initializing workers on {:?}", config.worker_urls);
    info!("🚧 Policy Config: {:?}", config.policy_config);
    info!(
        "🚧 Max payload size: {} MB",
        config.max_payload_size / (1024 * 1024)
    );

164
165
166
167
    let client = reqwest::Client::builder()
        .build()
        .expect("Failed to create HTTP client");

168
169
170
171
172
173
174
    let app_state = web::Data::new(
        AppState::new(
            config.worker_urls.clone(),
            client,
            config.policy_config.clone(),
        )
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?,
175
    );
176

177
178
179
    info!("✅ Serving router on {}:{}", config.host, config.port);
    info!("✅ Serving workers on {:?}", config.worker_urls);

180
181
182
    HttpServer::new(move || {
        App::new()
            .app_data(app_state.clone())
183
184
            .app_data(web::JsonConfig::default().limit(config.max_payload_size))
            .app_data(web::PayloadConfig::default().limit(config.max_payload_size))
185
            .service(generate)
186
187
188
            .service(v1_chat_completions)
            .service(v1_completions)
            .service(v1_models)
189
            .service(get_model_info)
190
191
            .service(health)
            .service(health_generate)
192
            .service(get_server_info)
193
            .service(add_worker)
194
            .service(remove_worker)
195
    })
196
    .bind((config.host, config.port))?
197
198
    .run()
    .await
199
}