fix: disabling sse keep-alive (NVIDIA#408)

Author: ryanolson

lib/llm/src/http/service.rs

@@ -52,6 +52,7 @@ use crate::types::openai::{
 use std::{
     collections::HashMap,
     sync::{Arc, Mutex},
+    time::Duration,
 };
 
 #[derive(Clone)]
@@ -191,6 +192,7 @@ pub struct DeploymentState {
     completion_engines: Arc<Mutex<ModelEngines<OpenAICompletionsStreamingEngine>>>,
     chat_completion_engines: Arc<Mutex<ModelEngines<OpenAIChatCompletionsStreamingEngine>>>,
     metrics: Arc<Metrics>,
+    sse_keep_alive: Option<Duration>,
 }
 
 impl DeploymentState {
@@ -199,6 +201,7 @@ impl DeploymentState {
             completion_engines: Arc::new(Mutex::new(ModelEngines::default())),
             chat_completion_engines: Arc::new(Mutex::new(ModelEngines::default())),
             metrics: Arc::new(Metrics::default()),
+            sse_keep_alive: None,
         }
     }
 

lib/llm/src/http/service/openai.rs

@@ -182,9 +182,13 @@ async fn completions(
         let stream = stream.map(|response| Event::try_from(EventConverter::from(response)));
         let stream = monitor_for_disconnects(stream.boxed(), ctx, inflight).await;
 
-        Ok(Sse::new(stream)
-            .keep_alive(KeepAlive::default())
-            .into_response())
+        let mut sse_stream = Sse::new(stream);
+
+        if let Some(keep_alive) = state.sse_keep_alive {
+            sse_stream = sse_stream.keep_alive(KeepAlive::default().interval(keep_alive));
+        }
+
+        Ok(sse_stream.into_response())
     } else {
         let response = CompletionResponse::from_annotated_stream(stream.into())
             .await
@@ -270,9 +274,13 @@ async fn chat_completions(
         let stream = stream.map(|response| Event::try_from(EventConverter::from(response)));
         let stream = monitor_for_disconnects(stream.boxed(), ctx, inflight).await;
 
-        Ok(Sse::new(stream)
-            .keep_alive(KeepAlive::default())
-            .into_response())
+        let mut sse_stream = Sse::new(stream);
+
+        if let Some(keep_alive) = state.sse_keep_alive {
+            sse_stream = sse_stream.keep_alive(KeepAlive::default().interval(keep_alive));
+        }
+
+        Ok(sse_stream.into_response())
     } else {
         let response = NvCreateChatCompletionResponse::from_annotated_stream(stream.into())
             .await