/build/source/nativelink-scheduler/src/worker.rs

Source
// Copyright 2024 The NativeLink Authors. All rights reserved.
//
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    See LICENSE file for details
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use core::hash::{Hash, Hasher};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};

use nativelink_error::{Code, Error, ResultExt};
use nativelink_metric::MetricsComponent;
use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{
    ConnectionResult, StartExecute, UpdateForWorker, update_for_worker,
};
use nativelink_util::action_messages::{ActionInfo, OperationId, WorkerId};
use nativelink_util::metrics_utils::{AsyncCounterWrapper, CounterWithTime, FuncCounterWrapper};
use nativelink_util::origin_event::OriginMetadata;
use nativelink_util::platform_properties::{PlatformProperties, PlatformPropertyValue};
use tokio::sync::mpsc::UnboundedSender;

pub type WorkerTimestamp = u64;

/// Represents the action info and the platform properties of the action.
/// These platform properties have the type of the properties as well as
/// the value of the properties, unlike `ActionInfo`, which only has the
/// string value of the properties.
#[derive(Clone, Debug, MetricsComponent)]
pub struct ActionInfoWithProps {
    /// The action info of the action.
    #[metric(group = "action_info")]
    pub inner: Arc<ActionInfo>,
    /// The platform properties of the action.
    #[metric(group = "platform_properties")]
    pub platform_properties: PlatformProperties,
    /// Origin metadata used when publishing scheduler-side telemetry for this action.
    pub origin_metadata: OriginMetadata,
    /// `OriginEvent` id for the `scheduler_start_execute` request.
    pub scheduler_start_execute_event_id: Option<String>,
}

/// Notifications to send worker about a requested state change.
#[derive(Debug)]
pub enum WorkerUpdate {
    /// Requests that the worker begin executing this action.
    RunAction(Box<(OperationId, ActionInfoWithProps)>),

    /// Request that the worker is no longer in the pool and may discard any jobs.
    Disconnect,
}

#[derive(Debug, MetricsComponent)]
pub struct PendingActionInfoData {
    #[metric]
    pub action_info: ActionInfoWithProps,
}

/// Represents a connection to a worker and used as the medium to
/// interact with the worker from the client/scheduler.
#[derive(Debug, MetricsComponent)]
pub struct Worker {
    /// Unique identifier of the worker.
    #[metric(help = "The unique identifier of the worker.")]
    pub id: WorkerId,

    /// Properties that describe the capabilities of this worker.
    #[metric(group = "platform_properties")]
    pub platform_properties: PlatformProperties,

    /// Channel to send commands from scheduler to worker.
    pub tx: UnboundedSender<UpdateForWorker>,

    /// The action info of the running actions on the worker.
    #[metric(group = "running_action_infos")]
    pub running_action_infos: HashMap<OperationId, PendingActionInfoData>,

    /// If the properties were restored already then it's added to this set.
    pub restored_platform_properties: HashSet<OperationId>,

    /// Timestamp of last time this worker had been communicated with.
    // Warning: Do not update this timestamp without updating the placement of the worker in
    // the LRUCache in the Workers struct.
    #[metric(help = "Last time this worker was communicated with.")]
    pub last_update_timestamp: WorkerTimestamp,

    /// Whether the worker rejected the last action due to back pressure.
    #[metric(help = "If the worker is paused.")]
    pub is_paused: bool,

    /// Whether the worker is draining.
    #[metric(help = "If the worker is draining.")]
    pub is_draining: bool,

    /// Maximum inflight tasks for this worker (or 0 for unlimited)
    #[metric(help = "Maximum inflight tasks for this worker (or 0 for unlimited)")]
    pub max_inflight_tasks: u64,

    /// Stats about the worker.
    #[metric]
    metrics: Arc<Metrics>,
}

fn send_msg_to_worker(
    tx: &UnboundedSender<UpdateForWorker>,
    msg: update_for_worker::Update,
) -> Result<(), Error> {
    tx.send(UpdateForWorker { update: Some(msg) })
        .map_err(|err| Error::from_std_err7(Code::Internal7, &err7).append7("Worker disconnected"))
}

