kserve.rs 32.5 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
GuanLuo's avatar
GuanLuo committed
2
3
4
5
6
7
8
9
10
// SPDX-License-Identifier: Apache-2.0

use std::pin::Pin;
use std::sync::Arc;

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;
11
use crate::http::service::service_v2 as http_service;
GuanLuo's avatar
GuanLuo committed
12
13

use crate::discovery::ModelManager;
14
15
use crate::local_model::runtime_config::ModelRuntimeConfig;
use crate::protocols::tensor::TensorModelConfig;
16
use crate::protocols::tensor::{NvCreateTensorRequest, NvCreateTensorResponse};
GuanLuo's avatar
GuanLuo committed
17
18
19
20
21
22
23
24
use crate::request_template::RequestTemplate;
use anyhow::Result;
use derive_builder::Builder;
use futures::pin_mut;
use tokio::task::JoinHandle;
use tokio_stream::{Stream, StreamExt};
use tokio_util::sync::CancellationToken;

25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
/// Optional HTTP/2 window size configuration from environment variables.
///
/// # Environment Variables
///
/// - `DYN_GRPC_INITIAL_CONNECTION_WINDOW_SIZE`: HTTP/2 connection window size in bytes
/// - `DYN_GRPC_INITIAL_STREAM_WINDOW_SIZE`: HTTP/2 per-stream window size in bytes
///
/// If set, these override tonic defaults. If not set, tonic defaults are used.
#[derive(Debug, Clone, Default)]
pub struct GrpcTuningConfig {
    /// HTTP/2 connection-level flow control window size in bytes.
    /// If None, uses tonic default.
    pub initial_connection_window_size: Option<u32>,

    /// HTTP/2 stream-level flow control window size in bytes.
    /// If None, uses tonic default.
    pub initial_stream_window_size: Option<u32>,
}

impl GrpcTuningConfig {
    /// Create configuration from environment variables.
    ///
    /// Reads `DYN_GRPC_INITIAL_CONNECTION_WINDOW_SIZE` and `DYN_GRPC_INITIAL_STREAM_WINDOW_SIZE`.
    /// If not set, the values remain None and tonic defaults are used.
    pub fn from_env() -> Self {
        let mut config = Self::default();

        if let Ok(val) = std::env::var("DYN_GRPC_INITIAL_CONNECTION_WINDOW_SIZE")
            && let Ok(size) = val.parse::<u32>()
        {
            config.initial_connection_window_size = Some(size);
        }

        if let Ok(val) = std::env::var("DYN_GRPC_INITIAL_STREAM_WINDOW_SIZE")
            && let Ok(size) = val.parse::<u32>()
        {
            config.initial_stream_window_size = Some(size);
        }

        config
    }
}

68
use crate::grpc::service::openai::completion_response_stream;
69
use crate::grpc::service::tensor::{ExtendedNvCreateTensorResponse, tensor_response_stream};
70
use std::convert::{TryFrom, TryInto};
GuanLuo's avatar
GuanLuo committed
71
72
73
74
75
76
77
78
79
80
81
use tonic::{Request, Response, Status, transport::Server};

use crate::protocols::openai::completions::{
    NvCreateCompletionRequest, NvCreateCompletionResponse,
};

pub mod inference {
    tonic::include_proto!("inference");
}
use inference::grpc_inference_service_server::{GrpcInferenceService, GrpcInferenceServiceServer};
use inference::{
82
83
    ModelConfig, ModelConfigRequest, ModelConfigResponse, ModelInferRequest, ModelInferResponse,
    ModelMetadataRequest, ModelMetadataResponse, ModelStreamInferResponse,
GuanLuo's avatar
GuanLuo committed
84
85
};

86
87
use prost::Message;

88
/// gRPC service state - shares metrics with HTTP service for unified metrics collection
GuanLuo's avatar
GuanLuo committed
89
90
91
92
93
pub struct State {
    metrics: Arc<Metrics>,
    manager: Arc<ModelManager>,
}

