feat: add Prometheus support to Tabby. (#838)

* Added Prometheus support to Tabby.

1) Added `axum-prometheus` to Cargo.toml

2) Added `metrics.rs` as a simple route to produce a metrics endpoint with prometheus-formatted events

3) Added `/v1/metrics` endpoint with API doc entry

4) Added the `PrometheusLayer` to the root layers.

This change effectively allows for external Prometheus to scrape metrics from Tabby during execution in order to monitor operations (resources, timings and overall usage).

* changed application/text to text/plain to adhere to rfc1341

* Update Makefile

* Update Makefile

* Update Makefile

* Update Makefile

* Update Makefile

* Update Makefile

* - Reworked code to fit upstream changes
- Added Prometheus layer and handle to worker.rs

* Update CHANGELOG.md

* [autofix.ci] apply automated fixes

* Simplified worker metrics route definition

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 2/3)

---------

Co-authored-by: Meng Zhang <meng@tabbyml.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 * Add distribution support (running completion / chat model on different process / machine).
 * Add conversation history in chat playground.
+* Add `/v1/metrics` endpoint for prometheus metrics collection. 
 
 ## Fixes and Improvements
 

Makefile

@@ -24,8 +24,9 @@ bump-release-version:
 	cargo ws version --allow-branch "r*" --no-individual-tags --force "*"
 
 update-openapi-doc:
-	curl http://localhost:8080/api-docs/openapi.json | jq '                                                       \
-	  delpaths([                                                                                                  \
+	curl http://localhost:8080/api-docs/openapi.json | jq '                                                             \
+	  delpaths([                                                                                                        \
+		  ["paths", "/v1/metrics"],                                                                                 \
 		  ["paths", "/v1beta/chat/completions"],                                                                    \
 		  ["paths", "/v1beta/search"],                                                                              \
 		  ["components", "schemas", "CompletionRequest", "properties", "prompt"],                                   \
@@ -37,4 +38,4 @@ update-openapi-doc:
 			    > website/static/openapi.json
 
 update-graphql-schema:
-	cargo run --package tabby-webserver --example update-schema
+	cargo run --package tabby-webserver --example update-schema

crates/tabby/Cargo.toml

@@ -49,6 +49,7 @@ async-trait.workspace = true
 tabby-webserver = { path = "../../ee/tabby-webserver", optional = true }
 thiserror.workspace = true
 chrono = "0.4.31"
+axum-prometheus = "0.4.0"
 
 [dependencies.uuid]
 version = "1.3.3"

crates/tabby/src/routes/metrics.rs

@@ -0,0 +1,16 @@
+use std::sync::Arc;
+
+use axum::extract::State;
+use axum_prometheus::metrics_exporter_prometheus::PrometheusHandle;
+
+#[utoipa::path(
+    get,
+    path = "/v1/metrics",
+    tag = "v1",
+    responses(
+        (status = 200, description = "Success", body = String, content_type = "text/plain"),
+    )
+)]
+pub async fn metrics(State(state): State<Arc<PrometheusHandle>>) -> String {
+    state.render()
+}

crates/tabby/src/routes/mod.rs

@@ -2,10 +2,12 @@ mod chat;
 mod completions;
 mod events;
 mod health;
+mod metrics;
 mod search;
 
 pub use chat::*;
 pub use completions::*;
 pub use events::*;
 pub use health::*;
+pub use metrics::*;
 pub use search::*;

crates/tabby/src/serve.rs

@@ -5,6 +5,7 @@ use std::{
 };
 
 use axum::{routing, Router, Server};
+use axum_prometheus::{metrics_exporter_prometheus::PrometheusHandle, PrometheusMetricLayer};
 use axum_tracing_opentelemetry::opentelemetry_tracing_layer;
 use clap::Args;
 use tabby_common::{
@@ -49,7 +50,7 @@ Install following IDE / Editor extensions to get started with [Tabby](https://gi
     servers(
         (url = "/", description = "Server"),
     ),
-    paths(routes::log_event, routes::completions, routes::completions, routes::health, routes::search),
+    paths(routes::log_event, routes::completions, routes::completions, routes::health, routes::search, routes::metrics),
     components(schemas(
         api::event::LogEventRequest,
         completion::CompletionRequest,
@@ -108,9 +109,11 @@ pub async fn main(config: &Config, args: &ServeArgs) {
 
     let logger = Arc::new(create_logger());
     let code = Arc::new(create_code_search());
+    let (prometheus_layer, prometheus_handle) = PrometheusMetricLayer::pair();
+    let metrics_handle = Arc::new(prometheus_handle);
 
     let app = Router::new()
-        .merge(api_router(args, config, logger.clone(), code.clone()).await)
+        .merge(api_router(args, config, logger.clone(), code.clone(), metrics_handle).await)
         .merge(SwaggerUi::new("/swagger-ui").url("/api-docs/openapi.json", ApiDoc::openapi()));
 
     #[cfg(feature = "ee")]
@@ -121,7 +124,8 @@ pub async fn main(config: &Config, args: &ServeArgs) {
 
     let app = app
         .layer(CorsLayer::permissive())
-        .layer(opentelemetry_tracing_layer());
+        .layer(opentelemetry_tracing_layer())
+        .layer(prometheus_layer);
 
     let address = SocketAddr::from((Ipv4Addr::UNSPECIFIED, args.port));
     info!("Listening at {}", address);
@@ -148,6 +152,7 @@ async fn api_router(
     config: &Config,
     logger: Arc<dyn EventLogger>,
     code: Arc<dyn CodeSearch>,
+    metrics_handle: Arc<PrometheusHandle>,
 ) -> Router {
     let completion_state = if let Some(model) = &args.model {
         Some(Arc::new(
@@ -179,6 +184,7 @@ async fn api_router(
         args.chat_model.as_deref(),
         &args.device,
     ));
+
     routers.push({
         Router::new()
             .route(
@@ -193,6 +199,10 @@ async fn api_router(
                 "/v1/health",
                 routing::get(routes::health).with_state(health_state),
             )
+            .route(
+                "/v1/metrics",
+                routing::get(routes::metrics).with_state(metrics_handle),
+            )
     });
 
     if let Some(completion_state) = completion_state {

crates/tabby/src/worker.rs

@@ -6,6 +6,7 @@ use std::{
 
 use anyhow::Result;
 use axum::{routing, Router};
+use axum_prometheus::PrometheusMetricLayer;
 use axum_tracing_opentelemetry::opentelemetry_tracing_layer;
 use clap::Args;
 use hyper::Server;
@@ -84,14 +85,22 @@ pub async fn main(kind: WorkerKind, args: &WorkerArgs) {
     info!("Starting worker, this might takes a few minutes...");
 
     let context = WorkerContext::new(&args.url).await;
+
+    let (prometheus_layer, prometheus_handle) = PrometheusMetricLayer::pair();
+
     let app = match kind {
         WorkerKind::Completion => make_completion_route(context, args).await,
         WorkerKind::Chat => make_chat_route(context, args).await,
     };
 
     let app = app
+        .route(
+            "/v1/metrics",
+            routing::get(routes::metrics).with_state(Arc::new(prometheus_handle)),
+        )
         .layer(CorsLayer::permissive())
-        .layer(opentelemetry_tracing_layer());
+        .layer(opentelemetry_tracing_layer())
+        .layer(prometheus_layer);
 
     let address = SocketAddr::from((Ipv4Addr::UNSPECIFIED, args.port));
     info!("Listening at {}", address);