/// Reduces the platform properties available on the worker based on the platform properties provided.
/// This is used because we allow more than 1 job to run on a worker at a time, and this is how the
/// scheduler knows if more jobs can run on a given worker.
fn reduce_platform_properties(
    parent_props: &mut PlatformProperties,
    reduction_props: &PlatformProperties,
) {
    debug_assert!(reduction_props0.is_satisfied_by0(parent_props0, false));
    for (property8, prop_value8) in &reduction_props.properties {
        if let PlatformPropertyValue::Minimum(value7) = prop_value {
            let worker_props = &mut parent_props.properties;
            if let &mut PlatformPropertyValue::Minimum(worker_value) =
                &mut worker_props.get_mut(property).unwrap()
            {
                *worker_value -= value;
            }0
        }
    }
}

impl Worker {
    pub fn new(
        id: WorkerId,
        platform_properties: PlatformProperties,
        tx: UnboundedSender<UpdateForWorker>,
        timestamp: WorkerTimestamp,
        max_inflight_tasks: u64,
    ) -> Self {
        Self {
            id,
            platform_properties,
            tx,
            running_action_infos: HashMap::new(),
            restored_platform_properties: HashSet::new(),
            last_update_timestamp: timestamp,
            is_paused: false,
            is_draining: false,
            max_inflight_tasks,
            metrics: Arc::new(Metrics {
                connected_timestamp: SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .unwrap()
                    .as_secs(),
                actions_completed: CounterWithTime::default(),
                run_action: AsyncCounterWrapper::default(),
                keep_alive: FuncCounterWrapper::default(),
                notify_disconnect: CounterWithTime::default(),
            }),
        }
    }

    /// Sends the initial connection information to the worker. This generally is just meta info.
    /// This should only be sent once and should always be the first item in the stream.
    pub fn send_initial_connection_result(&mut self) -> Result<(), Error> {
        send_msg_to_worker(
            &self.tx,
            update_for_worker::Update::ConnectionResult(ConnectionResult {
                worker_id: self.id.clone().into(),
            }),
        )
        .err_tip(|| format!0("Failed to send ConnectionResult to worker : {}", self.id))
    }

    /// Notifies the worker of a requested state change.
    pub async fn notify_update(&mut self, worker_update: WorkerUpdate) -> Result<(), Error> {
        match worker_update {
            WorkerUpdate::RunAction(action) => {
                let (operation_id, action_info) = *action;
                self.run_action(operation_id, action_info).await
            }
            WorkerUpdate::Disconnect => {
                self.metrics.notify_disconnect.inc();
                send_msg_to_worker(&self.tx, update_for_worker::Update::Disconnect(()))
            }
        }
    }

    pub fn keep_alive(&mut self) -> Result<(), Error> {
        let tx = &mut self.tx;
        let id = &self.id;
        self.metrics.keep_alive.wrap(move || {
            send_msg_to_worker(tx, update_for_worker::Update::KeepAlive(()))
                .err_tip(|| format!0("Failed to send KeepAlive to worker : {id}"))
        })
    }

    async fn run_action(
        &mut self,
        operation_id: OperationId,
        action_info: ActionInfoWithProps,
    ) -> Result<(), Error> {
        let tx = &mut self.tx;
        let worker_platform_properties = &mut self.platform_properties;
        let running_action_infos = &mut self.running_action_infos;
        let worker_id = self.id.clone().into();
        self.metrics
            .run_action
            .wrap(async move {
                let action_info_clone = action_info.clone();
                let operation_id_string = operation_id.to_string();
                let start_execute = StartExecute {
                    execute_request: Some(action_info_clone.inner.as_ref().into()),
                    operation_id: operation_id_string,
                    queued_timestamp: Some(action_info.inner.insert_timestamp.into()),
                    platform: Some((&action_info.platform_properties).into()),
                    worker_id,
                };
                reduce_platform_properties(
                    worker_platform_properties,
                    &action_info.platform_properties,
                );
                running_action_infos.insert(operation_id, PendingActionInfoData { action_info });

                send_msg_to_worker(tx, update_for_worker::Update::StartAction(start_execute))
            })
            .await
    }

