/build/source/nativelink-worker/src/worker_utils.rs
Line | Count | Source |
1 | | // Copyright 2024 The NativeLink Authors. All rights reserved. |
2 | | // |
3 | | // Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (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 | | // See LICENSE file for details |
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 core::hash::BuildHasher; |
16 | | use std::collections::HashMap; |
17 | | use std::io::{BufRead, BufReader, Cursor}; |
18 | | use std::process::Stdio; |
19 | | |
20 | | use futures::future::try_join_all; |
21 | | use nativelink_config::cas_server::WorkerProperty; |
22 | | use nativelink_error::{Error, ResultExt, make_err, make_input_err}; |
23 | | use nativelink_proto::build::bazel::remote::execution::v2::platform::Property; |
24 | | use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::ConnectWorkerRequest; |
25 | | use tokio::process; |
26 | | use tracing::{info, warn}; |
27 | | |
28 | | #[expect(clippy::future_not_send)] // TODO(jhpratt) remove this |
29 | 19 | pub async fn make_connect_worker_request<S: BuildHasher>( |
30 | 19 | worker_id_prefix: String, |
31 | 19 | worker_properties: &HashMap<String, WorkerProperty, S>, |
32 | 19 | extra_envs: &HashMap<String, String, S>, |
33 | 19 | max_inflight_tasks: u64, |
34 | 19 | ) -> Result<ConnectWorkerRequest, Error> { |
35 | 19 | let mut futures = vec![]; |
36 | 19 | for (property_name7 , worker_property7 ) in worker_properties { |
37 | 7 | futures.push(async move { |
38 | 7 | match worker_property { |
39 | 1 | WorkerProperty::Values(values) => { |
40 | 1 | let mut props = Vec::with_capacity(values.len()); |
41 | 1 | for value in values { |
42 | 1 | props.push(Property { |
43 | 1 | name: property_name.clone(), |
44 | 1 | value: value.clone(), |
45 | 1 | }); |
46 | 1 | } |
47 | 1 | Ok(props) |
48 | | } |
49 | 6 | WorkerProperty::QueryCmd(cmd) => { |
50 | 6 | let maybe_split_cmd = shlex::split(cmd); |
51 | 6 | let (command5 , args5 ) = match &maybe_split_cmd { |
52 | 5 | Some(split_cmd) => (&split_cmd[0], &split_cmd[1..]), |
53 | | None => { |
54 | 1 | return Err(make_input_err!( |
55 | 1 | "Could not parse the value of worker property: {}: '{}'", |
56 | 1 | property_name, |
57 | 1 | cmd |
58 | 1 | )); |
59 | | } |
60 | | }; |
61 | 5 | let mut process = process::Command::new(command); |
62 | 5 | process.env_clear(); |
63 | 5 | process.envs(extra_envs); |
64 | 5 | process.args(args); |
65 | 5 | process.stdin(Stdio::null()); |
66 | 5 | let err_fn = || {1 |
67 | 1 | format!("Error executing property_name {property_name} command: '{cmd}'") |
68 | 1 | }; |
69 | 5 | info!(cmd, property_name, "Spawning process",); |
70 | 5 | let process_output = process.output().await.err_tip(err_fn)?0 ; |
71 | 5 | if !process_output.stderr.is_empty() { |
72 | 2 | warn!( |
73 | 2 | stderr = ?String::from_utf8_lossy(&process_output.stderr), |
74 | | cmd = cmd, |
75 | | property_name = property_name, |
76 | | "Got stderr when running query cmd" |
77 | | ); |
78 | 3 | } |
79 | 5 | if !process_output.status.success() { |
80 | 1 | return Err(make_err!( |
81 | 1 | process_output.status.code().unwrap().into(), |
82 | 1 | "{}", |
83 | 1 | err_fn() |
84 | 1 | )); |
85 | 4 | } |
86 | 4 | let reader = BufReader::new(Cursor::new(process_output.stdout)); |
87 | | |
88 | 4 | let mut props = vec![]; |
89 | 4 | for value3 in reader.lines() { |
90 | 3 | props.push(Property { |
91 | 3 | name: property_name.clone(), |
92 | 3 | value: value |
93 | 3 | .err_tip(|| "Could split input by lines")?0 |
94 | 3 | .trim() |
95 | 3 | .to_string(), |
96 | | }); |
97 | | } |
98 | 4 | Ok(props) |
99 | | } |
100 | | } |
101 | 7 | }); |
102 | | } |
103 | | |
104 | | Ok(ConnectWorkerRequest { |
105 | 19 | worker_id_prefix, |
106 | 19 | properties: try_join_all(futures).await?2 .into_iter17 ().flatten17 ().collect17 (), |
107 | 17 | max_inflight_tasks, |
108 | | }) |
109 | 19 | } |