94
95
96
97
98
99
100
101
102
103
104
105
106
107
#[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>>,
}

GuanLuo's avatar
GuanLuo committed
108
impl State {
109
110
    pub fn builder() -> StateBuilder {
        StateBuilder::default()
GuanLuo's avatar
GuanLuo committed
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
    }

    /// Get the Prometheus [`Metrics`] object which tracks request counts and inflight requests
    pub fn metrics_clone(&self) -> Arc<Metrics> {
        self.metrics.clone()
    }

    pub fn manager(&self) -> &ModelManager {
        Arc::as_ref(&self.manager)
    }

    pub fn manager_clone(&self) -> Arc<ModelManager> {
        self.manager.clone()
    }

126
127
128
    fn is_tensor_model(&self, model: &String) -> bool {
        self.manager.list_tensor_models().contains(model)
    }
GuanLuo's avatar
GuanLuo committed
129
130
}

131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
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())),
        })
    }
}

GuanLuo's avatar
GuanLuo committed
146
147
148
149
150
#[derive(Clone)]
pub struct KserveService {
    // The state we share with every request handler
    state: Arc<State>,

151
152
153
    // HTTP service for metrics endpoint
    http_service: http_service::HttpService,

GuanLuo's avatar
GuanLuo committed
154
155
156
    port: u16,
    host: String,
    request_template: Option<RequestTemplate>,
157
158
159

    // gRPC server tuning configuration
    grpc_tuning: GrpcTuningConfig,
GuanLuo's avatar
GuanLuo committed
160
161
162
163
164
165
166
167
168
169
170
171
172
}

#[derive(Clone, Builder)]
#[builder(pattern = "owned", build_fn(private, name = "build_internal"))]
pub struct KserveServiceConfig {
    #[builder(default = "8787")]
    port: u16,

    #[builder(setter(into), default = "String::from(\"0.0.0.0\")")]
    host: String,

    #[builder(default = "None")]
    request_template: Option<RequestTemplate>,
173
174
175
176
177
178

    #[builder(default = "8788")]
    http_metrics_port: u16,

    #[builder(setter(into), default = "String::from(\"0.0.0.0\")")]
    http_metrics_host: String,
179

180
181
182
    #[builder(default = "None")]
    http_cancel_token: Option<CancellationToken>,

183
184
185
186
    /// gRPC server tuning configuration.
    /// Default: GrpcTuningConfig::from_env() - reads from environment variables with fallback to defaults.
    #[builder(default = "GrpcTuningConfig::from_env()")]
    grpc_tuning: GrpcTuningConfig,
GuanLuo's avatar
GuanLuo committed
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
}

impl KserveService {
    pub fn builder() -> KserveServiceConfigBuilder {
        KserveServiceConfigBuilder::default()
    }

    pub fn state_clone(&self) -> Arc<State> {
        self.state.clone()
    }

    pub fn state(&self) -> &State {
        Arc::as_ref(&self.state)
    }

    pub fn model_manager(&self) -> &ModelManager {
        self.state().manager()
    }

206
207
208
209
    pub fn http_service(&self) -> &http_service::HttpService {
        &self.http_service
    }

GuanLuo's avatar
GuanLuo committed
210
211
212
213
214
215
216
217
218
    pub async fn spawn(&self, cancel_token: CancellationToken) -> JoinHandle<Result<()>> {
        let this = self.clone();
        tokio::spawn(async move { this.run(cancel_token).await })
    }

