From 21ec60eddfc4d2948a771acf486691e04cafc5c1 Mon Sep 17 00:00:00 2001 From: Meng Zhang Date: Wed, 25 Oct 2023 15:05:23 -0700 Subject: [PATCH] feat: add server.completion_timeout to control timeout of /v1/completion (#637) * feat: add server.completion_timeout to control timeout of /v1/completion * Update config.rs --- crates/tabby-common/src/config.rs | 29 +++++++++++++++---- crates/tabby-scheduler/src/dataset.rs | 4 +-- crates/tabby-scheduler/src/repository.rs | 4 +-- .../tabby-scheduler/tests/integration_test.rs | 5 ++-- crates/tabby/src/serve/mod.rs | 10 ++++--- 5 files changed, 36 insertions(+), 16 deletions(-) diff --git a/crates/tabby-common/src/config.rs b/crates/tabby-common/src/config.rs index 21e23cc81..e344f5324 100644 --- a/crates/tabby-common/src/config.rs +++ b/crates/tabby-common/src/config.rs @@ -11,7 +11,10 @@ use crate::path::{config_file, repositories_dir}; #[derive(Serialize, Deserialize, Default)] pub struct Config { #[serde(default)] - pub repositories: Vec, + pub repositories: Vec, + + #[serde(default)] + pub server: ServerConfig, } impl Config { @@ -37,11 +40,11 @@ impl Config { } #[derive(Serialize, Deserialize)] -pub struct Repository { +pub struct RepositoryConfig { pub git_url: String, } -impl Repository { +impl RepositoryConfig { pub fn dir(&self) -> PathBuf { if self.is_local_dir() { let path = self.git_url.strip_prefix("file://").unwrap(); @@ -56,9 +59,23 @@ impl Repository { } } +#[derive(Serialize, Deserialize)] +pub struct ServerConfig { + /// The timeout in seconds for the /v1/completion api. + pub completion_timeout: u64, +} + +impl Default for ServerConfig { + fn default() -> Self { + Self { + completion_timeout: 30, + } + } +} + #[cfg(test)] mod tests { - use super::{Config, Repository}; + use super::{Config, RepositoryConfig}; #[test] fn it_parses_empty_config() { @@ -68,13 +85,13 @@ mod tests { #[test] fn it_parses_local_dir() { - let repo = Repository { + let repo = RepositoryConfig { git_url: "file:///home/user".to_owned(), }; assert!(repo.is_local_dir()); assert_eq!(repo.dir().display().to_string(), "/home/user"); - let repo = Repository { + let repo = RepositoryConfig { git_url: "https://github.com/TabbyML/tabby".to_owned(), }; assert!(!repo.is_local_dir()); diff --git a/crates/tabby-scheduler/src/dataset.rs b/crates/tabby-scheduler/src/dataset.rs index a224a6daa..a4852d35f 100644 --- a/crates/tabby-scheduler/src/dataset.rs +++ b/crates/tabby-scheduler/src/dataset.rs @@ -11,7 +11,7 @@ use ignore::{DirEntry, Walk}; use lazy_static::lazy_static; use serde_jsonlines::WriteExt; use tabby_common::{ - config::{Config, Repository}, + config::{Config, RepositoryConfig}, path::dataset_dir, SourceFile, }; @@ -22,7 +22,7 @@ trait RepositoryExt { fn create_dataset(&self, writer: &mut impl Write) -> Result<()>; } -impl RepositoryExt for Repository { +impl RepositoryExt for RepositoryConfig { fn create_dataset(&self, writer: &mut impl Write) -> Result<()> { let dir = self.dir(); diff --git a/crates/tabby-scheduler/src/repository.rs b/crates/tabby-scheduler/src/repository.rs index 6cf0fd572..ce4a4472d 100644 --- a/crates/tabby-scheduler/src/repository.rs +++ b/crates/tabby-scheduler/src/repository.rs @@ -1,7 +1,7 @@ use std::process::Command; use anyhow::{anyhow, Result}; -use tabby_common::config::{Config, Repository}; +use tabby_common::config::{Config, RepositoryConfig}; trait ConfigExt { fn sync_repositories(&self) -> Result<()>; @@ -27,7 +27,7 @@ trait RepositoryExt { fn sync(&self) -> Result<()>; } -impl RepositoryExt for Repository { +impl RepositoryExt for RepositoryConfig { fn sync(&self) -> Result<()> { let dir = self.dir(); let dir_string = dir.display().to_string(); diff --git a/crates/tabby-scheduler/tests/integration_test.rs b/crates/tabby-scheduler/tests/integration_test.rs index 54fe60d84..a80f701a3 100644 --- a/crates/tabby-scheduler/tests/integration_test.rs +++ b/crates/tabby-scheduler/tests/integration_test.rs @@ -3,7 +3,7 @@ mod tests { use std::fs::create_dir_all; use tabby_common::{ - config::{Config, Repository}, + config::{Config, RepositoryConfig, ServerConfig}, path::set_tabby_root, }; use temp_testdir::*; @@ -17,9 +17,10 @@ mod tests { set_tabby_root(root.to_path_buf()); let config = Config { - repositories: vec![Repository { + repositories: vec![RepositoryConfig { git_url: "https://github.com/TabbyML/interview-questions".to_owned(), }], + server: ServerConfig::default(), }; config.save(); diff --git a/crates/tabby/src/serve/mod.rs b/crates/tabby/src/serve/mod.rs index 70faad8b2..4bd6d2c38 100644 --- a/crates/tabby/src/serve/mod.rs +++ b/crates/tabby/src/serve/mod.rs @@ -125,7 +125,7 @@ fn should_download_ggml_files(device: &Device) -> bool { *device == Device::Metal } -pub async fn main(_config: &Config, args: &ServeArgs) { +pub async fn main(config: &Config, args: &ServeArgs) { valid_args(args); if args.device != Device::ExperimentalHttp { @@ -146,7 +146,7 @@ pub async fn main(_config: &Config, args: &ServeArgs) { .route("/", routing::get(playground::handler)) .route("/index.txt", routing::get(playground::handler)) .route("/_next/*path", routing::get(playground::handler)) - .merge(api_router(args)) + .merge(api_router(args, config)) .merge(SwaggerUi::new("/swagger-ui").url("/api-docs/openapi.json", doc)); let app = if args.chat_model.is_some() { @@ -166,7 +166,7 @@ pub async fn main(_config: &Config, args: &ServeArgs) { .unwrap_or_else(|err| fatal!("Error happens during serving: {}", err)) } -fn api_router(args: &ServeArgs) -> Router { +fn api_router(args: &ServeArgs, config: &Config) -> Router { let index_server = Arc::new(IndexServer::new()); let completion_state = { let ( @@ -218,7 +218,9 @@ fn api_router(args: &ServeArgs) -> Router { "/v1/completions", routing::post(completions::completions).with_state(completion_state), ) - .layer(TimeoutLayer::new(Duration::from_secs(3))) + .layer(TimeoutLayer::new(Duration::from_secs( + config.server.completion_timeout, + ))) }); if let Some(chat_state) = chat_state {