Implement NlpService embeddings methods (#9)

Co-authored-by: Nick Hill <[email protected]>
Signed-off-by: Prashant <[email protected]>

Author: 2 people

.gitignore

@@ -1,3 +1,4 @@
+.DS_Store
 .idea
 target
 fmaas-router/src/pb

Cargo.lock

@@ -20,12 +20,13 @@ clap = { version = "^4.5.3", features = ["derive", "env"] }
 futures = "^0.3.30"
 tonic = { version = "=0.11.0", features = ["tls"] }
 ginepro = "=0.7.1"
-lazy_static = "^1.4.0"
 tokio = { version = "^1.36.0", features = ["rt", "rt-multi-thread", "parking_lot", "signal", "sync", "fs"] }
 tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.18", features = ["json", "env-filter"] }
 prost = "^0.12.3"
+prost-types = "^0.12.3"
 serde_yaml = "^0.9.33"
+serde = { version = "^1.0.197", features = ["derive"] }
 
 mio = "^0.8.11" # Override to address CVE-2024-27308
 rustls-webpki = "^0.102.2" # Override to address WS-2023-0305, CVE-2018-16875

fmaas-router/Cargo.toml

@@ -7,7 +7,10 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         .build_server(true)
         .out_dir("src/pb")
         .include_file("mod.rs")
-        .compile(&["../proto/generation.proto"], &["../proto"])
+        .compile(
+            &["../proto/generation.proto", "../proto/nlpservice.proto"],
+            &["../proto"],
+        )
         .unwrap_or_else(|e| panic!("protobuf compilation failed: {}", e));
 
     Ok(())

fmaas-router/build.rs

@@ -1,5 +1,121 @@
-extern crate core;
-extern crate lazy_static;
+use std::{collections::HashMap, path::Path};
+use anyhow::Context;
+use futures::future::try_join_all;
+use ginepro::LoadBalancedChannel;
 
+use serde::{Deserialize, Deserializer};
+use tonic::transport::ClientTlsConfig;
+use tracing::info;
+
+#[allow(clippy::enum_variant_names)]
 mod pb;
+pub mod rpc;
 pub mod server;
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct ServiceAddr {
+    pub hostname: String,
+    pub port: Option<u16>,
+}
+
+/// Old format without top-level keys, generation models only.
+#[derive(Debug, Clone, Deserialize)]
+pub struct ModelMapV1(#[serde(deserialize_with = "de_service_addr")] HashMap<String, ServiceAddr>);
+
+/// New format with top-level keys for generation and embeddings models.
+#[derive(Debug, Clone, Deserialize)]
+pub struct ModelMapV2 {
+    #[serde(deserialize_with = "de_service_addr", default = "HashMap::default")]
+    generation: HashMap<String, ServiceAddr>,
+    #[serde(deserialize_with = "de_service_addr", default = "HashMap::default")]
+    embeddings: HashMap<String, ServiceAddr>,
+}
+
+/// Maps model names to service address.
+#[derive(Debug, Clone, Deserialize)]
+#[serde(untagged)]
+pub enum ModelMap {
+    V1(ModelMapV1),
+    V2(ModelMapV2),
+}
+
+impl ModelMap {
+    pub fn load(path: impl AsRef<Path>) -> Self {
+        let s = std::fs::read_to_string(path).expect("Failed to load model map config");
+        serde_yaml::from_str(&s).expect("Invalid model map config")
+    }
+
+    pub fn generation(&self) -> Option<&HashMap<String, ServiceAddr>> {
+        match self {
+            ModelMap::V1(v1) => Some(&v1.0),
+            ModelMap::V2(v2) => (!v2.generation.is_empty()).then_some(&v2.generation),
+        }
+    }
+
+    pub fn embeddings(&self) -> Option<&HashMap<String, ServiceAddr>> {
+        match self {
+            ModelMap::V1(_) => None,
+            ModelMap::V2(v2) => (!v2.embeddings.is_empty()).then_some(&v2.embeddings),
+        }
+    }
+}
+
+fn service_addr_from_str<'de, D>(deserializer: D) -> Result<ServiceAddr, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    let s = String::deserialize(deserializer).map_err(serde::de::Error::custom)?;
+    let mut parts = s.split(':');
+    let hostname = parts.next().unwrap().to_string();
+    let port = parts.next().map(|p| {
+        p.parse::<u16>()
+            .unwrap_or_else(|_| panic!("Invalid port in configured service name: {p}"))
+    });
+    if parts.next().is_some() {
+        panic!("Configured service name contains more than one : character");
+    }
+    Ok(ServiceAddr { hostname, port })
+}
+
+fn de_service_addr<'de, D>(deserializer: D) -> Result<HashMap<String, ServiceAddr>, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    #[derive(Deserialize)]
+    struct Wrapper(#[serde(deserialize_with = "service_addr_from_str")] ServiceAddr);
+
+    let v = HashMap::<String, Wrapper>::deserialize(deserializer)?;
+    Ok(v.into_iter().map(|(k, Wrapper(v))| (k, v)).collect())
+}
+
+async fn create_clients<C>(
+    default_target_port: u16,
+    client_tls: Option<&ClientTlsConfig>,
+    model_map: &HashMap<String, ServiceAddr>,
+    new: fn(LoadBalancedChannel) -> C,
+) -> HashMap<String, C> {
+    let clients = model_map
+        .iter()
+        .map(|(name, service_addr)| async move {
+            info!("Configuring client for model name: [{name}]");
+            // Build a load-balanced channel given a service name and a port.
+            let mut builder = LoadBalancedChannel::builder((
+                service_addr.hostname.clone(),
+                service_addr.port.unwrap_or(default_target_port),
+            ));
+            if let Some(tls_config) = client_tls {
+                builder = builder.with_tls(tls_config.clone());
+            }
+            let channel = builder
+                .channel()
+                .await
+                .context(format!("Channel failed for service {name}"))?;
+            Ok((name.clone(), new(channel))) as Result<(String, C), anyhow::Error>
+        })
+        .collect::<Vec<_>>();
+    try_join_all(clients)
+        .await
+        .expect("Error creating upstream service clients")
+        .into_iter()
+        .collect()
+}

fmaas-router/src/lib.rs

@@ -1,7 +1,7 @@
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 
 use clap::Parser;
-use fmaas_router::server;
+use fmaas_router::{server, ModelMap};
 use tracing_subscriber::EnvFilter;
 
 /// App Configuration
@@ -17,6 +17,8 @@ struct Args {
     #[clap(long, env)]
     json_output: bool,
     #[clap(long, env)]
+    model_map_config: String,
+    #[clap(long, env)]
     tls_cert_path: Option<String>,
     #[clap(long, env)]
     tls_key_path: Option<String>,
@@ -53,19 +55,19 @@ fn main() -> Result<(), std::io::Error> {
     if args.tls_key_path.is_some() != args.tls_cert_path.is_some() {
         panic!("tls: must provide both cert and key")
     }
-
     if args.tls_client_ca_cert_path.is_some() && args.tls_cert_path.is_none() {
         panic!("tls: cannot provide client ca cert without keypair")
     }
 
+    // Load model map config
+    let model_map = ModelMap::load(args.model_map_config);
+
     // Launch Tokio runtime
     tokio::runtime::Builder::new_multi_thread()
         .enable_all()
         .build()
         .unwrap()
         .block_on(async {
-            //TODO initialize clients
-
             let grpc_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), args.grpc_port);
             let http_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), args.probe_port);
 
@@ -78,6 +80,7 @@ fn main() -> Result<(), std::io::Error> {
                 args.default_upstream_port,
                 args.upstream_tls,
                 args.upstream_tls_ca_cert_path,
+                model_map,
             )
             .await;
 

fmaas-router/src/main.rs

@@ -0,0 +1,2 @@
+pub mod generation;
+pub mod nlp;

fmaas-router/src/rpc.rs

@@ -0,0 +1,108 @@
+use std::collections::HashMap;
+
+use ginepro::LoadBalancedChannel;
+use tonic::{transport::ClientTlsConfig, Request, Response, Status, Streaming};
+use tracing::{debug, instrument};
+
+use crate::{pb::fmaas::{
+    generation_service_client::GenerationServiceClient,
+    generation_service_server::GenerationService, BatchedGenerationRequest,
+    BatchedGenerationResponse, BatchedTokenizeRequest, BatchedTokenizeResponse,
+    GenerationResponse, ModelInfoRequest, ModelInfoResponse, SingleGenerationRequest,
+}, create_clients, ServiceAddr};
+
+#[derive(Debug, Default)]
+pub struct GenerationServicer {
+    clients: HashMap<String, GenerationServiceClient<LoadBalancedChannel>>,
+}
+
+impl GenerationServicer {
+    pub async fn new(
+        default_target_port: u16,
+        client_tls: Option<&ClientTlsConfig>,
+        model_map: &HashMap<String, ServiceAddr>,
+    ) -> Self {
+        let clients = create_clients(
+            default_target_port, client_tls, model_map, GenerationServiceClient::new
+        ).await;
+        Self { clients }
+    }
+
+    async fn client(
+        &self,
+        model_id: &str,
+    ) -> Result<GenerationServiceClient<LoadBalancedChannel>, Status> {
+        Ok(self
+            .clients
+            .get(model_id)
+            .ok_or_else(|| Status::not_found(format!("Unrecognized model_id: {model_id}")))?
+            .clone())
+    }
+}
+
+#[tonic::async_trait]
+impl GenerationService for GenerationServicer {
+    #[instrument(skip_all)]
+    async fn generate(
+        &self,
+        request: Request<BatchedGenerationRequest>,
+    ) -> Result<Response<BatchedGenerationResponse>, Status> {
+        let br = request.get_ref();
+        if br.requests.is_empty() {
+            return Ok(Response::new(BatchedGenerationResponse {
+                responses: vec![],
+            }));
+        }
+        debug!("Routing generation request for Model ID {}", &br.model_id);
+        self.client(&br.model_id).await?.generate(request).await
+    }
+
+    type GenerateStreamStream = Streaming<GenerationResponse>;
+
+    #[instrument(skip_all)]
+    async fn generate_stream(
+        &self,
+        request: Request<SingleGenerationRequest>,
+    ) -> Result<Response<Self::GenerateStreamStream>, Status> {
+        let sr = request.get_ref();
+        if sr.request.is_none() {
+            return Err(Status::invalid_argument("missing request"));
+        }
+        debug!(
+            "Routing streaming generation request for Model ID {}",
+            &sr.model_id
+        );
+        self.client(&sr.model_id)
+            .await?
+            .generate_stream(request)
+            .await
+    }
+
+    #[instrument(skip_all)]
+    async fn tokenize(
+        &self,
+        request: Request<BatchedTokenizeRequest>,
+    ) -> Result<Response<BatchedTokenizeResponse>, Status> {
+        let br = request.get_ref();
+        if br.requests.is_empty() {
+            return Ok(Response::new(BatchedTokenizeResponse { responses: vec![] }));
+        }
+        debug!("Routing tokenization request for Model ID {}", &br.model_id);
+        self.client(&br.model_id).await?.tokenize(request).await
+    }
+
+    #[instrument(skip_all)]
+    async fn model_info(
+        &self,
+        request: Request<ModelInfoRequest>,
+    ) -> Result<Response<ModelInfoResponse>, Status> {
+        debug!(
+            "Routing model info request for Model ID {}",
+            &request.get_ref().model_id
+        );
+        self.client(&request.get_ref().model_id)
+            .await?
+            .model_info(request)
+            .await
+    }
+}