    pub async fn run(&self, cancel_token: CancellationToken) -> Result<()> {
        let address = format!("{}:{}", self.host, self.port);
        tracing::info!(address, "Starting KServe gRPC service on: {address}");

219
220
221
222
223
224
225
226
227
228
229
230
231
        let tuning = &self.grpc_tuning;

        // Log tuning settings if configured via environment variables
        if tuning.initial_connection_window_size.is_some()
            || tuning.initial_stream_window_size.is_some()
        {
            tracing::info!(
                "gRPC tuning: connection_window={:?}, stream_window={:?}",
                tuning.initial_connection_window_size,
                tuning.initial_stream_window_size
            );
        }

GuanLuo's avatar
GuanLuo committed
232
        let observer = cancel_token.child_token();
233
234
235
236
237
238
239
240
241
242
243
244

        // Build server - only override window sizes if set via env vars
        let mut builder = Server::builder();

        if let Some(size) = tuning.initial_connection_window_size {
            builder = builder.initial_connection_window_size(size);
        }
        if let Some(size) = tuning.initial_stream_window_size {
            builder = builder.initial_stream_window_size(size);
        }

        builder
GuanLuo's avatar
GuanLuo committed
245
246
247
248
249
250
251
252
253
254
255
256
257
            .add_service(GrpcInferenceServiceServer::new(self.clone()))
            .serve_with_shutdown(address.parse()?, observer.cancelled_owned())
            .await
            .inspect_err(|_| cancel_token.cancel())?;

        Ok(())
    }
}

impl KserveServiceConfigBuilder {
    pub fn build(self) -> Result<KserveService, anyhow::Error> {
        let config: KserveServiceConfig = self.build_internal()?;

258
259
260
261
262
        // 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())
263
            .cancel_token(config.http_cancel_token)
264
265
266
267
268
            // 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)
269
            .enable_anthropic_endpoints(false)
270
271
272
273
274
275
276
277
278
            .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()?,
        );
GuanLuo's avatar
GuanLuo committed
279
280
281

        Ok(KserveService {
            state,
282
            http_service,
GuanLuo's avatar
GuanLuo committed
283
284
285
            port: config.port,
            host: config.host,
            request_template: config.request_template,
286
            grpc_tuning: config.grpc_tuning,
GuanLuo's avatar
GuanLuo committed
287
288
289
290
291
292
293
294
295
        })
    }

    pub fn with_request_template(mut self, request_template: Option<RequestTemplate>) -> Self {
        self.request_template = Some(request_template);
        self
    }
}

296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
#[allow(clippy::large_enum_variant)]
enum Config {
    Dynamo(TensorModelConfig),
    Triton(ModelConfig),
}

impl Config {
    fn from_runtime_config(runtime_config: &ModelRuntimeConfig) -> Result<Config, anyhow::Error> {
        if let Some(tensor_model_config) = runtime_config.tensor_model_config.as_ref() {
            if let Some(triton_model_config) = tensor_model_config.triton_model_config.as_ref() {
                let model_config = ModelConfig::decode(triton_model_config.as_slice())?;
                Ok(Config::Triton(model_config))
            } else {
                Ok(Config::Dynamo(tensor_model_config.clone()))
            }
        } else {
            Err(anyhow::anyhow!("no model config is provided"))
        }
    }
}

