chore: cherry-pick KServe readiness (#4708) and metrics (#4400) feature commits (#4730)

Signed-off-by: Guan Luo <gluo@nvidia.com>
Signed-off-by: GuanLuo <41310872+GuanLuo@users.noreply.github.com>
Co-authored-by: Ryan McCormick <rmccormick@nvidia.com>
Co-authored-by: Tanmay Verma <tanmayv@nvidia.com>

Author: 3 people

lib/llm/src/entrypoint/input/grpc.rs

@@ -152,9 +152,18 @@ pub async fn run(
             grpc_service
         }
     };
-    grpc_service
-        .run(distributed_runtime.primary_token())
-        .await?;
+
+    // Run both HTTP (for metrics) and gRPC servers concurrently
+    let http_service = grpc_service.http_service().clone();
+    let shutdown_token = distributed_runtime.primary_token();
+
+    // Wait for both servers to complete, propagating the first error if any occurs
+    // Both tasks should run indefinitely until cancelled by the shutdown token
+    tokio::try_join!(
+        grpc_service.run(shutdown_token.clone()),
+        http_service.run(shutdown_token)
+    )?;
+
     distributed_runtime.shutdown(); // Cancel primary token
     Ok(())
 }

lib/llm/src/grpc/protos/kserve.proto

@@ -16,6 +16,27 @@ import "model_config.proto";
 //@@
 service GRPCInferenceService
 {
+  //@@  .. cpp:var:: rpc ServerLive(ServerLiveRequest) returns
+  //@@       (ServerLiveResponse)
+  //@@
+  //@@     Check liveness of the inference server.
+  //@@
+  rpc ServerLive(ServerLiveRequest) returns (ServerLiveResponse) {}
+
+  //@@  .. cpp:var:: rpc ServerReady(ServerReadyRequest) returns
+  //@@       (ServerReadyResponse)
+  //@@
+  //@@     Check readiness of the inference server.
+  //@@
+  rpc ServerReady(ServerReadyRequest) returns (ServerReadyResponse) {}
+
+  //@@  .. cpp:var:: rpc ModelReady(ModelReadyRequest) returns
+  //@@       (ModelReadyResponse)
+  //@@
+  //@@     Check readiness of a model in the inference server.
+  //@@
+  rpc ModelReady(ModelReadyRequest) returns (ModelReadyResponse) {}
+
   //@@  .. cpp:var:: rpc ModelMetadata(ModelMetadataRequest) returns
   //@@       (ModelMetadataResponse)
   //@@
@@ -45,6 +66,89 @@ service GRPCInferenceService
   rpc ModelConfig(ModelConfigRequest) returns (ModelConfigResponse) {}
 }
 
+//@@
+//@@.. cpp:var:: message ServerLiveRequest
+//@@
+//@@   Request message for ServerLive.
+//@@
+message ServerLiveRequest {}
+
+//@@
+//@@.. cpp:var:: message ServerLiveResponse
+//@@
+//@@   Response message for ServerLive.
+//@@
+message ServerLiveResponse
+{
+  //@@
+  //@@  .. cpp:var:: bool live
+  //@@
+  //@@     True if the inference server is live, false if not live.
+  //@@
+  bool live = 1;
+}
+
+//@@
+//@@.. cpp:var:: message ServerReadyRequest
+//@@
+//@@   Request message for ServerReady.
+//@@
+message ServerReadyRequest {}
+
+//@@
+//@@.. cpp:var:: message ServerReadyResponse
+//@@
+//@@   Response message for ServerReady.
+//@@
+message ServerReadyResponse
+{
+  //@@
+  //@@  .. cpp:var:: bool ready
+  //@@
+  //@@     True if the inference server is ready, false if not ready. The server
+  //@@     is considered ready if it has any registered models, since models
+  //@@     can freely be registered and unregistered at runtime.
+  //@@
+  bool ready = 1;
+}
+
+//@@
+//@@.. cpp:var:: message ModelReadyRequest
+//@@
+//@@   Request message for ModelReady.
+//@@
+message ModelReadyRequest
+{
+  //@@
+  //@@  .. cpp:var:: string name
+  //@@
+  //@@     The name of the model to check for readiness.
+  //@@
+  string name = 1;
+
+  //@@  .. cpp:var:: string version
+  //@@
+  //@@     The version of the model to check for readiness. If not given the
+  //@@     server will choose a version based on the model and internal policy.
+  //@@
+  string version = 2;
+}
+
+//@@
+//@@.. cpp:var:: message ModelReadyResponse
+//@@
+//@@   Response message for ModelReady.
+//@@
+message ModelReadyResponse
+{
+  //@@
+  //@@  .. cpp:var:: bool ready
+  //@@
+  //@@     True if the model is ready, false if not ready.
+  //@@
+  bool ready = 1;
+}
+
 //@@
 //@@.. cpp:var:: message ModelMetadataRequest
 //@@