    pub(crate) fn execution_complete(&mut self, operation_id: &OperationId) {
        if let Some((operation_id, pending_action_info)) =
            self.running_action_infos.remove_entry(operation_id)
        {
            self.restored_platform_properties
                .insert(operation_id.clone());
            self.restore_platform_properties(&pending_action_info.action_info.platform_properties);
            self.running_action_infos
                .insert(operation_id, pending_action_info);
        }
    }

    pub(crate) async fn complete_action(
        &mut self,
        operation_id: &OperationId,
    ) -> Result<(), Error> {
        let pending_action_info = self.running_action_infos.remove(operation_id).err_tip(|| {0
            format!(
                "Worker {} tried to complete operation {} that was not running",
                self.id, operation_id
            )
        })?;
        if !self.restored_platform_properties.remove(operation_id) {
            self.restore_platform_properties(&pending_action_info.action_info.platform_properties);
        }0
        self.is_paused = false;
        self.metrics.actions_completed.inc();
        Ok(())
    }

    pub fn has_actions(&self) -> bool {
        !self.running_action_infos.is_empty()
    }

    fn restore_platform_properties(&mut self, props: &PlatformProperties) {
        for (property3, prop_value3) in &props.properties {
            if let PlatformPropertyValue::Minimum(value) = prop_value {
                let worker_props = &mut self.platform_properties.properties;
                if let PlatformPropertyValue::Minimum(worker_value) =
                    worker_props.get_mut(property).unwrap()
                {
                    *worker_value += value;
                }0
            }
        }
    }

    pub fn can_accept_work(&self) -> bool {
        !self.is_paused
            && !self.is_draining
            && (self.max_inflight_tasks == 0
                || u64::try_from(self.running_action_infos.len()).unwrap_or(u64::MAX)
                    < self.max_inflight_tasks)
    }
}

impl PartialEq for Worker {
    fn eq(&self, other: &Self) -> bool {
        self.id == other.id
    }
}

impl Eq for Worker {}

impl Hash for Worker {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.id.hash(state);
    }
}

#[derive(Debug, Default, MetricsComponent)]
struct Metrics {
    #[metric(help = "The timestamp of when this worker connected.")]
    connected_timestamp: u64,
    #[metric(help = "The number of actions completed for this worker.")]
    actions_completed: CounterWithTime,
    #[metric(help = "The number of actions started for this worker.")]
    run_action: AsyncCounterWrapper,
    #[metric(help = "The number of keep_alive sent to this worker.")]
    keep_alive: FuncCounterWrapper,
    #[metric(help = "The number of notify_disconnect sent to this worker.")]
    notify_disconnect: CounterWithTime,
}

Coverage Report