GuanLuo's avatar
GuanLuo committed
317
318
319
320
321
322
#[tonic::async_trait]
impl GrpcInferenceService for KserveService {
    async fn model_infer(
        &self,
        request: Request<ModelInferRequest>,
    ) -> Result<Response<ModelInferResponse>, Status> {
323
        let model = request.get_ref().model_name.clone();
GuanLuo's avatar
GuanLuo committed
324
325
        let request = request.into_inner();
        let request_id = request.id.clone();
326
327
328

        // [gluo TODO] refactor to reuse code, inference logic is largely the same
        if self.state().is_tensor_model(&model) {
329
            let set_raw_output_contents = !request.raw_input_contents.is_empty();
330
331
332
333
334
            let tensor_request: NvCreateTensorRequest = NvCreateTensorRequest::try_from(request)
                .map_err(|e| Status::invalid_argument(format!("Failed to parse request: {}", e)))?;

            let stream = tensor_response_stream(self.state_clone(), tensor_request, false).await?;

335
336
337
338
339
340
341
342
343
            let tensor_response = ExtendedNvCreateTensorResponse {
                response: NvCreateTensorResponse::from_annotated_stream(stream)
                    .await
                    .map_err(|e| {
                        tracing::error!("Failed to fold completions stream: {:?}", e);
                        Status::internal(format!("Failed to fold completions stream: {}", e))
                    })?,
                set_raw_output_contents,
            };
344
345
346
347
348
349
350
351
352

            let mut reply: ModelInferResponse = tensor_response.try_into().map_err(|e| {
                Status::invalid_argument(format!("Failed to parse response: {}", e))
            })?;
            reply.id = request_id;

            return Ok(Response::new(reply));
        }

353
354
        // [gluo FIXME] check model existence first, otherwise the true error
        // is masked by "Failed to parse request" below.
355
        // Fallback handling by assuming the model is OpenAI Completions model
GuanLuo's avatar
GuanLuo committed
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
        let mut completion_request: NvCreateCompletionRequest = request
            .try_into()
            .map_err(|e| Status::invalid_argument(format!("Failed to parse request: {}", e)))?;

        if completion_request.inner.stream.unwrap_or(false) {
            // return error that streaming is not supported
            return Err(Status::invalid_argument(
                "Streaming is not supported for this endpoint",
            ));
        }

        // Apply template values if present
        if let Some(template) = self.request_template.as_ref() {
            if completion_request.inner.model.is_empty() {
                completion_request.inner.model = template.model.clone();
            }
            if completion_request.inner.temperature.unwrap_or(0.0) == 0.0 {
                completion_request.inner.temperature = Some(template.temperature);
            }
            if completion_request.inner.max_tokens.unwrap_or(0) == 0 {
                completion_request.inner.max_tokens = Some(template.max_completion_tokens);
            }
        }

380
381
        let (stream, parsing_options) =
            completion_response_stream(self.state_clone(), completion_request).await?;
GuanLuo's avatar
GuanLuo committed
382
383
384
385
386
387

        let completion_response =
            NvCreateCompletionResponse::from_annotated_stream(stream, parsing_options)
                .await
                .map_err(|e| {
                    tracing::error!("Failed to fold completions stream: {:?}", e);
388
                    Status::internal(format!("Failed to fold completions stream: {}", e))
GuanLuo's avatar
GuanLuo committed
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
                })?;

        let mut reply: ModelInferResponse = completion_response
            .try_into()
            .map_err(|e| Status::invalid_argument(format!("Failed to parse response: {}", e)))?;
        reply.id = request_id;

        Ok(Response::new(reply))
    }

    type ModelStreamInferStream =
        Pin<Box<dyn Stream<Item = Result<ModelStreamInferResponse, Status>> + Send + 'static>>;

