Revert ExecutionComplete early scheduling optimization (#1929)

* Revert "Early scheduling (#1904)"

This reverts commit 85c279a.

* Revert Early Scheduling to fix the Build Panics

Author: amankrx

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

@@ -57,10 +57,6 @@ service WorkerApi {
 
     /// Informs the scheduler about the result of an execution request.
     rpc ExecutionResponse(ExecuteResult) returns (google.protobuf.Empty);
-
-    /// Notify the scheduler that an execution request is in the upload phase
-    /// and therefore a new execution may be scheduled.
-    rpc ExecutionComplete(ExecuteComplete) returns (google.protobuf.Empty);
 }
 
 /// Request object for keep alive requests.
@@ -127,21 +123,6 @@ message ExecuteResult {
     reserved 9; // NextId.
 }
 
-/// A notification that an ExecutionRequest is in the upload phase.
-message ExecuteComplete {
-    /// ID of the worker making the request.
-    string worker_id = 1;
-
-    /// The `instance_name` this task was initially assigned to. This is set by the client
-    /// that initially sent the job as part of the BRE protocol.
-    string instance_name = 2;
-
-    /// The operation ID that was executed.
-    string operation_id = 3;
-
-    reserved 4; // NextId.
-}
-
 /// Result sent back from the server when a node connects.
 message ConnectionResult {
     /// The internal ID given to the newly connected node.

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

@@ -86,20 +86,6 @@ pub mod execute_result {
         InternalError(super::super::super::super::super::super::google::rpc::Status),
     }
 }
-/// / A notification that an ExecutionRequest is in the upload phase.
-#[derive(Clone, PartialEq, ::prost::Message)]
-pub struct ExecuteComplete {
-    /// / ID of the worker making the request.
-    #[prost(string, tag = "1")]
-    pub worker_id: ::prost::alloc::string::String,
-    /// / The `instance_name` this task was initially assigned to. This is set by the client
-    /// / that initially sent the job as part of the BRE protocol.
-    #[prost(string, tag = "2")]
-    pub instance_name: ::prost::alloc::string::String,
-    /// / The operation ID that was executed.
-    #[prost(string, tag = "3")]
-    pub operation_id: ::prost::alloc::string::String,
-}
 /// / Result sent back from the server when a node connects.
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct ConnectionResult {
@@ -402,34 +388,6 @@ pub mod worker_api_client {
                 );
             self.inner.unary(req, path, codec).await
         }
-        /// / Notify the scheduler that an execution request is in the upload phase
-        /// / and therefore a new execution may be scheduled.
-        pub async fn execution_complete(
-            &mut self,
-            request: impl tonic::IntoRequest<super::ExecuteComplete>,
-        ) -> std::result::Result<tonic::Response<()>, tonic::Status> {
-            self.inner
-                .ready()
-                .await
-                .map_err(|e| {
-                    tonic::Status::unknown(
-                        format!("Service was not ready: {}", e.into()),
-                    )
-                })?;
-            let codec = tonic::codec::ProstCodec::default();
-            let path = http::uri::PathAndQuery::from_static(
-                "/com.github.trace_machina.nativelink.remote_execution.WorkerApi/ExecutionComplete",
-            );
-            let mut req = request.into_request();
-            req.extensions_mut()
-                .insert(
-                    GrpcMethod::new(
-                        "com.github.trace_machina.nativelink.remote_execution.WorkerApi",
-                        "ExecutionComplete",
-                    ),
-                );
-            self.inner.unary(req, path, codec).await
-        }
     }
 }
 /// Generated server implementations.
@@ -491,12 +449,6 @@ pub mod worker_api_server {
             &self,
             request: tonic::Request<super::ExecuteResult>,
         ) -> std::result::Result<tonic::Response<()>, tonic::Status>;
-        /// / Notify the scheduler that an execution request is in the upload phase
-        /// / and therefore a new execution may be scheduled.
-        async fn execution_complete(
-            &self,
-            request: tonic::Request<super::ExecuteComplete>,
-        ) -> std::result::Result<tonic::Response<()>, tonic::Status>;
     }
     /// / This API describes how schedulers communicate with Worker nodes.
     /// /
@@ -760,51 +712,6 @@ pub mod worker_api_server {
                     };
                     Box::pin(fut)
                 }
