inference 0.3.0

A crate for managing the machine learning inference process
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

//! Client for interfacing with the Triton inference server over gRPC

use tonic::{Request, Status};
use tracing::{event, instrument, Level};

use crate::inference::{ModelConfigRequest, ModelMetadataRequest, ModelStatisticsRequest};
use crate::TritonClient;

use crate::inference::{
    ModelConfigResponse, ModelMetadataResponse, ModelReadyRequest, ModelStatisticsResponse,
};

impl TritonClient {
    /// Is a specific model & version ready to accept inference requests?
    ///
    /// RPC ModelReady
    ///
    /// Returns an error if the model doesn't exist on the server or the request fails
    #[instrument]
    pub async fn is_model_ready(&self, name: &str, version: &str) -> Result<bool, Status> {
        let message = ModelReadyRequest {
            name: name.to_owned(),
            version: version.to_owned(),
        };
        let request = Request::new(message);

        match self.client().model_ready(request).await {
            Ok(r) => {
                let model_readiness = r.into_inner().ready;
                if model_readiness {
                    event!(Level::INFO, "model is ready");
                } else {
                    event!(Level::INFO, "model is not ready");
                }
                Ok(model_readiness)
            }
            Err(s) => {
                event!(Level::ERROR, "failed to get model readiness, {s}");
                Err(s)
            }
        }
    }

    /// Get a model's metadata
    ///
    /// RPC ModelMetadata
    #[instrument]
    pub async fn model_metadata(
        &self,
        name: &str,
        version: &str,
    ) -> Result<ModelMetadataResponse, Status> {
        let msg = ModelMetadataRequest {
            name: name.to_owned(),
            version: version.to_owned(),
        };
        let request = Request::new(msg);
        match self.client().model_metadata(request).await {
            Ok(r) => {
                let model_metadata = r.into_inner();
                event!(Level::INFO, ?model_metadata);
                Ok(model_metadata)
            }
            Err(s) => {
                event!(Level::ERROR, "failed to get model metadata, {s}");
                Err(s)
            }
        }
    }

    /// Get model configuration info
    ///
    /// RPC ModelConfig
    #[instrument]
    pub async fn model_config(
        &self,
        name: &str,
        version: &str,
    ) -> Result<ModelConfigResponse, Status> {
        let msg = ModelConfigRequest {
            name: name.to_owned(),
            version: version.to_owned(),
        };
        let request = Request::new(msg);
        match self.client().model_config(request).await {
            Ok(r) => {
                let model_config = r.into_inner();
                event!(Level::INFO, ?model_config);
                Ok(model_config)
            }
            Err(s) => {
                event!(Level::ERROR, "failed to get model configuration, {s}");
                Err(s)
            }
        }
    }

    /// Get the cumulative inference statistics for a model
    ///
    /// RPC ModelStatistics
    #[instrument]
    pub async fn model_statistics(
        &self,
        name: &str,
        version: &str,
    ) -> Result<ModelStatisticsResponse, Status> {
        let msg = ModelStatisticsRequest {
            name: name.to_owned(),
            version: version.to_owned(),
        };
        let request = Request::new(msg);
        match self.client().model_statistics(request).await {
            Ok(r) => {
                let model_statistics = r.into_inner();
                event!(Level::INFO, ?model_statistics);
                Ok(model_statistics)
            }
            Err(s) => {
                event!(Level::ERROR, "failed to get model statistics, {s}");
                Err(s)
            }
        }
    }
}