    async fn model_stream_infer(
        &self,
        request: Request<tonic::Streaming<ModelInferRequest>>,
    ) -> Result<Response<Self::ModelStreamInferStream>, Status> {
        let mut request_stream = request.into_inner();
        let state = self.state_clone();
        let template = self.request_template.clone();
        let output = async_stream::try_stream! {
            // [gluo FIXME] should be able to demux request / response streaming
            // await requests in a separate task until cancellation / completion,
            // and passing AsyncEngineStream for each request to the response stream
            // which will be collectively polling.
            while let Some(request) = request_stream.next().await {
415
                let request = match request {
GuanLuo's avatar
GuanLuo committed
416
417
418
419
420
421
422
423
424
                    Err(e) => {
                        tracing::error!("Unexpected gRPC failed to read request: {}", e);
                        yield ModelStreamInferResponse {
                            error_message: e.to_string(),
                            infer_response: None
                        };
                        continue;
                    }
                    Ok(request) => {
425
                        request
GuanLuo's avatar
GuanLuo committed
426
427
428
                    }
                };

429
430
431
432
433
434
                let model = request.model_name.clone();

                // [gluo TODO] refactor to reuse code, inference logic is largely the same
                if state.is_tensor_model(&model) {
                    // Must keep track of 'request_id' which will be returned in corresponding response
                    let request_id = request.id.clone();
435
                    let set_raw_output_contents = !request.raw_input_contents.is_empty();
436
437
438
439
440
441
442
                    let tensor_request: NvCreateTensorRequest = request.try_into().map_err(|e| {
                        Status::invalid_argument(format!("Failed to parse request: {}", e))
                    })?;

                    let stream = tensor_response_stream(state.clone(), tensor_request, true).await?;

                    pin_mut!(stream);
443
444
445
446
447
448
449
450
451
452
453
                    while let Some(delta) = stream.next().await {
                        let response = match delta.ok() {
                            Err(e) => {
                                yield ModelStreamInferResponse {
                                    error_message: e.to_string(),
                                    infer_response: None
                                };
                                continue;
                            }
                            Ok(response) => response,
                        };
454
455
                        match response.data {
                            Some(data) => {
456
457
458
                                let data = ExtendedNvCreateTensorResponse {response: data,
                                    set_raw_output_contents,
                                };
459
460
461
                                let mut reply = ModelStreamInferResponse::try_from(data).map_err(|e| {
                                    Status::invalid_argument(format!("Failed to parse response: {}", e))
                                })?;
462
463
                                if let Some(infer_response) = reply.infer_response.as_mut() {
                                    infer_response.id = request_id.clone();
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
                                }
                                yield reply;
                            },
                            None => {
                                // Skip if no data is present, the response is for annotation
                            },
                        }
                    }
                    continue;
                }

                // Fallback handling by assuming the model is OpenAI Completions model
                // Must keep track of 'request_id' which will be returned in corresponding response
                let request_id = request.id.clone();
                let mut completion_request: NvCreateCompletionRequest = request.try_into().map_err(|e| {
                    Status::invalid_argument(format!("Failed to parse request: {}", e))
                })?;

GuanLuo's avatar
GuanLuo committed
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
                // Apply template values if present
                if let Some(template) = &template {
                    if completion_request.inner.model.is_empty() {
                        completion_request.inner.model = template.model.clone();
                    }
                    if completion_request.inner.temperature.unwrap_or(0.0) == 0.0 {
                        completion_request.inner.temperature = Some(template.temperature);
                    }
                    if completion_request.inner.max_tokens.unwrap_or(0) == 0 {
                        completion_request.inner.max_tokens = Some(template.max_completion_tokens);
                    }
                }

                let streaming = completion_request.inner.stream.unwrap_or(false);

497
                let (stream, parsing_options) = completion_response_stream(state.clone(), completion_request).await?;
GuanLuo's avatar
GuanLuo committed
498
499
500

                if streaming {
                    pin_mut!(stream);
501
502
503
504
505
506
507
508
509
510
511
                    while let Some(delta) = stream.next().await {
                        let response = match delta.ok() {
                            Err(e) => {
                                yield ModelStreamInferResponse {
                                    error_message: e.to_string(),
                                    infer_response: None
                                };
                                continue;
                            }
                            Ok(response) => response,
                        };
GuanLuo's avatar
GuanLuo committed
512
513
514
515
516
                        match response.data {
                            Some(data) => {
                                let mut reply = ModelStreamInferResponse::try_from(data).map_err(|e| {
                                    Status::invalid_argument(format!("Failed to parse response: {}", e))
                                })?;
517
518
                                if let Some(infer_response) = reply.infer_response.as_mut() {
                                    infer_response.id = request_id.clone();
GuanLuo's avatar
GuanLuo committed
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
                                }
                                yield reply;
                            },
                            None => {
                                // Skip if no data is present, the response is for annotation
                            },
                        }
                    }
                } else {
                    let completion_response = NvCreateCompletionResponse::from_annotated_stream(stream, parsing_options)
                        .await
                        .map_err(|e| {
                            tracing::error!(
                                "Failed to fold completions stream: {:?}",
                                e
                            );
535
                            Status::internal(format!("Failed to fold completions stream: {}", e))
GuanLuo's avatar
GuanLuo committed
536
537
538
539
540
                        })?;

                    let mut response: ModelStreamInferResponse = completion_response.try_into().map_err(|e| {
                        Status::invalid_argument(format!("Failed to parse response: {}", e))
                    })?;
541
542
                    if let Some(infer_response) = response.infer_response.as_mut() {
                        infer_response.id = request_id.clone();
GuanLuo's avatar
GuanLuo committed
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
                    }
                    yield response;
                }
            }
        };

        Ok(Response::new(
            Box::pin(output) as Self::ModelStreamInferStream
        ))
    }

