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

use crate::manager::Manager;
use actix_web::http::header::ContentType;
use actix_web::http::StatusCode;
use actix_web::HttpResponse;
use actix_web::{middleware, web, App, HttpServer};
use anyhow::anyhow;
use config::Config;
use std::collections::HashMap;
use std::io;
use std::sync::RwLock;

pub mod routes;

pub struct Server;

impl Server {
    pub async fn new(model: &str, config: Config) -> io::Result<()> {
        let manager = web::Data::new(RwLock::new(
            Manager::new(model, config.clone()).await.unwrap(),
        ));

        // Start the HTTP server
        let cfg = config.clone();
        HttpServer::new(move || {
            App::new()
                .app_data(manager.clone())
                .app_data(web::Data::new(cfg.clone()))
                .wrap(middleware::Logger::default())
                .service(routes::inference)
                .service(routes::worker_status)
                .service(routes::all_workers)
                .service(routes::worker_info)
        })
        .bind(format!(
            "0.0.0.0:{}",
            config.clone().get_int("http_server.port").unwrap()
        ))?
        .run()
        .await
    }
}

#[derive(Debug)]
pub struct WebError {
    err: anyhow::Error,
}

impl std::fmt::Display for WebError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.err)
    }
}

impl actix_web::error::ResponseError for WebError {
    fn error_response(&self) -> HttpResponse {
        let err = HashMap::from([("errors", vec![self.to_string()])]);

        HttpResponse::build(self.status_code())
            .insert_header(ContentType::json())
            .json(err)
    }

    fn status_code(&self) -> StatusCode {
        StatusCode::INTERNAL_SERVER_ERROR
    }
}

impl From<anyhow::Error> for WebError {
    fn from(err: anyhow::Error) -> WebError {
        WebError { err }
    }
}

impl From<config::ConfigError> for WebError {
    fn from(err: config::ConfigError) -> WebError {
        WebError { err: anyhow!(err) }
    }
}

impl From<base64::DecodeError> for WebError {
    fn from(err: base64::DecodeError) -> Self {
        WebError { err: anyhow!(err) }
    }
}