/build/source/nativelink-scheduler/src/grpc_scheduler.rs
Line | Count | Source |
1 | | // Copyright 2024 The NativeLink Authors. All rights reserved. |
2 | | // |
3 | | // Licensed under the Apache License, Version 2.0 (the "License"); |
4 | | // you may not use this file except in compliance with the License. |
5 | | // You may obtain a copy of the License at |
6 | | // |
7 | | // http://www.apache.org/licenses/LICENSE-2.0 |
8 | | // |
9 | | // Unless required by applicable law or agreed to in writing, software |
10 | | // distributed under the License is distributed on an "AS IS" BASIS, |
11 | | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | | // See the License for the specific language governing permissions and |
13 | | // limitations under the License. |
14 | | |
15 | | use std::collections::HashMap; |
16 | | use std::future::Future; |
17 | | use std::sync::Arc; |
18 | | use std::time::Duration; |
19 | | |
20 | | use async_trait::async_trait; |
21 | | use futures::stream::unfold; |
22 | | use futures::{StreamExt, TryFutureExt}; |
23 | | use nativelink_config::schedulers::GrpcSpec; |
24 | | use nativelink_error::{error_if, make_err, Code, Error, ResultExt}; |
25 | | use nativelink_metric::{MetricsComponent, RootMetricsComponent}; |
26 | | use nativelink_proto::build::bazel::remote::execution::v2::capabilities_client::CapabilitiesClient; |
27 | | use nativelink_proto::build::bazel::remote::execution::v2::execution_client::ExecutionClient; |
28 | | use nativelink_proto::build::bazel::remote::execution::v2::{ |
29 | | ExecuteRequest, ExecutionPolicy, GetCapabilitiesRequest, WaitExecutionRequest, |
30 | | }; |
31 | | use nativelink_proto::google::longrunning::Operation; |
32 | | use nativelink_util::action_messages::{ |
33 | | ActionInfo, ActionState, ActionUniqueQualifier, OperationId, DEFAULT_EXECUTION_PRIORITY, |
34 | | }; |
35 | | use nativelink_util::connection_manager::ConnectionManager; |
36 | | use nativelink_util::known_platform_property_provider::KnownPlatformPropertyProvider; |
37 | | use nativelink_util::operation_state_manager::{ |
38 | | ActionStateResult, ActionStateResultStream, ClientStateManager, OperationFilter, |
39 | | }; |
40 | | use nativelink_util::retry::{Retrier, RetryResult}; |
41 | | use nativelink_util::{background_spawn, tls_utils}; |
42 | | use parking_lot::Mutex; |
43 | | use rand::rngs::OsRng; |
44 | | use rand::Rng; |
45 | | use tokio::select; |
46 | | use tokio::sync::watch; |
47 | | use tokio::time::sleep; |
48 | | use tonic::{Request, Streaming}; |
49 | | use tracing::{event, Level}; |
50 | | |
51 | | struct GrpcActionStateResult { |
52 | | client_operation_id: OperationId, |
53 | | rx: watch::Receiver<Arc<ActionState>>, |
54 | | } |
55 | | |
56 | | #[async_trait] |
57 | | impl ActionStateResult for GrpcActionStateResult { |
58 | 0 | async fn as_state(&self) -> Result<Arc<ActionState>, Error> { |
59 | 0 | let mut action_state = self.rx.borrow().clone(); |
60 | 0 | Arc::make_mut(&mut action_state).client_operation_id = self.client_operation_id.clone(); |
61 | 0 | Ok(action_state) |
62 | 0 | } |
63 | | |
64 | 0 | async fn changed(&mut self) -> Result<Arc<ActionState>, Error> { |
65 | 0 | self.rx.changed().await.map_err(|_| { |
66 | 0 | make_err!( |
67 | 0 | Code::Internal, |
68 | 0 | "Channel closed in GrpcActionStateResult::changed" |
69 | 0 | ) |
70 | 0 | })?; |
71 | 0 | let mut action_state = self.rx.borrow().clone(); |
72 | 0 | Arc::make_mut(&mut action_state).client_operation_id = self.client_operation_id.clone(); |
73 | 0 | Ok(action_state) |
74 | 0 | } |
75 | | |
76 | 0 | async fn as_action_info(&self) -> Result<Arc<ActionInfo>, Error> { |
77 | | // TODO(allada) We should probably remove as_action_info() |
78 | | // or implement it properly. |
79 | 0 | return Err(make_err!( |
80 | 0 | Code::Unimplemented, |
81 | 0 | "as_action_info not implemented for GrpcActionStateResult::as_action_info" |
82 | 0 | )); |
83 | 0 | } |
84 | | } |
85 | | |
86 | | #[derive(MetricsComponent)] |
87 | | pub struct GrpcScheduler { |
88 | | #[metric(group = "property_managers")] |
89 | | supported_props: Mutex<HashMap<String, Vec<String>>>, |
90 | | retrier: Retrier, |
91 | | connection_manager: ConnectionManager, |
92 | | } |
93 | | |
94 | | impl GrpcScheduler { |
95 | 0 | pub fn new(spec: &GrpcSpec) -> Result<Self, Error> { |
96 | 0 | let jitter_amt = spec.retry.jitter; |
97 | 0 | Self::new_with_jitter( |
98 | 0 | spec, |
99 | 0 | Box::new(move |delay: Duration| { |
100 | 0 | if jitter_amt == 0. { Branch (100:20): [True: 0, False: 0]
Branch (100:20): [Folded - Ignored]
|
101 | 0 | return delay; |
102 | 0 | } |
103 | 0 | let min = 1. - (jitter_amt / 2.); |
104 | 0 | let max = 1. + (jitter_amt / 2.); |
105 | 0 | delay.mul_f32(OsRng.gen_range(min..max)) |
106 | 0 | }), |
107 | 0 | ) |
108 | 0 | } |
109 | | |
110 | 0 | pub fn new_with_jitter( |
111 | 0 | spec: &GrpcSpec, |
112 | 0 | jitter_fn: Box<dyn Fn(Duration) -> Duration + Send + Sync>, |
113 | 0 | ) -> Result<Self, Error> { |
114 | 0 | let endpoint = tls_utils::endpoint(&spec.endpoint)?; |
115 | 0 | let jitter_fn = Arc::new(jitter_fn); |
116 | 0 | Ok(Self { |
117 | 0 | supported_props: Mutex::new(HashMap::new()), |
118 | 0 | retrier: Retrier::new( |
119 | 0 | Arc::new(|duration| Box::pin(sleep(duration))), |
120 | 0 | jitter_fn.clone(), |
121 | 0 | spec.retry.clone(), |
122 | 0 | ), |
123 | 0 | connection_manager: ConnectionManager::new( |
124 | 0 | std::iter::once(endpoint), |
125 | 0 | spec.connections_per_endpoint, |
126 | 0 | spec.max_concurrent_requests, |
127 | 0 | spec.retry.clone(), |
128 | 0 | jitter_fn, |
129 | 0 | ), |
130 | 0 | }) |
131 | 0 | } |
132 | | |
133 | 0 | async fn perform_request<F, Fut, R, I>(&self, input: I, mut request: F) -> Result<R, Error> |
134 | 0 | where |
135 | 0 | F: FnMut(I) -> Fut + Send + Copy, |
136 | 0 | Fut: Future<Output = Result<R, Error>> + Send, |
137 | 0 | R: Send, |
138 | 0 | I: Send + Clone, |
139 | 0 | { |
140 | 0 | self.retrier |
141 | 0 | .retry(unfold(input, move |input| async move { |
142 | 0 | let input_clone = input.clone(); |
143 | 0 | Some(( |
144 | 0 | request(input_clone) |
145 | 0 | .await |
146 | 0 | .map_or_else(RetryResult::Retry, RetryResult::Ok), |
147 | 0 | input, |
148 | 0 | )) |
149 | 0 | })) |
150 | 0 | .await |
151 | 0 | } |
152 | | |
153 | 0 | async fn stream_state( |
154 | 0 | mut result_stream: Streaming<Operation>, |
155 | 0 | ) -> Result<Box<dyn ActionStateResult>, Error> { |
156 | 0 | if let Some(initial_response) = result_stream Branch (156:16): [True: 0, False: 0]
Branch (156:16): [Folded - Ignored]
|
157 | 0 | .message() |
158 | 0 | .await |
159 | 0 | .err_tip(|| "Recieving response from upstream scheduler")? |
160 | | { |
161 | 0 | let client_operation_id = OperationId::from(initial_response.name.as_str()); |
162 | 0 | // Our operation_id is not needed here is just a place holder to recycle existing object. |
163 | 0 | // The only thing that actually matters is the operation_id. |
164 | 0 | let operation_id = OperationId::default(); |
165 | 0 | let action_state = |
166 | 0 | ActionState::try_from_operation(initial_response, operation_id.clone()) |
167 | 0 | .err_tip(|| "In GrpcScheduler::stream_state")?; |
168 | 0 | let (tx, mut rx) = watch::channel(Arc::new(action_state)); |
169 | 0 | rx.mark_changed(); |
170 | 0 | background_spawn!("grpc_scheduler_stream_state", async move { |
171 | | loop { |
172 | 0 | select!( |
173 | 0 | () = tx.closed() => { |
174 | 0 | event!( |
175 | 0 | Level::INFO, |
176 | 0 | "Client disconnected in GrpcScheduler" |
177 | | ); |
178 | 0 | return; |
179 | | } |
180 | 0 | response = result_stream.message() => { |
181 | | // When the upstream closes the channel, close the |
182 | | // downstream too. |
183 | 0 | let Ok(Some(response)) = response else { Branch (183:33): [True: 0, False: 0]
Branch (183:33): [Folded - Ignored]
|
184 | 0 | return; |
185 | | }; |
186 | 0 | let maybe_action_state = ActionState::try_from_operation(response, operation_id.clone()); |
187 | 0 | match maybe_action_state { |
188 | 0 | Ok(response) => { |
189 | 0 | if let Err(err) = tx.send(Arc::new(response)) { Branch (189:44): [True: 0, False: 0]
Branch (189:44): [Folded - Ignored]
|
190 | 0 | event!( |
191 | 0 | Level::INFO, |
192 | | ?err, |
193 | 0 | "Client error in GrpcScheduler" |
194 | | ); |
195 | 0 | return; |
196 | 0 | } |
197 | | } |
198 | 0 | Err(err) => { |
199 | 0 | event!( |
200 | 0 | Level::ERROR, |
201 | | ?err, |
202 | 0 | "Error converting response to ActionState in GrpcScheduler" |
203 | | ); |
204 | | }, |
205 | | } |
206 | | } |
207 | | ); |
208 | | } |
209 | 0 | }); |
210 | | |
211 | 0 | return Ok(Box::new(GrpcActionStateResult { |
212 | 0 | client_operation_id, |
213 | 0 | rx, |
214 | 0 | })); |
215 | 0 | } |
216 | 0 | Err(make_err!( |
217 | 0 | Code::Internal, |
218 | 0 | "Upstream scheduler didn't accept action." |
219 | 0 | )) |
220 | 0 | } |
221 | | |
222 | 0 | async fn inner_get_known_properties(&self, instance_name: &str) -> Result<Vec<String>, Error> { |
223 | 0 | if let Some(supported_props) = self.supported_props.lock().get(instance_name) { Branch (223:16): [True: 0, False: 0]
Branch (223:16): [Folded - Ignored]
|
224 | 0 | return Ok(supported_props.clone()); |
225 | 0 | } |
226 | 0 |
|
227 | 0 | self.perform_request(instance_name, |instance_name| async move { |
228 | | // Not in the cache, lookup the capabilities with the upstream. |
229 | 0 | let channel = self |
230 | 0 | .connection_manager |
231 | 0 | .connection() |
232 | 0 | .await |
233 | 0 | .err_tip(|| "in get_platform_property_manager()")?; |
234 | 0 | let capabilities_result = CapabilitiesClient::new(channel) |
235 | 0 | .get_capabilities(GetCapabilitiesRequest { |
236 | 0 | instance_name: instance_name.to_string(), |
237 | 0 | }) |
238 | 0 | .await |
239 | 0 | .err_tip(|| "Retrieving upstream GrpcScheduler capabilities"); |
240 | 0 | let capabilities = capabilities_result?.into_inner(); |
241 | 0 | let supported_props = capabilities |
242 | 0 | .execution_capabilities |
243 | 0 | .err_tip(|| "Unable to get execution properties in GrpcScheduler")? |
244 | | .supported_node_properties |
245 | 0 | .into_iter() |
246 | 0 | .collect::<Vec<String>>(); |
247 | 0 |
|
248 | 0 | self.supported_props |
249 | 0 | .lock() |
250 | 0 | .insert(instance_name.to_string(), supported_props.clone()); |
251 | 0 | Ok(supported_props) |
252 | 0 | }) |
253 | 0 | .await |
254 | 0 | } |
255 | | |
256 | 0 | async fn inner_add_action( |
257 | 0 | &self, |
258 | 0 | _client_operation_id: OperationId, |
259 | 0 | action_info: Arc<ActionInfo>, |
260 | 0 | ) -> Result<Box<dyn ActionStateResult>, Error> { |
261 | 0 | let execution_policy = if action_info.priority == DEFAULT_EXECUTION_PRIORITY { Branch (261:35): [True: 0, False: 0]
Branch (261:35): [Folded - Ignored]
|
262 | 0 | None |
263 | | } else { |
264 | 0 | Some(ExecutionPolicy { |
265 | 0 | priority: action_info.priority, |
266 | 0 | }) |
267 | | }; |
268 | 0 | let skip_cache_lookup = match action_info.unique_qualifier { |
269 | 0 | ActionUniqueQualifier::Cachable(_) => false, |
270 | 0 | ActionUniqueQualifier::Uncachable(_) => true, |
271 | | }; |
272 | 0 | let request = ExecuteRequest { |
273 | 0 | instance_name: action_info.instance_name().clone(), |
274 | 0 | skip_cache_lookup, |
275 | 0 | action_digest: Some(action_info.digest().into()), |
276 | 0 | execution_policy, |
277 | 0 | // TODO: Get me from the original request, not very important as we ignore it. |
278 | 0 | results_cache_policy: None, |
279 | 0 | digest_function: action_info |
280 | 0 | .unique_qualifier |
281 | 0 | .digest_function() |
282 | 0 | .proto_digest_func() |
283 | 0 | .into(), |
284 | 0 | }; |
285 | 0 | let result_stream = self |
286 | 0 | .perform_request(request, |request| async move { |
287 | 0 | let channel = self |
288 | 0 | .connection_manager |
289 | 0 | .connection() |
290 | 0 | .await |
291 | 0 | .err_tip(|| "in add_action()")?; |
292 | 0 | ExecutionClient::new(channel) |
293 | 0 | .execute(Request::new(request)) |
294 | 0 | .await |
295 | 0 | .err_tip(|| "Sending action to upstream scheduler") |
296 | 0 | }) |
297 | 0 | .await? |
298 | 0 | .into_inner(); |
299 | 0 | Self::stream_state(result_stream).await |
300 | 0 | } |
301 | | |
302 | 0 | async fn inner_filter_operations( |
303 | 0 | &self, |
304 | 0 | filter: OperationFilter, |
305 | 0 | ) -> Result<ActionStateResultStream, Error> { |
306 | 0 | error_if!(filter != OperationFilter { Branch (306:19): [True: 0, False: 0]
Branch (306:19): [Folded - Ignored]
|
307 | 0 | client_operation_id: filter.client_operation_id.clone(), |
308 | 0 | ..Default::default() |
309 | 0 | }, "Unsupported filter in GrpcScheduler::filter_operations. Only client_operation_id is supported - {filter:?}"); |
310 | 0 | let client_operation_id = filter.client_operation_id.ok_or_else(|| { |
311 | 0 | make_err!(Code::InvalidArgument, "`client_operation_id` is the only supported filter in GrpcScheduler::filter_operations") |
312 | 0 | })?; |
313 | 0 | let request = WaitExecutionRequest { |
314 | 0 | name: client_operation_id.to_string(), |
315 | 0 | }; |
316 | 0 | let result_stream = self |
317 | 0 | .perform_request(request, |request| async move { |
318 | 0 | let channel = self |
319 | 0 | .connection_manager |
320 | 0 | .connection() |
321 | 0 | .await |
322 | 0 | .err_tip(|| "in find_by_client_operation_id()")?; |
323 | 0 | ExecutionClient::new(channel) |
324 | 0 | .wait_execution(Request::new(request)) |
325 | 0 | .await |
326 | 0 | .err_tip(|| "While getting wait_execution stream") |
327 | 0 | }) |
328 | 0 | .and_then(|result_stream| Self::stream_state(result_stream.into_inner())) |
329 | 0 | .await; |
330 | 0 | match result_stream { |
331 | 0 | Ok(result_stream) => Ok(unfold( |
332 | 0 | Some(result_stream), |
333 | 0 | |maybe_result_stream| async move { maybe_result_stream.map(|v| (v, None)) }, |
334 | 0 | ) |
335 | 0 | .boxed()), |
336 | 0 | Err(err) => { |
337 | 0 | event!( |
338 | 0 | Level::WARN, |
339 | | ?err, |
340 | 0 | "Error looking up action with upstream scheduler" |
341 | | ); |
342 | 0 | Ok(futures::stream::empty().boxed()) |
343 | | } |
344 | | } |
345 | 0 | } |
346 | | } |
347 | | |
348 | | #[async_trait] |
349 | | impl ClientStateManager for GrpcScheduler { |
350 | | async fn add_action( |
351 | | &self, |
352 | | client_operation_id: OperationId, |
353 | | action_info: Arc<ActionInfo>, |
354 | 0 | ) -> Result<Box<dyn ActionStateResult>, Error> { |
355 | 0 | self.inner_add_action(client_operation_id, action_info) |
356 | 0 | .await |
357 | 0 | } |
358 | | |
359 | | async fn filter_operations<'a>( |
360 | | &'a self, |
361 | | filter: OperationFilter, |
362 | 0 | ) -> Result<ActionStateResultStream<'a>, Error> { |
363 | 0 | self.inner_filter_operations(filter).await |
364 | 0 | } |
365 | | |
366 | 0 | fn as_known_platform_property_provider(&self) -> Option<&dyn KnownPlatformPropertyProvider> { |
367 | 0 | Some(self) |
368 | 0 | } |
369 | | } |
370 | | |
371 | | #[async_trait] |
372 | | impl KnownPlatformPropertyProvider for GrpcScheduler { |
373 | 0 | async fn get_known_properties(&self, instance_name: &str) -> Result<Vec<String>, Error> { |
374 | 0 | self.inner_get_known_properties(instance_name).await |
375 | 0 | } |
376 | | } |
377 | | |
378 | | impl RootMetricsComponent for GrpcScheduler {} |