Merge branch 'TraceMachina:main' into main

Author: MarcusSorealheis

nativelink-config/examples/worker_with_redis_scheduler.json5

@@ -78,6 +78,7 @@
         worker_api_endpoint: {
           uri: "grpc://127.0.0.1:50061",
         },
+        max_inflight_tasks: 5,
         cas_fast_slow_store: "WORKER_FAST_SLOW_STORE",
         upload_action_result: {
           ac_store: "AC_MAIN_STORE",

nativelink-config/src/cas_server.rs

@@ -727,6 +727,12 @@ pub struct LocalWorkerConfig {
     #[serde(default, deserialize_with = "convert_duration_with_shellexpand")]
     pub max_action_timeout: usize,
 
+    /// Maximum number of inflight tasks this worker can cope with.
+    ///
+    /// Default: 0 (infinite tasks)
+    #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")]
+    pub max_inflight_tasks: u64,
+
     /// If timeout is handled in `entrypoint` or another wrapper script.
     /// If set to true `NativeLink` will not honor the timeout the action requested
     /// and instead will always force kill the action after `max_action_timeout`

nativelink-proto/com/github/trace_machina/nativelink/remote_execution/worker_api.proto

@@ -71,7 +71,11 @@ message ConnectWorkerRequest {
     /// append this prefix to the assigned worker_id followed by a UUIDv6.
     string worker_id_prefix = 2;
 
-    reserved 3; // NextId.
+    /// Maximum number of inflight tasks this worker can cope with at one time
+    /// The default (0) means unlimited.
+    uint64 max_inflight_tasks = 3;
+
+    reserved 4; // NextId.
 }
 
 /// The result of an ExecutionRequest.

nativelink-proto/genproto/com.github.trace_machina.nativelink.remote_execution.pb.rs

@@ -42,6 +42,10 @@ pub struct ConnectWorkerRequest {
     /// / append this prefix to the assigned worker_id followed by a UUIDv6.
     #[prost(string, tag = "2")]
     pub worker_id_prefix: ::prost::alloc::string::String,
+    /// / Maximum number of inflight tasks this worker can cope with at one time
+    /// / The default (0) means unlimited.
+    #[prost(uint64, tag = "3")]
+    pub max_inflight_tasks: u64,
 }
 /// / The result of an ExecutionRequest.
 #[derive(Clone, PartialEq, ::prost::Message)]

nativelink-scheduler/src/api_worker_scheduler.rs

@@ -259,8 +259,11 @@ impl ApiWorkerSchedulerImpl {
             if !w.can_accept_work() {
                 if full_worker_logging {
                     info!(
-                        "Worker {worker_id} cannot accept work: is_paused={}, is_draining={}",
-                        w.is_paused, w.is_draining
+                        "Worker {worker_id} cannot accept work: is_paused={}, is_draining={}, inflight={}/{}",
+                        w.is_paused,
+                        w.is_draining,
+                        w.running_action_infos.len(),
+                        w.max_inflight_tasks
                     );
                 }
                 return false;

nativelink-scheduler/src/worker.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use core::hash::{Hash, Hasher};
+use core::u64;
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};
@@ -95,6 +96,10 @@ pub struct Worker {
     #[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>,
@@ -134,6 +139,7 @@ impl Worker {
         platform_properties: PlatformProperties,
         tx: UnboundedSender<UpdateForWorker>,
         timestamp: WorkerTimestamp,
+        max_inflight_tasks: u64,
     ) -> Self {
         Self {
             id,
@@ -144,6 +150,7 @@ impl Worker {
             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)
@@ -270,8 +277,12 @@ impl Worker {
         }
     }
 
-    pub const fn can_accept_work(&self) -> bool {
-        !self.is_paused && !self.is_draining
+    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)
     }
 }
 

nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs

@@ -369,7 +369,7 @@ async fn setup_new_worker(
     props: PlatformProperties,
 ) -> Result<mpsc::UnboundedReceiver<UpdateForWorker>, Error> {
     let (tx, mut rx) = mpsc::unbounded_channel();
-    let worker = Worker::new(worker_id.clone(), props, tx, NOW_TIME);
+    let worker = Worker::new(worker_id.clone(), props, tx, NOW_TIME, 0);
     scheduler
         .add_worker(worker)
         .await

nativelink-scheduler/tests/simple_scheduler_test.rs

@@ -92,7 +92,7 @@ async fn setup_new_worker(
     props: PlatformProperties,
 ) -> Result<mpsc::UnboundedReceiver<UpdateForWorker>, Error> {
     let (tx, mut rx) = mpsc::unbounded_channel();
-    let worker = Worker::new(worker_id.clone(), props, tx, NOW_TIME);
+    let worker = Worker::new(worker_id.clone(), props, tx, NOW_TIME, 0);
     scheduler
         .add_worker(worker)
         .await

nativelink-service/src/worker_api_server.rs

@@ -189,6 +189,7 @@ impl WorkerApiServer {
                 platform_properties,
                 tx,
                 (self.now_fn)()?.as_secs(),
+                connect_worker_request.max_inflight_tasks,
             );
             self.scheduler
                 .add_worker(worker)

nativelink-service/tests/worker_api_server_test.rs

@@ -45,6 +45,7 @@ use nativelink_util::action_messages::{
 use nativelink_util::common::DigestInfo;
 use nativelink_util::digest_hasher::DigestHasherFunc;
 use nativelink_util::operation_state_manager::{UpdateOperationType, WorkerStateManager};
+use nativelink_util::platform_properties::PlatformProperties;
 use pretty_assertions::assert_eq;
 use tokio::join;
 use tokio::sync::{Notify, mpsc};
@@ -143,6 +144,14 @@ const fn static_now_fn() -> Result<Duration, Error> {
 }
 
 async fn setup_api_server(worker_timeout: u64, now_fn: NowFn) -> Result<TestContext, Error> {
+    setup_api_server_with_task_limit(worker_timeout, now_fn, 0).await
+}
+
+async fn setup_api_server_with_task_limit(
+    worker_timeout: u64,
+    now_fn: NowFn,
+    max_worker_tasks: u64,
+) -> Result<TestContext, Error> {
     const SCHEDULER_NAME: &str = "DUMMY_SCHEDULE_NAME";
 
     const UUID_SIZE: usize = 36;
@@ -172,7 +181,10 @@ async fn setup_api_server(worker_timeout: u64, now_fn: NowFn) -> Result<TestCont
     )
     .err_tip(|| "Error creating WorkerApiServer")?;
 
-    let connect_worker_request = ConnectWorkerRequest::default();
+    let connect_worker_request = ConnectWorkerRequest {
+        max_inflight_tasks: max_worker_tasks,
+        ..Default::default()
+    };
     let (tx, rx) = mpsc::channel(1);
     tx.send(Update::ConnectWorkerRequest(connect_worker_request))
         .await
@@ -545,3 +557,75 @@ pub async fn execution_response_success_test() -> Result<(), Box<dyn core::error
     }
     Ok(())
 }
+
+#[nativelink_test]
+pub async fn workers_only_allow_max_tasks() -> Result<(), Box<dyn core::error::Error>> {
+    let test_context =
+        setup_api_server_with_task_limit(BASE_WORKER_TIMEOUT_S, Box::new(static_now_fn), 1).await?;
+
+    let selected_worker = test_context
+        .scheduler
+        .find_worker_for_action(&PlatformProperties::new(HashMap::new()), true)
+        .await;
+    assert_eq!(
+        selected_worker,
+        Some(test_context.worker_id.clone()),
+        "Expected worker to permit tasks to begin with"
+    );
+
+    let action_digest = DigestInfo::new([7u8; 32], 123);
+    let instance_name = "instance_name".to_string();
+
+    let unique_qualifier = ActionUniqueQualifier::Uncacheable(ActionUniqueKey {
+        instance_name: instance_name.clone(),
+        digest_function: DigestHasherFunc::Sha256,
+        digest: action_digest,
+    });
+
+    let action_info = Arc::new(ActionInfo {
+        command_digest: DigestInfo::new([0u8; 32], 0),
+        input_root_digest: DigestInfo::new([0u8; 32], 0),
+        timeout: Duration::MAX,
+        platform_properties: HashMap::new(),
+        priority: 0,
+        load_timestamp: make_system_time(0),
+        insert_timestamp: make_system_time(0),
+        unique_qualifier,
+    });
+
+    let platform_properties = test_context
+        .scheduler
+        .get_platform_property_manager()
+        .make_platform_properties(action_info.platform_properties.clone())
+        .err_tip(|| "Failed to make platform properties in SimpleScheduler::do_try_match")?;
+
+    let expected_operation_id = OperationId::default();
+
+    test_context
+        .scheduler
+        .worker_notify_run_action(
+            test_context.worker_id.clone(),
+            expected_operation_id,
+            ActionInfoWithProps {
+                inner: action_info,
+                platform_properties,
+            },
+        )
+        .await
+        .unwrap();
+
+    let selected_worker = test_context
+        .scheduler
+        .find_worker_for_action(&PlatformProperties::new(HashMap::new()), true)
+        .await;
+    assert_eq!(
+        selected_worker, None,
+        "Expected not to be able to give worker a second task"
+    );
+
+    assert!(logs_contain(
+        "cannot accept work: is_paused=false, is_draining=false, inflight=1/1"
+    ));
+
+    Ok(())
+}