lib/llm/src/grpc/service/kserve.rs

@@ -8,7 +8,7 @@ use crate::grpc::service::kserve::inference::DataType;
 use crate::grpc::service::kserve::inference::ModelInput;
 use crate::grpc::service::kserve::inference::ModelOutput;
 use crate::http::service::Metrics;
-use crate::http::service::metrics;
+use crate::http::service::service_v2 as http_service;
 
 use crate::discovery::ModelManager;
 use crate::local_model::runtime_config::ModelRuntimeConfig;
@@ -42,20 +42,29 @@ use inference::{
 
 use prost::Message;
 
-/// [gluo TODO] 'metrics' are for HTTP service and there is HTTP endpoint
-/// for it as part of HTTP service. Should we always start HTTP service up
-/// for non-inference?
+/// gRPC service state - shares metrics with HTTP service for unified metrics collection
 pub struct State {
     metrics: Arc<Metrics>,
     manager: Arc<ModelManager>,
 }
 
+#[derive(Default, Builder)]
+#[builder(
+    pattern = "owned",
+    build_fn(private, name = "build_internal"),
+    name = "StateBuilder",
+    vis = "pub"
+)]
+pub(crate) struct StateConfig {
+    #[builder(default, setter(strip_option))]
+    metrics: Option<Arc<Metrics>>,
+    #[builder(default, setter(strip_option))]
+    manager: Option<Arc<ModelManager>>,
+}
+
 impl State {
-    pub fn new(manager: Arc<ModelManager>) -> Self {
-        Self {
-            manager,
-            metrics: Arc::new(Metrics::default()),
-        }
+    pub fn builder() -> StateBuilder {
+        StateBuilder::default()
     }
 
     /// Get the Prometheus [`Metrics`] object which tracks request counts and inflight requests
@@ -76,11 +85,29 @@ impl State {
     }
 }
 
+impl StateBuilder {
+    pub fn build(self) -> Result<State, anyhow::Error> {
+        let config = self.build_internal()?;
+
+        Ok(State {
+            manager: config
+                .manager
+                .unwrap_or_else(|| Arc::new(ModelManager::new())),
+            metrics: config
+                .metrics
+                .unwrap_or_else(|| Arc::new(Metrics::default())),
+        })
+    }
+}
+
 #[derive(Clone)]
 pub struct KserveService {
     // The state we share with every request handler
     state: Arc<State>,
 
+    // HTTP service for metrics endpoint
+    http_service: http_service::HttpService,
+
     port: u16,
     host: String,
     request_template: Option<RequestTemplate>,
@@ -97,6 +124,12 @@ pub struct KserveServiceConfig {
 
     #[builder(default = "None")]
     request_template: Option<RequestTemplate>,
+
+    #[builder(default = "8788")]
+    http_metrics_port: u16,
+
+    #[builder(setter(into), default = "String::from(\"0.0.0.0\")")]
+    http_metrics_host: String,
 }
 
 impl KserveService {
@@ -116,6 +149,10 @@ impl KserveService {
         self.state().manager()
     }
 
+    pub fn http_service(&self) -> &http_service::HttpService {
+        &self.http_service
+    }
+
     pub async fn spawn(&self, cancel_token: CancellationToken) -> JoinHandle<Result<()>> {
         let this = self.clone();
         tokio::spawn(async move { this.run(cancel_token).await })
@@ -140,15 +177,29 @@ impl KserveServiceConfigBuilder {
     pub fn build(self) -> Result<KserveService, anyhow::Error> {
         let config: KserveServiceConfig = self.build_internal()?;
 
-        let model_manager = Arc::new(ModelManager::new());
-        let state = Arc::new(State::new(model_manager));
-
-        // enable prometheus metrics
-        let registry = metrics::Registry::new();
-        state.metrics_clone().register(&registry)?;
+        // Create HTTP service with only non-inference endpoints (metrics, health, models list)
+        // This provides the metrics endpoint and shared metrics object
+        let http_service = http_service::HttpService::builder()
+            .port(config.http_metrics_port)
+            .host(config.http_metrics_host.clone())
+            // Disable all inference endpoints - only use for metrics/health
+            .enable_chat_endpoints(false)
+            .enable_cmpl_endpoints(false)
+            .enable_embeddings_endpoints(false)
+            .enable_responses_endpoints(false)
+            .build()?;
+
+        // Share the HTTP service's model manager and metrics object with gRPC state
+        let state = Arc::new(
+            State::builder()
+                .manager(http_service.state().manager_clone())
+                .metrics(http_service.state().metrics_clone())
+                .build()?,
+        );
 
         Ok(KserveService {
             state,
+            http_service,
             port: config.port,
             host: config.host,
             request_template: config.request_template,
@@ -624,4 +675,38 @@ impl GrpcInferenceService for KserveService {
             request_model_name
         )))
     }
+
+    async fn server_live(
+        &self,
+        _request: Request<inference::ServerLiveRequest>,
+    ) -> Result<Response<inference::ServerLiveResponse>, Status> {
+        // server is live if we can respond
+        Ok(Response::new(inference::ServerLiveResponse { live: true }))
+    }
+
+    async fn server_ready(
+        &self,
+        _request: Request<inference::ServerReadyRequest>,
+    ) -> Result<Response<inference::ServerReadyResponse>, Status> {
+        let has_models = !self.state.manager().get_model_cards().is_empty();
+        Ok(Response::new(inference::ServerReadyResponse {
+            ready: has_models,
+        }))
+    }
+
+    async fn model_ready(
+        &self,
+        request: Request<inference::ModelReadyRequest>,
+    ) -> Result<Response<inference::ModelReadyResponse>, Status> {
+        let request_model_name = &request.into_inner().name;
+        let is_ready = self
+            .state
+            .manager()
+            .get_model_cards()
+            .into_iter()
+            .any(|card| request_model_name == &card.display_name);
+        Ok(Response::new(inference::ModelReadyResponse {
+            ready: is_ready,
+        }))
+    }
 }

lib/llm/src/grpc/service/tensor.rs

@@ -15,11 +15,11 @@ use crate::types::Annotated;
 use super::kserve;
 
 // [gluo NOTE] These are common utilities that should be shared between frontends
+use crate::http::service::metrics::InflightGuard;
 use crate::http::service::{
     disconnect::{ConnectionHandle, create_connection_monitor},
-    metrics::{Endpoint, ResponseMetricCollector},
+    metrics::{Endpoint, process_response_and_observe_metrics},
 };
-use crate::{http::service::metrics::InflightGuard, preprocessor::LLMMetricAnnotation};
 
 use crate::protocols::tensor;
 use crate::protocols::tensor::{
@@ -76,6 +76,8 @@ pub async fn tensor_response_stream(
         .get_tensor_engine(model)
         .map_err(|_| Status::not_found("model not found"))?;
 
+    let http_queue_guard = state.metrics_clone().create_http_queue_guard(model);
+
     let inflight_guard =
         state
             .metrics_clone()
@@ -115,9 +117,15 @@ pub async fn tensor_response_stream(
     // apply any annotations to the front of the stream
     let stream = stream::iter(annotations).chain(stream);
 
-    // Tap on the stream to collect response metrics
+    // Tap on the stream to collect response metrics and handle http_queue_guard
+    let mut http_queue_guard = Some(http_queue_guard);
     let stream = stream.inspect(move |response| {
-        process_metrics_only(response, &mut response_collector);
+        // Calls observe_response() on each token - drops http_queue_guard on first token
+        process_response_and_observe_metrics(
+            response,
+            &mut response_collector,
+            &mut http_queue_guard,
+        );
     });
 
     let stream = grpc_monitor_for_disconnects(stream, ctx, inflight_guard, stream_handle);
@@ -170,17 +178,6 @@ pub fn grpc_monitor_for_disconnects<T>(
     }
 }
 
-fn process_metrics_only<T>(
-    annotated: &Annotated<T>,
-    response_collector: &mut ResponseMetricCollector,
-) {
-    // update metrics
-    if let Ok(Some(metrics)) = LLMMetricAnnotation::from_annotation(annotated) {
-        response_collector.observe_current_osl(metrics.output_tokens);
-        response_collector.observe_response(metrics.input_tokens, metrics.chunk_tokens);
-    }
-}
-
 /// Get the request ID from a primary source, or lastly create a new one if not present
 fn get_or_create_request_id(primary: Option<&str>) -> String {
     // Try to get the request ID from the primary source