-                "/com.github.trace_machina.nativelink.remote_execution.WorkerApi/ExecutionComplete" => {
-                    #[allow(non_camel_case_types)]
-                    struct ExecutionCompleteSvc<T: WorkerApi>(pub Arc<T>);
-                    impl<
-                        T: WorkerApi,
-                    > tonic::server::UnaryService<super::ExecuteComplete>
-                    for ExecutionCompleteSvc<T> {
-                        type Response = ();
-                        type Future = BoxFuture<
-                            tonic::Response<Self::Response>,
-                            tonic::Status,
-                        >;
-                        fn call(
-                            &mut self,
-                            request: tonic::Request<super::ExecuteComplete>,
-                        ) -> Self::Future {
-                            let inner = Arc::clone(&self.0);
-                            let fut = async move {
-                                <T as WorkerApi>::execution_complete(&inner, request).await
-                            };
-                            Box::pin(fut)
-                        }
-                    }
-                    let accept_compression_encodings = self.accept_compression_encodings;
-                    let send_compression_encodings = self.send_compression_encodings;
-                    let max_decoding_message_size = self.max_decoding_message_size;
-                    let max_encoding_message_size = self.max_encoding_message_size;
-                    let inner = self.inner.clone();
-                    let fut = async move {
-                        let method = ExecutionCompleteSvc(inner);
-                        let codec = tonic::codec::ProstCodec::default();
-                        let mut grpc = tonic::server::Grpc::new(codec)
-                            .apply_compression_config(
-                                accept_compression_encodings,
-                                send_compression_encodings,
-                            )
-                            .apply_max_message_size_config(
-                                max_decoding_message_size,
-                                max_encoding_message_size,
-                            );
-                        let res = grpc.unary(method, req).await;
-                        Ok(res)
-                    };
-                    Box::pin(fut)
-                }
                 _ => {
                     Box::pin(async move {
                         let mut response = http::Response::new(

nativelink-scheduler/src/api_worker_scheduler.rs

@@ -13,7 +13,6 @@
 // limitations under the License.
 
 use core::ops::{Deref, DerefMut};
-use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use async_lock::Mutex;
@@ -82,10 +81,6 @@ struct ApiWorkerSchedulerImpl {
     #[metric(group = "workers")]
     workers: Workers,
 
-    /// A set of operations that have notified completion, but are still
-    /// uploading results.
-    pending_results: HashMap<WorkerId, HashSet<OperationId>>,
-
     /// The worker state manager.
     #[metric(group = "worker_state_manager")]
     worker_state_manager: Arc<dyn WorkerStateManager>,
@@ -173,7 +168,6 @@ impl ApiWorkerSchedulerImpl {
     /// running.
     fn remove_worker(&mut self, worker_id: &WorkerId) -> Option<Worker> {
         let result = self.workers.pop(worker_id);
-        self.pending_results.remove(worker_id);
         self.worker_change_notify.notify_one();
         result
     }
@@ -247,17 +241,13 @@ impl ApiWorkerSchedulerImpl {
         })?;
 
         // Ensure the worker is supposed to be running the operation.
-        let pending_completion = !worker.running_action_infos.contains_key(operation_id);
-        if pending_completion
-            && !self
-                .pending_results
-                .get(worker_id)
-                .is_some_and(|pending_operations| pending_operations.contains(operation_id))
-        {
-            return Err(make_err!(
+        if !worker.running_action_infos.contains_key(operation_id) {
+            let err = make_err!(
                 Code::Internal,
                 "Operation {operation_id} should not be running on worker {worker_id} in SimpleScheduler::update_action"
-            ));
+            );
+            return Result::<(), _>::Err(err.clone())
+                .merge(self.immediate_evict_worker(worker_id, err, false).await);
         }
 
         let (is_finished, due_to_backpressure) = match &update {
@@ -293,32 +283,6 @@ impl ApiWorkerSchedulerImpl {
             return Ok(());
         }
 
-        if pending_completion {
-            // This is absolutely always true, but pattern match anyway.
-            if let Some(pending_operations) = self.pending_results.get_mut(worker_id) {
-                pending_operations.remove(operation_id);
-                if pending_operations.is_empty() {
-                    self.pending_results.remove(worker_id);
-                }
-                return Ok(());
-            }
-        }
-
-        Self::complete_worker_action(
-            worker,
-            operation_id,
-            due_to_backpressure,
-            &self.worker_change_notify,
-        )
-        .await
-    }
-
-    async fn complete_worker_action(
-        worker: &mut Worker,
-        operation_id: &OperationId,
-        due_to_backpressure: bool,
-        notify: &Notify,
-    ) -> Result<(), Error> {
         // Clear this action from the current worker if finished.
         let complete_action_res = {
             let was_paused = !worker.can_accept_work();
@@ -333,43 +297,11 @@ impl ApiWorkerSchedulerImpl {
             complete_action_res
         };
 
-        notify.notify_one();
+        self.worker_change_notify.notify_one();
 
         complete_action_res
     }
 
-    async fn notify_complete(
-        &mut self,
-        worker_id: &WorkerId,
-        operation_id: &OperationId,
-    ) -> Result<(), Error> {
-        let worker = self.workers.get_mut(worker_id).err_tip(|| {
-            format!("Worker {worker_id} does not exist in SimpleScheduler::notify_complete")
-        })?;
-
-        // Ensure the worker is supposed to be running the operation.
-        if !worker.running_action_infos.contains_key(operation_id) {
-            let err = make_err!(
-                Code::Internal,
-                "Operation {operation_id} should not be running on worker {worker_id} in SimpleScheduler::update_action"
-            );
-            return Err(err);
-        }
-
-        match self.pending_results.entry(worker_id.clone()) {
-            std::collections::hash_map::Entry::Occupied(mut occupied_entry) => {
-                occupied_entry.get_mut().insert(operation_id.clone());
-            }
-            std::collections::hash_map::Entry::Vacant(vacant_entry) => {
-                vacant_entry
-                    .insert(HashSet::new())
-                    .insert(operation_id.clone());
-            }
-        }
-
-        Self::complete_worker_action(worker, operation_id, false, &self.worker_change_notify).await
-    }
-
     /// Notifies the specified worker to run the given action and handles errors by evicting
     /// the worker if the notification fails.
     async fn worker_notify_run_action(
@@ -476,7 +408,6 @@ impl ApiWorkerScheduler {
         Arc::new(Self {
             inner: Mutex::new(ApiWorkerSchedulerImpl {
                 workers: Workers(LruCache::unbounded()),
-                pending_results: HashMap::new(),
                 worker_state_manager: worker_state_manager.clone(),
                 allocation_strategy,
                 worker_change_notify,
@@ -591,15 +522,6 @@ impl WorkerScheduler for ApiWorkerScheduler {
         inner.update_action(worker_id, operation_id, update).await
     }
 
-    async fn notify_complete(
-        &self,
-        worker_id: &WorkerId,
-        operation_id: &OperationId,
-    ) -> Result<(), Error> {
-        let mut inner = self.inner.lock().await;
-        inner.notify_complete(worker_id, operation_id).await
-    }
-
     async fn worker_keep_alive_received(
         &self,
         worker_id: &WorkerId,

nativelink-scheduler/src/simple_scheduler.rs

@@ -518,16 +518,6 @@ impl WorkerScheduler for SimpleScheduler {
             .await
     }
 
-    async fn notify_complete(
-        &self,
-        worker_id: &WorkerId,
-        operation_id: &OperationId,
-    ) -> Result<(), Error> {
-        self.worker_scheduler
-            .notify_complete(worker_id, operation_id)
-            .await
-    }
-
     async fn worker_keep_alive_received(
         &self,
         worker_id: &WorkerId,

nativelink-scheduler/src/worker_scheduler.rs

@@ -40,13 +40,6 @@ pub trait WorkerScheduler: Sync + Send + Unpin + RootMetricsComponent + 'static
         update: UpdateOperationType,
     ) -> Result<(), Error>;
 
-    /// Notify that the operation has completed execution, but not uploaded yet.
-    async fn notify_complete(
-        &self,
-        worker_id: &WorkerId,
-        operation_id: &OperationId,
-    ) -> Result<(), Error>;
-
     /// Event for when the keep alive message was received from the worker.
     async fn worker_keep_alive_received(
         &self,

nativelink-service/src/worker_api_server.rs

@@ -27,7 +27,7 @@ use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::
     WorkerApi, WorkerApiServer as Server,
 };
 use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{
-    execute_result, ConnectWorkerRequest, ExecuteComplete, ExecuteResult, GoingAwayRequest, KeepAliveRequest, UpdateForWorker
+    execute_result, ConnectWorkerRequest, ExecuteResult, GoingAwayRequest, KeepAliveRequest, UpdateForWorker
 };
 use nativelink_scheduler::worker::Worker;
 use nativelink_scheduler::worker_scheduler::WorkerScheduler;
@@ -258,19 +258,6 @@ impl WorkerApiServer {
         }
         Ok(Response::new(()))
     }
-
-    async fn inner_execution_complete(
-        &self,
-        execute_complete: ExecuteComplete,
-    ) -> Result<Response<()>, Error> {
-        let worker_id: WorkerId = execute_complete.worker_id.into();
-        let operation_id = OperationId::from(execute_complete.operation_id);
-        self.scheduler
-            .notify_complete(&worker_id, &operation_id)
-            .await
-            .err_tip(|| format!("Failed to operation {operation_id:?}"))?;
-        Ok(Response::new(()))
-    }
 }
 
 #[tonic::async_trait]
@@ -344,20 +331,4 @@ impl WorkerApi for WorkerApiServer {
             .await
             .map_err(Into::into)
     }
-
-    #[instrument(
-        err,
-        ret(level = Level::DEBUG),
-        level = Level::ERROR,
-        skip_all,
-        fields(request = ?grpc_request.get_ref())
-    )]
-    async fn execution_complete(
-        &self,
-        grpc_request: Request<ExecuteComplete>,
-    ) -> Result<Response<()>, Status> {
-        self.inner_execution_complete(grpc_request.into_inner())
-            .await
-            .map_err(Into::into)
-    }
 }