Skip to content

Commit

Permalink
feat: add server.completion_timeout to control timeout of /v1/complet…
Browse files Browse the repository at this point in the history
…ion (#637)

* feat: add server.completion_timeout to control timeout of /v1/completion

* Update config.rs
  • Loading branch information
wsxiaoys authored Oct 25, 2023
1 parent d6296bb commit 21ec60e
Show file tree
Hide file tree
Showing 5 changed files with 36 additions and 16 deletions.
29 changes: 23 additions & 6 deletions crates/tabby-common/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,10 @@ use crate::path::{config_file, repositories_dir};
#[derive(Serialize, Deserialize, Default)]
pub struct Config {
#[serde(default)]
pub repositories: Vec<Repository>,
pub repositories: Vec<RepositoryConfig>,

#[serde(default)]
pub server: ServerConfig,
}

impl Config {
Expand All @@ -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();
Expand All @@ -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() {
Expand All @@ -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());
Expand Down
4 changes: 2 additions & 2 deletions crates/tabby-scheduler/src/dataset.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};
Expand All @@ -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();

Expand Down
4 changes: 2 additions & 2 deletions crates/tabby-scheduler/src/repository.rs
Original file line number Diff line number Diff line change
@@ -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<()>;
Expand All @@ -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();
Expand Down
5 changes: 3 additions & 2 deletions crates/tabby-scheduler/tests/integration_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;
Expand All @@ -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();
Expand Down
10 changes: 6 additions & 4 deletions crates/tabby/src/serve/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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() {
Expand All @@ -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 (
Expand Down Expand Up @@ -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 {
Expand Down

0 comments on commit 21ec60e

Please sign in to comment.