    async fn model_metadata(
        &self,
        request: Request<ModelMetadataRequest>,
    ) -> Result<Response<ModelMetadataResponse>, Status> {
558
        let cards = self.state.manager().get_model_cards();
GuanLuo's avatar
GuanLuo committed
559
        let request_model_name = &request.into_inner().name;
560
        if let Some(card) = cards
561
            .into_iter()
562
            .find(|card| request_model_name == &card.display_name)
563
        {
564
            if card.model_type.supports_tensor() {
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
                let config = Config::from_runtime_config(&card.runtime_config).map_err(|e| {
                    Status::invalid_argument(format!(
                        "Model '{}' has type Tensor but: {}",
                        request_model_name, e
                    ))
                })?;
                match config {
                    Config::Triton(model_config) => {
                        return Ok(Response::new(ModelMetadataResponse {
                            name: model_config.name,
                            versions: vec!["1".to_string()],
                            platform: model_config.platform,
                            inputs: model_config
                                .input
                                .iter()
                                .map(|input| inference::model_metadata_response::TensorMetadata {
                                    name: input.name.clone(),
                                    datatype: match inference::DataType::try_from(input.data_type) {
                                        Ok(dt) => dt.as_str_name().to_string(),
                                        Err(_) => "TYPE_INVALID".to_string(),
                                    },
                                    shape: input.dims.clone(),
                                })
                                .collect(),
                            outputs: model_config
                                .output
                                .iter()
                                .map(
                                    |output| inference::model_metadata_response::TensorMetadata {
                                        name: output.name.clone(),
                                        datatype: match inference::DataType::try_from(
                                            output.data_type,
                                        ) {
                                            Ok(dt) => dt.as_str_name().to_string(),
                                            Err(_) => "TYPE_INVALID".to_string(),
                                        },
                                        shape: output.dims.clone(),
                                    },
                                )
                                .collect(),
                        }));
                    }
                    Config::Dynamo(model_config) => {
                        return Ok(Response::new(ModelMetadataResponse {
                            name: model_config.name.clone(),
                            versions: vec!["1".to_string()],
                            platform: "dynamo".to_string(),
                            inputs: model_config
                                .inputs
                                .iter()
                                .map(|input| inference::model_metadata_response::TensorMetadata {
                                    name: input.name.clone(),
                                    datatype: input.data_type.to_string(),
                                    shape: input.shape.clone(),
                                })
                                .collect(),
                            outputs: model_config
                                .outputs
                                .iter()
                                .map(
                                    |output| inference::model_metadata_response::TensorMetadata {
                                        name: output.name.clone(),
                                        datatype: output.data_type.to_string(),
                                        shape: output.shape.clone(),
                                    },
                                )
                                .collect(),
                        }));
                    }
634
                }
635
            } else if card.model_type.supports_completions() {
636
                return Ok(Response::new(ModelMetadataResponse {
637
                    name: card.display_name,
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
                    versions: vec!["1".to_string()],
                    platform: "dynamo".to_string(),
                    inputs: vec![
                        inference::model_metadata_response::TensorMetadata {
                            name: "text_input".to_string(),
                            datatype: "BYTES".to_string(),
                            shape: vec![1],
                        },
                        inference::model_metadata_response::TensorMetadata {
                            name: "streaming".to_string(),
                            datatype: "BOOL".to_string(),
                            shape: vec![1],
                        },
                    ],
                    outputs: vec![
                        inference::model_metadata_response::TensorMetadata {
                            name: "text_output".to_string(),
                            datatype: "BYTES".to_string(),
                            shape: vec![-1],
                        },
                        inference::model_metadata_response::TensorMetadata {
                            name: "finish_reason".to_string(),
                            datatype: "BYTES".to_string(),
                            shape: vec![-1],
                        },
                    ],
                }));
            }
GuanLuo's avatar
GuanLuo committed
666
667
668
669
670
671
672
673
674
675
676
        }
        Err(Status::not_found(format!(
            "Model '{}' not found",
            request_model_name
        )))
    }

