Add model and adapter headings

router/Cargo.toml

@@ -41,6 +41,7 @@ tracing-subscriber = { version = "0.3.16", features = ["json", "env-filter"] }
 utoipa = { version = "3.0.1", features = ["axum_extras"] }
 utoipa-swagger-ui = { version = "3.0.2", features = ["axum"] }
 ngrok = { version = "0.12.3", features = ["axum"], optional = true }
+once_cell = "1.19.0"
 
 [build-dependencies]
 vergen = { version = "8.0.0", features = ["build", "git", "gitcl"] }

router/src/adapter.rs

@@ -73,3 +73,24 @@ impl PartialEq for Adapter {
         self.index == other.index
     }
 }
+
+pub(crate) fn extract_adapter_params(
+    adapter_id: Option<String>,
+    adapter_source: Option<String>,
+    adapter_parameters: Option<AdapterParameters>,
+) -> (Option<String>, AdapterParameters) {
+    let mut adapter_id = adapter_id.clone();
+    if adapter_id.is_none() || adapter_id.as_ref().unwrap().is_empty() {
+        adapter_id = Some(BASE_MODEL_ADAPTER_ID.to_string());
+    }
+    let mut adapter_source = adapter_source.clone();
+    if adapter_source.is_none() {
+        adapter_source = Some(DEFAULT_ADAPTER_SOURCE.to_string());
+    }
+
+    let adapter_parameters = adapter_parameters.clone().unwrap_or(AdapterParameters {
+        adapter_ids: vec![adapter_id.clone().unwrap()],
+        ..Default::default()
+    });
+    return (adapter_source, adapter_parameters);
+}

router/src/infer.rs

@@ -1,5 +1,5 @@
 /// Batching and inference logic
-use crate::adapter::{Adapter, BASE_MODEL_ADAPTER_ID, DEFAULT_ADAPTER_SOURCE};
+use crate::adapter::{extract_adapter_params, Adapter};
 use crate::queue::AdapterEvent;
 use crate::scheduler::AdapterScheduler;
 use crate::validation::{Validation, ValidationError};
@@ -123,24 +123,11 @@ impl Infer {
                 err
             })?;
 
-        let mut adapter_id = request.parameters.adapter_id.clone();
-        if adapter_id.is_none() || adapter_id.as_ref().unwrap().is_empty() {
-            adapter_id = Some(BASE_MODEL_ADAPTER_ID.to_string());
-        }
-        let mut adapter_source = request.parameters.adapter_source.clone();
-        if adapter_source.is_none() {
-            adapter_source = Some(DEFAULT_ADAPTER_SOURCE.to_string());
-        }
-
-        let adapter_parameters =
-            request
-                .parameters
-                .adapter_parameters
-                .clone()
-                .unwrap_or(AdapterParameters {
-                    adapter_ids: vec![adapter_id.clone().unwrap()],
-                    ..Default::default()
-                });
+        let (adapter_source, adapter_parameters) = extract_adapter_params(
+            request.parameters.adapter_id.clone(),
+            request.parameters.adapter_source.clone(),
+            request.parameters.adapter_parameters.clone(),
+        );
 
         let adapter_idx;
         {

router/src/server.rs

@@ -1,4 +1,5 @@
 /// HTTP Server logic
+use crate::adapter::{extract_adapter_params, BASE_MODEL_ADAPTER_ID};
 use crate::health::Health;
 use crate::infer::{InferError, InferResponse, InferStreamResponse};
 use crate::validation::ValidationError;
@@ -19,6 +20,7 @@ use futures::stream::StreamExt;
 use futures::Stream;
 use lorax_client::{ShardInfo, ShardedClient};
 use metrics_exporter_prometheus::{Matcher, PrometheusBuilder, PrometheusHandle};
+use once_cell::sync::OnceCell;
 use std::convert::Infallible;
 use std::net::SocketAddr;
 use std::sync::atomic::AtomicBool;
@@ -31,6 +33,8 @@ use tracing::{info_span, instrument, Instrument};
 use utoipa::OpenApi;
 use utoipa_swagger_ui::SwaggerUi;
 
+static MODEL_ID: OnceCell<String> = OnceCell::new();
+
 /// Generate tokens if `stream == false` or a stream of token if `stream == true`
 #[utoipa::path(
 post,
@@ -286,6 +290,11 @@ async fn generate(
     }
 
     let details = req.0.parameters.details || req.0.parameters.decoder_input_details;
+    let (adapter_source, adapter_parameters) = extract_adapter_params(
+        req.0.parameters.adapter_id.clone(),
+        req.0.parameters.adapter_source.clone(),
+        req.0.parameters.adapter_parameters.clone(),
+    );
 
     // Inference
     let (response, best_of_responses) = match req.0.parameters.best_of {
@@ -397,6 +406,22 @@ async fn generate(
         time_per_token.as_millis().to_string().parse().unwrap(),
     );
 
+    headers.insert("x-model-id", MODEL_ID.get().unwrap().parse().unwrap());
+
+    let adapter_id_string = adapter_parameters
+        .adapter_ids
+        .iter()
+        .map(|id| id.as_str())
+        // filter out base model adapter id
+        .filter(|id| *id != BASE_MODEL_ADAPTER_ID)
+        .collect::<Vec<_>>()
+        .join(",");
+
+    if adapter_id_string.len() > 0 {
+        headers.insert("x-adapter-ids", adapter_id_string.parse().unwrap());
+        headers.insert("x-adapter-source", adapter_source.unwrap().parse().unwrap());
+    }
+
     // Metrics
     metrics::increment_counter!("lorax_request_success");
     metrics::histogram!("lorax_request_duration", total_time.as_secs_f64());
@@ -752,6 +777,8 @@ pub async fn run(
         generation_health,
     );
 
+    let model_id = model_info.model_id.clone();
+
     // Duration buckets
     let duration_matcher = Matcher::Suffix(String::from("duration"));
     let n_duration_buckets = 35;
@@ -826,6 +853,10 @@ pub async fn run(
         docker_label: option_env!("DOCKER_LABEL"),
     };
 
+    MODEL_ID.set(model_id.clone()).unwrap_or_else(|_| {
+        panic!("MODEL_ID was already set!");
+    });
+
     // Create router
     let app = Router::new()
         .merge(SwaggerUi::new("/docs").url("/api-doc/openapi.json", ApiDoc::openapi()))