Created: 2026-08-12 00:28

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::{Hash, Hasher};
16		use std::collections::{HashMap, HashSet};
17		use std::sync::Arc;
18		use std::time::{SystemTime, UNIX_EPOCH};
19
20		use nativelink_error::{Code, Error, ResultExt};
21		use nativelink_metric::MetricsComponent;
22		use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{
23		ConnectionResult, StartExecute, UpdateForWorker, update_for_worker,
24		};
25		use nativelink_util::action_messages::{ActionInfo, OperationId, WorkerId};
26		use nativelink_util::metrics_utils::{AsyncCounterWrapper, CounterWithTime, FuncCounterWrapper};
27		use nativelink_util::origin_event::OriginMetadata;
28		use nativelink_util::platform_properties::{PlatformProperties, PlatformPropertyValue};
29		use tokio::sync::mpsc::UnboundedSender;
30
31		pub type WorkerTimestamp = u64;
32
33		/// Represents the action info and the platform properties of the action.
34		/// These platform properties have the type of the properties as well as
35		/// the value of the properties, unlike `ActionInfo`, which only has the
36		/// string value of the properties.
37		#[derive(Clone, Debug, MetricsComponent)]
38		pub struct ActionInfoWithProps {
39		/// The action info of the action.
40		#[metric(group = "action_info")]
41		pub inner: Arc<ActionInfo>,
42		/// The platform properties of the action.
43		#[metric(group = "platform_properties")]
44		pub platform_properties: PlatformProperties,
45		/// Origin metadata used when publishing scheduler-side telemetry for this action.
46		pub origin_metadata: OriginMetadata,
47		/// `OriginEvent` id for the `scheduler_start_execute` request.
48		pub scheduler_start_execute_event_id: Option<String>,
49		}
50
51		/// Notifications to send worker about a requested state change.
52		#[derive(Debug)]
53		pub enum WorkerUpdate {
54		/// Requests that the worker begin executing this action.
55		RunAction(Box<(OperationId, ActionInfoWithProps)>),
56
57		/// Request that the worker is no longer in the pool and may discard any jobs.
58		Disconnect,
59		}
60
61		#[derive(Debug, MetricsComponent)]
62		pub struct PendingActionInfoData {
63		#[metric]
64		pub action_info: ActionInfoWithProps,
65		}
66
67		/// Represents a connection to a worker and used as the medium to
68		/// interact with the worker from the client/scheduler.
69		#[derive(Debug, MetricsComponent)]
70		pub struct Worker {
71		/// Unique identifier of the worker.
72		#[metric(help = "The unique identifier of the worker.")]
73		pub id: WorkerId,
74
75		/// Properties that describe the capabilities of this worker.
76		#[metric(group = "platform_properties")]
77		pub platform_properties: PlatformProperties,
78
79		/// Channel to send commands from scheduler to worker.
80		pub tx: UnboundedSender<UpdateForWorker>,
81
82		/// The action info of the running actions on the worker.
83		#[metric(group = "running_action_infos")]
84		pub running_action_infos: HashMap<OperationId, PendingActionInfoData>,
85
86		/// If the properties were restored already then it's added to this set.
87		pub restored_platform_properties: HashSet<OperationId>,
88
89		/// Timestamp of last time this worker had been communicated with.
90		// Warning: Do not update this timestamp without updating the placement of the worker in
91		// the LRUCache in the Workers struct.
92		#[metric(help = "Last time this worker was communicated with.")]
93		pub last_update_timestamp: WorkerTimestamp,
94
95		/// Whether the worker rejected the last action due to back pressure.
96		#[metric(help = "If the worker is paused.")]
97		pub is_paused: bool,
98
99		/// Whether the worker is draining.
100		#[metric(help = "If the worker is draining.")]
101		pub is_draining: bool,
102
103		/// Maximum inflight tasks for this worker (or 0 for unlimited)
104		#[metric(help = "Maximum inflight tasks for this worker (or 0 for unlimited)")]
105		pub max_inflight_tasks: u64,
106
107		/// Stats about the worker.
108		#[metric]
109		metrics: Arc<Metrics>,
110		}
111
112	94	fn send_msg_to_worker(
113	94	tx: &UnboundedSender<UpdateForWorker>,
114	94	msg: update_for_worker::Update,
115	94	) -> Result<(), Error> {
116	94	tx.send(UpdateForWorker { update: Some(msg) })
117	94	.map_err(\|err\| Error::from_std_err7 ( Code::Internal7 , &err7 ). append7 ("Worker disconnected"))
118	94	}
119
120		/// Reduces the platform properties available on the worker based on the platform properties provided.
121		/// This is used because we allow more than 1 job to run on a worker at a time, and this is how the
122		/// scheduler knows if more jobs can run on a given worker.
123	40	fn reduce_platform_properties(
124	40	parent_props: &mut PlatformProperties,
125	40	reduction_props: &PlatformProperties,
126	40	) {
127	40	debug_assert!( reduction_props0 . is_satisfied_by0 ( parent_props0 , false));
128	40	for ( property8 , prop_value8 ) in &reduction_props.properties {
129	8	if let PlatformPropertyValue::Minimum( value7 ) = prop_value {
130	7	let worker_props = &mut parent_props.properties;
131	7	if let &mut PlatformPropertyValue::Minimum(worker_value) =
132	7	&mut worker_props.get_mut(property).unwrap()
133	7	{
134	7	*worker_value -= value;
135	7	}0
136	1	}
137		}
138	40	}
139
140		impl Worker {
141	43	pub fn new(
142	43	id: WorkerId,
143	43	platform_properties: PlatformProperties,
144	43	tx: UnboundedSender<UpdateForWorker>,
145	43	timestamp: WorkerTimestamp,
146	43	max_inflight_tasks: u64,
147	43	) -> Self {
148	43	Self {
149	43	id,
150	43	platform_properties,
151	43	tx,
152	43	running_action_infos: HashMap::new(),
153	43	restored_platform_properties: HashSet::new(),
154	43	last_update_timestamp: timestamp,
155	43	is_paused: false,
156	43	is_draining: false,
157	43	max_inflight_tasks,
158	43	metrics: Arc::new(Metrics {
159	43	connected_timestamp: SystemTime::now()
160	43	.duration_since(UNIX_EPOCH)
161	43	.unwrap()
162	43	.as_secs(),
163	43	actions_completed: CounterWithTime::default(),
164	43	run_action: AsyncCounterWrapper::default(),
165	43	keep_alive: FuncCounterWrapper::default(),
166	43	notify_disconnect: CounterWithTime::default(),
167	43	}),
168	43	}
169	43	}
170
171		/// Sends the initial connection information to the worker. This generally is just meta info.
172		/// This should only be sent once and should always be the first item in the stream.
173	43	pub fn send_initial_connection_result(&mut self) -> Result<(), Error> {
174	43	send_msg_to_worker(
175	43	&self.tx,
176	43	update_for_worker::Update::ConnectionResult(ConnectionResult {
177	43	worker_id: self.id.clone().into(),
178	43	}),
179		)
180	43	.err_tip(\|\| format!0 ("Failed to send ConnectionResult to worker : {}", self.id))
181	43	}
182
183		/// Notifies the worker of a requested state change.
184	50	pub async fn notify_update(&mut self, worker_update: WorkerUpdate) -> Result<(), Error> {
185	50	match worker_update {
186	40	WorkerUpdate::RunAction(action) => {
187	40	let (operation_id, action_info) = *action;
188	40	self.run_action(operation_id, action_info).await
189		}
190		WorkerUpdate::Disconnect => {
191	10	self.metrics.notify_disconnect.inc();
192	10	send_msg_to_worker(&self.tx, update_for_worker::Update::Disconnect(()))
193		}
194		}
195	50	}
196
197	1	pub fn keep_alive(&mut self) -> Result<(), Error> {
198	1	let tx = &mut self.tx;
199	1	let id = &self.id;
200	1	self.metrics.keep_alive.wrap(move \|\| {
201	1	send_msg_to_worker(tx, update_for_worker::Update::KeepAlive(()))
202	1	.err_tip(\|\| format!0 ("Failed to send KeepAlive to worker : {id}"))
203	1	})
204	1	}
205
206	40	async fn run_action(
207	40	&mut self,
208	40	operation_id: OperationId,
209	40	action_info: ActionInfoWithProps,
210	40	) -> Result<(), Error> {
211	40	let tx = &mut self.tx;
212	40	let worker_platform_properties = &mut self.platform_properties;
213	40	let running_action_infos = &mut self.running_action_infos;
214	40	let worker_id = self.id.clone().into();
215	40	self.metrics
216	40	.run_action
217	40	.wrap(async move {
218	40	let action_info_clone = action_info.clone();
219	40	let operation_id_string = operation_id.to_string();
220	40	let start_execute = StartExecute {
221	40	execute_request: Some(action_info_clone.inner.as_ref().into()),
222	40	operation_id: operation_id_string,
223	40	queued_timestamp: Some(action_info.inner.insert_timestamp.into()),
224	40	platform: Some((&action_info.platform_properties).into()),
225	40	worker_id,
226	40	};
227	40	reduce_platform_properties(
228	40	worker_platform_properties,
229	40	&action_info.platform_properties,
230		);
231	40	running_action_infos.insert(operation_id, PendingActionInfoData { action_info });
232
233	40	send_msg_to_worker(tx, update_for_worker::Update::StartAction(start_execute))
234	40	})
235	40	.await
236	40	}
237
238	0	pub(crate) fn execution_complete(&mut self, operation_id: &OperationId) {
239	0	if let Some((operation_id, pending_action_info)) =
240	0	self.running_action_infos.remove_entry(operation_id)
241	0	{
242	0	self.restored_platform_properties
243	0	.insert(operation_id.clone());
244	0	self.restore_platform_properties(&pending_action_info.action_info.platform_properties);
245	0	self.running_action_infos
246	0	.insert(operation_id, pending_action_info);
247	0	}
248	0	}
249
250	11	pub(crate) async fn complete_action(
251	11	&mut self,
252	11	operation_id: &OperationId,
253	11	) -> Result<(), Error> {
254	11	let pending_action_info = self.running_action_infos.remove(operation_id).err_tip(\|\| {0
255	0	format!(
256		"Worker {} tried to complete operation {} that was not running",
257		self.id, operation_id
258		)
259	0	})?;
260	11	if !self.restored_platform_properties.remove(operation_id) {
261	11	self.restore_platform_properties(&pending_action_info.action_info.platform_properties);
262	11	}0
263	11	self.is_paused = false;
264	11	self.metrics.actions_completed.inc();
265	11	Ok(())
266	11	}
267
268	0	pub fn has_actions(&self) -> bool {
269	0	!self.running_action_infos.is_empty()
270	0	}
271
272	11	fn restore_platform_properties(&mut self, props: &PlatformProperties) {
273	11	for ( property3 , prop_value3 ) in &props.properties {
274	3	if let PlatformPropertyValue::Minimum(value) = prop_value {
275	3	let worker_props = &mut self.platform_properties.properties;
276	3	if let PlatformPropertyValue::Minimum(worker_value) =
277	3	worker_props.get_mut(property).unwrap()
278	3	{
279	3	*worker_value += value;
280	3	}0
281	0	}
282		}
283	11	}
284
285	116	pub fn can_accept_work(&self) -> bool {
286	116	!self.is_paused
287	116	&& !self.is_draining
288	115	&& (self.max_inflight_tasks == 0
289	3	\|\| u64::try_from(self.running_action_infos.len()).unwrap_or(u64::MAX)
290	3	< self.max_inflight_tasks)
291	116	}
292		}
293
294		impl PartialEq for Worker {
295	0	fn eq(&self, other: &Self) -> bool {
296	0	self.id == other.id
297	0	}
298		}
299
300		impl Eq for Worker {}
301
302		impl Hash for Worker {
303	0	fn hash<H: Hasher>(&self, state: &mut H) {
304	0	self.id.hash(state);
305	0	}
306		}
307
308		#[derive(Debug, Default, MetricsComponent)]
309		struct Metrics {
310		#[metric(help = "The timestamp of when this worker connected.")]
311		connected_timestamp: u64,
312		#[metric(help = "The number of actions completed for this worker.")]
313		actions_completed: CounterWithTime,
314		#[metric(help = "The number of actions started for this worker.")]
315		run_action: AsyncCounterWrapper,
316		#[metric(help = "The number of keep_alive sent to this worker.")]
317		keep_alive: FuncCounterWrapper,
318		#[metric(help = "The number of notify_disconnect sent to this worker.")]
319		notify_disconnect: CounterWithTime,
320		}