client.rs 8.57 KB
Newer Older
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1
/// Single shard Client
OlivierDehaene's avatar
OlivierDehaene committed
2
3
4
5
use crate::v2::pb;
use crate::{ClientError, Result};

use crate::WARMUP_IMAGE_BASE64;
6
use grpc_metadata::InjectTelemetryContext;
OlivierDehaene's avatar
OlivierDehaene committed
7
8
use pb::generate::v2::text_generation_service_client::TextGenerationServiceClient;
use pb::generate::v2::*;
9
use std::cmp::min;
10
use std::time::Duration;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
11
use tonic::transport::{Channel, Uri};
12
use tracing::instrument;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
13

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
14
/// Text Generation Inference gRPC client
15
#[derive(Debug, Clone)]
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
16
pub struct Client {
Olivier Dehaene's avatar
Olivier Dehaene committed
17
    stub: TextGenerationServiceClient<Channel>,
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
18
19
20
}

impl Client {
Olivier Dehaene's avatar
Olivier Dehaene committed
21
22
23
    /// Returns a client connected to the given url
    pub async fn connect(uri: Uri) -> Result<Self> {
        let channel = Channel::builder(uri).connect().await?;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
24

Olivier Dehaene's avatar
Olivier Dehaene committed
25
26
27
        Ok(Self {
            stub: TextGenerationServiceClient::new(channel),
        })
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
28
29
    }

Olivier Dehaene's avatar
Olivier Dehaene committed
30
31
    /// Returns a client connected to the given unix socket
    pub async fn connect_uds(path: String) -> Result<Self> {
Olivier Dehaene's avatar
Olivier Dehaene committed
32
        let channel = Channel::from_shared("http://[::]:50051".to_string())
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
33
34
35
36
            .unwrap()
            .connect_with_connector(tower::service_fn(move |_: Uri| {
                tokio::net::UnixStream::connect(path.clone())
            }))
Olivier Dehaene's avatar
Olivier Dehaene committed
37
            .await?;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
38

Olivier Dehaene's avatar
Olivier Dehaene committed
39
40
41
        Ok(Self {
            stub: TextGenerationServiceClient::new(channel),
        })
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
42
43
    }

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
44
    /// Returns a list of uris or unix sockets of all shards
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
45
46
    #[instrument(skip(self))]
    pub async fn service_discovery(&mut self) -> Result<Vec<String>> {
47
        let request = tonic::Request::new(ServiceDiscoveryRequest {}).inject_context();
OlivierDehaene's avatar
OlivierDehaene committed
48
49
50
        let response = self.stub.service_discovery(request).await.map_err(|_| {
            ClientError::Connection("Server does not support v2 interface".to_string())
        })?;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
51
52
53
54
        let urls = response
            .into_inner()
            .urls
            .into_iter()
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
55
            // Remove unix socket prefix
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
56
57
58
59
60
61
62
63
            .map(|url| match url.strip_prefix("unix://") {
                None => url,
                Some(stripped_url) => stripped_url.to_string(),
            })
            .collect();
        Ok(urls)
    }

64
65
66
67
68
69
70
71
    /// Get model info
    #[instrument(skip(self))]
    pub async fn info(&mut self) -> Result<InfoResponse> {
        let request = tonic::Request::new(InfoRequest {}).inject_context();
        let response = self.stub.info(request).await?.into_inner();
        Ok(response)
    }

72
73
74
75
76
77
78
79
    /// Get model health
    #[instrument(skip(self))]
    pub async fn health(&mut self) -> Result<HealthResponse> {
        let request = tonic::Request::new(HealthRequest {}).inject_context();
        let response = self.stub.health(request).await?.into_inner();
        Ok(response)
    }

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
80
    /// Clear the past generations cache
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
81
    #[instrument(skip(self))]
82
83
    pub async fn clear_cache(&mut self, batch_id: Option<u64>) -> Result<()> {
        let request = tonic::Request::new(ClearCacheRequest { id: batch_id }).inject_context();
84
        self.stub.clear_cache(request).await?;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
85
86
87
        Ok(())
    }

88
89
90
91
92
    /// Filter a cached batch
    #[instrument(skip(self))]
    pub async fn filter_batch(
        &mut self,
        batch_id: u64,
93
94
        request_ids: Vec<u64>,
    ) -> Result<Option<CachedBatch>> {
95
96
        let request = tonic::Request::new(FilterBatchRequest {
            batch_id,
97
            request_ids,
98
99
100
101
102
103
        })
        .inject_context();
        let filtered_batch = self.stub.filter_batch(request).await?.into_inner();
        Ok(filtered_batch.batch)
    }

104
105
106
    /// Warmup on a max size batch
    ///
    /// Returns the maximum amount of tokens supported by the hardware
107
    #[instrument(skip_all)]
108
109
110
111
    pub async fn warmup(
        &mut self,
        max_input_length: u32,
        max_prefill_tokens: u32,
OlivierDehaene's avatar
OlivierDehaene committed
112
        max_total_tokens: u32,
113
        max_batch_size: Option<usize>,
114
    ) -> Result<Option<u32>> {
115
116
117
118
        let mut n_tokens = 0;
        let mut requests = Vec::new();
        // Create requests
        while n_tokens < max_prefill_tokens {
OlivierDehaene's avatar
OlivierDehaene committed
119
            let truncate = min(max_input_length, max_prefill_tokens - n_tokens);
120
121
122

            let mut inputs = String::new();
            inputs.push_str(&"_test ".to_string().repeat(max_input_length as usize));
Nicolas Patry's avatar
Nicolas Patry committed
123
124
125
            if n_tokens == 0 {
                // 1 request is enough to test vision heads.
                // Sending images on other queries messes up easily with truncation.
126
127
128
                inputs.push_str(&format!(
                    "![](data:image/jpeg;base64,{WARMUP_IMAGE_BASE64})",
                ));
Nicolas Patry's avatar
Nicolas Patry committed
129
            }
130

131
132
            requests.push(Request {
                id: 0,
133
                inputs,
134
                // We truncate the input on the server side to be sure that it has the correct size
OlivierDehaene's avatar
OlivierDehaene committed
135
                truncate,
136
137
138
139
140
141
142
143
144
                // Set sampling parameters to also take these ops into account in the max memory
                parameters: Some(NextTokenChooserParameters {
                    temperature: 0.9,
                    top_k: 10,
                    top_p: 0.9,
                    typical_p: 0.9,
                    do_sample: false,
                    seed: 0,
                    repetition_penalty: 1.2,
145
                    frequency_penalty: 0.1,
146
                    watermark: true,
drbh's avatar
drbh committed
147
148
                    grammar: String::new(),
                    grammar_type: GrammarType::None as i32,
149
150
                }),
                stopping_parameters: Some(StoppingCriteriaParameters {
OlivierDehaene's avatar
OlivierDehaene committed
151
                    max_new_tokens: max_total_tokens - truncate,
152
                    stop_sequences: vec![],
OlivierDehaene's avatar
OlivierDehaene committed
153
                    ignore_eos_token: true,
154
155
                }),
                prefill_logprobs: true,
Nicolas Patry's avatar
Nicolas Patry committed
156
                top_n_tokens: 20,
157
158
            });
            n_tokens += max_input_length;
159
160
161
162
163

            // Check max_batch_size
            if Some(requests.len()) == max_batch_size {
                break;
            }
164
165
166
167
168
169
170
171
172
        }

        let batch = Batch {
            id: 0,
            size: requests.len() as u32,
            requests,
            max_tokens: 0,
        };

173
174
175
176
177
178
179
        let request = tonic::Request::new(WarmupRequest {
            batch: Some(batch),
            max_input_length,
            max_prefill_tokens,
            max_total_tokens,
        })
        .inject_context();
180
181
        let response = self.stub.warmup(request).await?.into_inner();
        Ok(response.max_supported_total_tokens)
182
183
    }

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
184
185
    /// Generate one token for each request in the given batch
    ///
186
    /// Returns Generation for each request in batch
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
187
    /// and the next cached batch
188
    #[instrument(skip_all, fields(id = &batch.id, size = &batch.size))]
189
190
191
    pub async fn prefill(
        &mut self,
        batch: Batch,
192
    ) -> Result<(Vec<Generation>, Option<CachedBatch>, PrefillTimings)> {
193
194
        let request = tonic::Request::new(PrefillRequest { batch: Some(batch) }).inject_context();
        let response = self.stub.prefill(request).await?.into_inner();
195
196
197
198
199
        Ok((
            response.generations,
            response.batch,
            PrefillTimings::new(response.forward_ns, response.decode_ns, response.total_ns),
        ))
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
200
201
    }

202
    /// Generate one token for each request in the given cached batches
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
203
    ///
204
    /// Returns Generation for each request in batches
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
205
    /// and the next cached batch
206
    #[instrument(skip_all, fields(size = batches.iter().map(|batch|{batch.size}).sum::<u32>()))]
207
    pub async fn decode(
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
208
        &mut self,
209
        batches: Vec<CachedBatch>,
210
    ) -> Result<(Vec<Generation>, Option<CachedBatch>, DecodeTimings)> {
211
212
        let request = tonic::Request::new(DecodeRequest { batches }).inject_context();
        let response = self.stub.decode(request).await?.into_inner();
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
        Ok((
            response.generations,
            response.batch,
            DecodeTimings::new(
                response.concat_ns,
                response.forward_ns,
                response.decode_ns,
                response.total_ns,
            ),
        ))
    }
}

pub struct PrefillTimings {
    pub forward: Duration,
    pub decode: Duration,
    pub total: Duration,
}

impl PrefillTimings {
    fn new(forward_ns: u64, decode_ns: u64, total_ns: u64) -> Self {
        Self {
            forward: Duration::from_nanos(forward_ns),
            decode: Duration::from_nanos(decode_ns),
            total: Duration::from_nanos(total_ns),
        }
    }
}

pub struct DecodeTimings {
    pub concat: Option<Duration>,
    pub forward: Duration,
    pub decode: Duration,
    pub total: Duration,
}

impl DecodeTimings {
    fn new(concat_ns: Option<u64>, forward_ns: u64, decode_ns: u64, total_ns: u64) -> Self {
        Self {
252
            concat: concat_ns.map(Duration::from_nanos),
253
254
255
256
            forward: Duration::from_nanos(forward_ns),
            decode: Duration::from_nanos(decode_ns),
            total: Duration::from_nanos(total_ns),
        }
Olivier Dehaene's avatar
Olivier Dehaene committed
257
    }
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
258
}