    async fn model_config(
        &self,
        request: Request<ModelConfigRequest>,
    ) -> Result<Response<ModelConfigResponse>, Status> {
677
        let cards = self.state.manager().get_model_cards();
GuanLuo's avatar
GuanLuo committed
678
        let request_model_name = &request.into_inner().name;
679
        if let Some(card) = cards
680
            .into_iter()
681
            .find(|card| request_model_name == &card.display_name)
682
        {
683
            if card.model_type.supports_tensor() {
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
                let config = Config::from_runtime_config(&card.runtime_config).map_err(|e| {
                    Status::invalid_argument(format!(
                        "Model '{}' has type Tensor but: {}",
                        request_model_name, e
                    ))
                })?;
                match config {
                    Config::Triton(model_config) => {
                        return Ok(Response::new(ModelConfigResponse {
                            config: Some(model_config),
                        }));
                    }
                    Config::Dynamo(tensor_model_config) => {
                        let model_config = ModelConfig {
                            name: tensor_model_config.name.clone(),
                            platform: "dynamo".to_string(),
                            backend: "dynamo".to_string(),
                            input: tensor_model_config
                                .inputs
                                .iter()
                                .map(|input| ModelInput {
                                    name: input.name.clone(),
                                    data_type: input.data_type.to_kserve(),
                                    dims: input.shape.clone(),
                                    ..Default::default()
                                })
                                .collect(),
                            output: tensor_model_config
                                .outputs
                                .iter()
                                .map(|output| ModelOutput {
                                    name: output.name.clone(),
                                    data_type: output.data_type.to_kserve(),
                                    dims: output.shape.clone(),
                                    ..Default::default()
                                })
                                .collect(),
                            ..Default::default()
                        };
                        return Ok(Response::new(ModelConfigResponse {
                            config: Some(model_config.clone()),
                        }));
                    }
727
                }
728
            } else if card.model_type.supports_completions() {
729
                let config = ModelConfig {
730
                    name: card.display_name,
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
                    platform: "dynamo".to_string(),
                    backend: "dynamo".to_string(),
                    input: vec![
                        ModelInput {
                            name: "text_input".to_string(),
                            data_type: DataType::TypeString as i32,
                            dims: vec![1],
                            ..Default::default()
                        },
                        ModelInput {
                            name: "streaming".to_string(),
                            data_type: DataType::TypeBool as i32,
                            dims: vec![1],
                            optional: true,
                            ..Default::default()
                        },
                    ],
                    output: vec![
                        ModelOutput {
                            name: "text_output".to_string(),
                            data_type: DataType::TypeString as i32,
                            dims: vec![-1],
                            ..Default::default()
                        },
                        ModelOutput {
                            name: "finish_reason".to_string(),
                            data_type: DataType::TypeString as i32,
                            dims: vec![-1],
                            ..Default::default()
                        },
                    ],
                    ..Default::default()
                };
                return Ok(Response::new(ModelConfigResponse {
                    config: Some(config),
                }));
            }
GuanLuo's avatar
GuanLuo committed
768
769
770
771
772
773
        }
        Err(Status::not_found(format!(
            "Model '{}' not found",
            request_model_name
        )))
    }
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807

    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,
        }))
    }
GuanLuo's avatar
GuanLuo committed
808
}