Skip to content

Commit

Permalink
feat: add ee feature flag to distinguish the OSS and EE offering (#789)
Browse files Browse the repository at this point in the history
* feat: add ee feature flag to distinguish the OSS offering and EE augmented offering

* [autofix.ci] apply automated fixes

* fix lint

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
  • Loading branch information
wsxiaoys and autofix-ci[bot] authored Nov 15, 2023
1 parent bd072d8 commit abe0411
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 8 deletions.
8 changes: 5 additions & 3 deletions crates/tabby/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@ version = "0.6.0-dev"
edition = "2021"

[features]
default = ["ee"]
ee = ["dep:tabby-webserver"]
cuda = ["llama-cpp-bindings/cuda"]
experimental-http = ["dep:http-api-bindings"]

Expand All @@ -15,7 +17,7 @@ tabby-inference = { path = "../tabby-inference" }
axum.workspace = true
hyper = { workspace = true }
tokio = { workspace = true, features = ["full"] }
utoipa = { workspace= true, features = ["axum_extras", "preserve_order"] }
utoipa = { workspace = true, features = ["axum_extras", "preserve_order"] }
utoipa-swagger-ui = { version = "3.1", features = ["axum"] }
serde = { workspace = true }
serdeconv = { workspace = true }
Expand All @@ -35,7 +37,7 @@ tantivy = { workspace = true }
anyhow = { workspace = true }
sysinfo = "0.29.8"
nvml-wrapper = "0.9.0"
http-api-bindings = { path = "../http-api-bindings", optional = true } # included when build with `experimental-http` feature
http-api-bindings = { path = "../http-api-bindings", optional = true } # included when build with `experimental-http` feature
async-stream = { workspace = true }
axum-streams = { version = "0.9.1", features = ["json"] }
minijinja = { version = "1.0.8", features = ["loader"] }
Expand All @@ -44,7 +46,7 @@ regex.workspace = true
llama-cpp-bindings = { path = "../llama-cpp-bindings" }
futures.workspace = true
async-trait.workspace = true
tabby-webserver = { path = "../../ee/tabby-webserver" }
tabby-webserver = { path = "../../ee/tabby-webserver", optional = true }
thiserror.workspace = true
chrono = "0.4.31"

Expand Down
15 changes: 12 additions & 3 deletions crates/tabby/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@ mod services;

mod download;
mod serve;

#[cfg(feature = "ee")]
mod worker;

use clap::{Parser, Subcommand};
Expand All @@ -14,7 +16,6 @@ use opentelemetry::{
};
use opentelemetry_otlp::WithExportConfig;
use tabby_common::config::Config;
use tabby_webserver::api::WorkerKind;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer};

#[derive(Parser)]
Expand All @@ -41,11 +42,13 @@ pub enum Commands {
Scheduler(SchedulerArgs),

/// Run completion model as worker
#[cfg(feature = "ee")]
#[clap(name = "worker::completion")]
#[command(arg_required_else_help = true)]
WorkerCompletion(worker::WorkerArgs),

/// Run chat model as worker
#[cfg(feature = "ee")]
#[clap(name = "worker::chat")]
#[command(arg_required_else_help = true)]
WorkerChat(worker::WorkerArgs),
Expand Down Expand Up @@ -106,8 +109,14 @@ async fn main() {
Commands::Scheduler(args) => tabby_scheduler::scheduler(args.now)
.await
.unwrap_or_else(|err| fatal!("Scheduler failed due to '{}'", err)),
Commands::WorkerCompletion(args) => worker::main(WorkerKind::Completion, args).await,
Commands::WorkerChat(args) => worker::main(WorkerKind::Chat, args).await,
#[cfg(feature = "ee")]
Commands::WorkerCompletion(args) => {
worker::main(tabby_webserver::api::WorkerKind::Completion, args).await
}
#[cfg(feature = "ee")]
Commands::WorkerChat(args) => {
worker::main(tabby_webserver::api::WorkerKind::Chat, args).await
}
}

opentelemetry::global::shutdown_tracer_provider();
Expand Down
7 changes: 5 additions & 2 deletions crates/tabby/src/serve.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ use axum::{routing, Router, Server};
use axum_tracing_opentelemetry::opentelemetry_tracing_layer;
use clap::Args;
use tabby_common::{config::Config, usage};
use tabby_webserver::attach_webserver;
use tokio::time::sleep;
use tower_http::{cors::CorsLayer, timeout::TimeoutLayer};
use tracing::info;
Expand Down Expand Up @@ -106,7 +105,11 @@ pub async fn main(config: &Config, args: &ServeArgs) {
.merge(api_router(args, config).await)
.merge(SwaggerUi::new("/swagger-ui").url("/api-docs/openapi.json", ApiDoc::openapi()));

let app = attach_webserver(app).await;
#[cfg(feature = "ee")]
let app = tabby_webserver::attach_webserver(app).await;

#[cfg(not(feature = "ee"))]
let app = app.fallback(|| async { axum::response::Redirect::permanent("/swagger-ui") });

let address = SocketAddr::from((Ipv4Addr::UNSPECIFIED, args.port));
info!("Listening at {}", address);
Expand Down

0 comments on commit abe0411

Please sign in to comment.