🐛 Fix model id configuration bug for nlp client and rename tgis-route…

…r to common-router (#89) * 🐛 Fix model id configuration bug for nlp client and rename tgis-router to common-router Signed-off-by: gkumbhat <[email protected]> * 🎨 Fix formatting Signed-off-by: gkumbhat <[email protected]> * 🎨🚚 Move common router key as const and re-use that Signed-off-by: gkumbhat <[email protected]> --------- Signed-off-by: gkumbhat <[email protected]>
foundation-model-stack · Jun 19, 2024 · e5b72c1 · e5b72c1
1 parent 19f7f59
commit e5b72c1
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 6 deletions.
diff --git a/src/clients.rs b/src/clients.rs
@@ -28,6 +28,7 @@ pub const DEFAULT_TGIS_PORT: u16 = 8033;
 pub const DEFAULT_CAIKIT_NLP_PORT: u16 = 8085;
 pub const DEFAULT_CHUNKER_PORT: u16 = 8085;
 pub const DEFAULT_DETECTOR_PORT: u16 = 8080;
+pub const COMMON_ROUTER_KEY: &str = "common-router";
 const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
 const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
 

diff --git a/src/clients/nlp.rs b/src/clients/nlp.rs
@@ -8,6 +8,7 @@ use tonic::Request;
 
 use super::{create_grpc_clients, Error};
 use crate::{
+    clients::COMMON_ROUTER_KEY,
     config::ServiceConfig,
     pb::{
         caikit::runtime::nlp::{
@@ -34,7 +35,9 @@ impl NlpClient {
         Self { clients }
     }
 
-    fn client(&self, model_id: &str) -> Result<NlpServiceClient<LoadBalancedChannel>, Error> {
+    fn client(&self, _model_id: &str) -> Result<NlpServiceClient<LoadBalancedChannel>, Error> {
+        // NOTE: We currently forward requests to common router, so we use a single client.
+        let model_id = COMMON_ROUTER_KEY;
         Ok(self
             .clients
             .get(model_id)

diff --git a/src/clients/tgis.rs b/src/clients/tgis.rs
@@ -7,6 +7,7 @@ use tokio_stream::wrappers::ReceiverStream;
 
 use super::{create_grpc_clients, Error};
 use crate::{
+    clients::COMMON_ROUTER_KEY,
     config::ServiceConfig,
     pb::fmaas::{
         generation_service_client::GenerationServiceClient, BatchedGenerationRequest,
@@ -32,8 +33,8 @@ impl TgisClient {
         &self,
         _model_id: &str,
     ) -> Result<GenerationServiceClient<LoadBalancedChannel>, Error> {
-        // NOTE: We currently forward requests to the tgis-router, so we use a single client.
-        let model_id = "tgis-router";
+        // NOTE: We currently forward requests to the common-router, so we use a single client.
+        let model_id = COMMON_ROUTER_KEY;
         Ok(self
             .clients
             .get(model_id)

diff --git a/src/orchestrator.rs b/src/orchestrator.rs
@@ -13,7 +13,7 @@ use uuid::Uuid;
 use crate::{
     clients::{
         self, detector::ContentAnalysisRequest, ChunkerClient, DetectorClient, GenerationClient,
-        NlpClient, TgisClient,
+        NlpClient, TgisClient, COMMON_ROUTER_KEY,
     },
     config::{GenerationProvider, OrchestratorConfig},
     models::{
@@ -601,15 +601,21 @@ async fn create_clients(
         GenerationProvider::Tgis => {
             let client = TgisClient::new(
                 clients::DEFAULT_TGIS_PORT,
-                &[("tgis-router".to_string(), config.generation.service.clone())],
+                &[(
+                    COMMON_ROUTER_KEY.to_string(),
+                    config.generation.service.clone(),
+                )],
             )
             .await;
             GenerationClient::Tgis(client)
         }
         GenerationProvider::Nlp => {
             let client = NlpClient::new(
                 clients::DEFAULT_CAIKIT_NLP_PORT,
-                &[("tgis-router".to_string(), config.generation.service.clone())],
+                &[(
+                    COMMON_ROUTER_KEY.to_string(),
+                    config.generation.service.clone(),
+                )],
             )
             .await;
             GenerationClient::Nlp(client)