[core] Cancel worker lease requests that are no longer needed (#7929)

* regression test

* Cancel lease requests

* unit tests

* update

* fix build

* Move unit test

* Set success

* Ref to shared_ptr

* debug

* Revert "debug"

This reverts commit 6b2c25805a8223b41ffcc2d88d903e16ea415089.

* Bad move

* Fix bad status handling

python/ray/tests/test_advanced_3.py

@@ -661,6 +661,34 @@ def test_move_log_files_to_old(shutdown_only):
     assert ray.services.remaining_processes_alive()
 
 
+def test_lease_request_leak(shutdown_only):
+    ray.init(
+        num_cpus=1,
+        _internal_config=json.dumps({
+            "initial_reconstruction_timeout_milliseconds": 200
+        }))
+    assert len(ray.objects()) == 0
+
+    @ray.remote
+    def f(x):
+        time.sleep(0.1)
+        return
+
+    # Submit pairs of tasks. Tasks in a pair can reuse the same worker leased
+    # from the raylet.
+    tasks = []
+    for _ in range(10):
+        oid = ray.put(1)
+        for _ in range(2):
+            tasks.append(f.remote(oid))
+        del oid
+    ray.get(tasks)
+
+    time.sleep(
+        1)  # Sleep for an amount longer than the reconstruction timeout.
+    assert len(ray.objects()) == 0, ray.objects()
+
+
 if __name__ == "__main__":
     import pytest
     sys.exit(pytest.main(["-v", __file__]))

src/ray/common/task/task.h

@@ -17,6 +17,8 @@ typedef std::function<void(const std::shared_ptr<void>, const std::string &, int
 typedef std::function<void(const ClientID &, const std::string &, int)>
     SpillbackTaskCallback;
 
+typedef std::function<void()> CancelTaskCallback;
+
 /// \class Task
 ///
 /// A Task represents a Ray task and a specification of its execution (e.g.,
@@ -56,6 +58,11 @@ class Task {
     on_spillback_ = callback;
   }
 
+  /// Override cancellation behaviour.
+  void OnCancellationInstead(const CancelTaskCallback &callback) {
+    on_cancellation_ = callback;
+  }
+
   /// Get the mutable specification for the task. This specification may be
   /// updated at runtime.
   ///
@@ -86,6 +93,9 @@ class Task {
   /// Returns the override spillback task callback, or nullptr.
   const SpillbackTaskCallback &OnSpillback() const { return on_spillback_; }
 
+  /// Returns the cancellation task callback, or nullptr.
+  const CancelTaskCallback &OnCancellation() const { return on_cancellation_; }
+
   std::string DebugString() const;
 
  private:
@@ -109,6 +119,9 @@ class Task {
   /// For direct task calls, overrides the spillback behaviour to send an RPC
   /// back to the submitting worker.
   mutable SpillbackTaskCallback on_spillback_ = nullptr;
+  /// For direct task calls, overrides the cancellation behaviour to send an
+  /// RPC back to the submitting worker.
+  mutable CancelTaskCallback on_cancellation_ = nullptr;
 };
 
 }  // namespace ray

src/ray/core_worker/test/direct_task_transport_test.cc

@@ -101,11 +101,21 @@ class MockRayletClient : public WorkerLeaseInterface {
     return Status::OK();
   }
 
+  ray::Status CancelWorkerLease(
+      const TaskID &task_id,
+      const rpc::ClientCallback<rpc::CancelWorkerLeaseReply> &callback) override {
+    num_leases_canceled += 1;
+    cancel_callbacks.push_back(callback);
+    return Status::OK();
+  }
+
   // Trigger reply to RequestWorkerLease.
   bool GrantWorkerLease(const std::string &address, int port,
-                        const ClientID &retry_at_raylet_id) {
+                        const ClientID &retry_at_raylet_id, bool cancel = false) {
     rpc::RequestWorkerLeaseReply reply;
-    if (!retry_at_raylet_id.IsNil()) {
+    if (cancel) {
+      reply.set_canceled(true);
+    } else if (!retry_at_raylet_id.IsNil()) {
       reply.mutable_retry_at_raylet_address()->set_ip_address(address);
       reply.mutable_retry_at_raylet_address()->set_port(port);
       reply.mutable_retry_at_raylet_address()->set_raylet_id(retry_at_raylet_id.Binary());
@@ -124,12 +134,27 @@ class MockRayletClient : public WorkerLeaseInterface {
     }
   }
 
+  bool ReplyCancelWorkerLease(bool success = true) {
+    rpc::CancelWorkerLeaseReply reply;
+    reply.set_success(success);
+    if (cancel_callbacks.size() == 0) {
+      return false;
+    } else {
+      auto callback = cancel_callbacks.front();
+      callback(Status::OK(), reply);
+      cancel_callbacks.pop_front();
+      return true;
+    }
+  }
+
   ~MockRayletClient() {}
 
   int num_workers_requested = 0;
   int num_workers_returned = 0;
   int num_workers_disconnected = 0;
+  int num_leases_canceled = 0;
   std::list<rpc::ClientCallback<rpc::RequestWorkerLeaseReply>> callbacks = {};
+  std::list<rpc::ClientCallback<rpc::CancelWorkerLeaseReply>> cancel_callbacks = {};
 };
 
 TEST(TestMemoryStore, TestPromoteToPlasma) {
@@ -327,6 +352,8 @@ TEST(DirectTaskTransportTest, TestSubmitOneTask) {
   ASSERT_EQ(raylet_client->num_workers_disconnected, 0);
   ASSERT_EQ(task_finisher->num_tasks_complete, 1);
   ASSERT_EQ(task_finisher->num_tasks_failed, 0);
+  ASSERT_EQ(raylet_client->num_leases_canceled, 0);
+  ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease());
 }
 
 TEST(DirectTaskTransportTest, TestHandleTaskFailure) {
@@ -352,6 +379,8 @@ TEST(DirectTaskTransportTest, TestHandleTaskFailure) {
   ASSERT_EQ(raylet_client->num_workers_disconnected, 1);
   ASSERT_EQ(task_finisher->num_tasks_complete, 0);
   ASSERT_EQ(task_finisher->num_tasks_failed, 1);
+  ASSERT_EQ(raylet_client->num_leases_canceled, 0);
+  ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease());
 }
 
 TEST(DirectTaskTransportTest, TestConcurrentWorkerLeases) {
@@ -398,6 +427,8 @@ TEST(DirectTaskTransportTest, TestConcurrentWorkerLeases) {
   ASSERT_EQ(raylet_client->num_workers_disconnected, 0);
   ASSERT_EQ(task_finisher->num_tasks_complete, 3);
   ASSERT_EQ(task_finisher->num_tasks_failed, 0);
+  ASSERT_EQ(raylet_client->num_leases_canceled, 0);
+  ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease());
 }
 
 TEST(DirectTaskTransportTest, TestReuseWorkerLease) {
@@ -425,16 +456,20 @@ TEST(DirectTaskTransportTest, TestReuseWorkerLease) {
   ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1000, ClientID::Nil()));
   ASSERT_EQ(worker_client->callbacks.size(), 1);
   ASSERT_EQ(raylet_client->num_workers_requested, 2);
+  ASSERT_EQ(raylet_client->num_leases_canceled, 0);
 
   // Task 1 finishes, Task 2 is scheduled on the same worker.
   ASSERT_TRUE(worker_client->ReplyPushTask());
   ASSERT_EQ(worker_client->callbacks.size(), 1);
   ASSERT_EQ(raylet_client->num_workers_returned, 0);
+  ASSERT_EQ(raylet_client->num_leases_canceled, 0);
 
   // Task 2 finishes, Task 3 is scheduled on the same worker.
   ASSERT_TRUE(worker_client->ReplyPushTask());
   ASSERT_EQ(worker_client->callbacks.size(), 1);
   ASSERT_EQ(raylet_client->num_workers_returned, 0);
+  ASSERT_EQ(raylet_client->num_leases_canceled, 1);
+  ASSERT_TRUE(raylet_client->ReplyCancelWorkerLease());
 
   // Task 3 finishes, the worker is returned.
   ASSERT_TRUE(worker_client->ReplyPushTask());
@@ -447,6 +482,111 @@ TEST(DirectTaskTransportTest, TestReuseWorkerLease) {
   ASSERT_EQ(raylet_client->num_workers_disconnected, 0);
   ASSERT_EQ(task_finisher->num_tasks_complete, 3);
   ASSERT_EQ(task_finisher->num_tasks_failed, 0);
+  ASSERT_EQ(raylet_client->num_leases_canceled, 1);
+  ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease());
+}
+
+TEST(DirectTaskTransportTest, TestRetryLeaseCancellation) {
+  rpc::Address address;
+  auto raylet_client = std::make_shared<MockRayletClient>();
+  auto worker_client = std::make_shared<MockWorkerClient>();
+  auto store = std::make_shared<CoreWorkerMemoryStore>();
+  auto factory = [&](const rpc::Address &addr) { return worker_client; };
+  auto task_finisher = std::make_shared<MockTaskFinisher>();
+  CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store,
+                                          task_finisher, ClientID::Nil(), kLongTimeout);
+  std::unordered_map<std::string, double> empty_resources;
+  ray::FunctionDescriptor empty_descriptor =
+      ray::FunctionDescriptorBuilder::BuildPython("", "", "", "");
+  TaskSpecification task1 = BuildTaskSpec(empty_resources, empty_descriptor);
+  TaskSpecification task2 = BuildTaskSpec(empty_resources, empty_descriptor);
+  TaskSpecification task3 = BuildTaskSpec(empty_resources, empty_descriptor);
+
+  ASSERT_TRUE(submitter.SubmitTask(task1).ok());
+  ASSERT_TRUE(submitter.SubmitTask(task2).ok());
+  ASSERT_TRUE(submitter.SubmitTask(task3).ok());
+  ASSERT_EQ(raylet_client->num_workers_requested, 1);
+
+  // Task 1 is pushed.
+  ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1000, ClientID::Nil()));
+  // Task 1 finishes, Task 2 is scheduled on the same worker.
+  ASSERT_TRUE(worker_client->ReplyPushTask());
+  // Task 2 finishes, Task 3 is scheduled on the same worker.
+  ASSERT_TRUE(worker_client->ReplyPushTask());
+  // Task 3 finishes, the worker is returned.
+  ASSERT_TRUE(worker_client->ReplyPushTask());
+  ASSERT_EQ(raylet_client->num_workers_returned, 1);
+
+  // Simulate the lease cancellation request failing because it arrives at the
+  // raylet before the last worker lease request has been received.
+  int i = 1;
+  for (; i <= 3; i++) {
+    ASSERT_EQ(raylet_client->num_leases_canceled, i);
+    ASSERT_TRUE(raylet_client->ReplyCancelWorkerLease(false));
+  }
+
+  // Simulate the lease cancellation request succeeding.
+  ASSERT_TRUE(raylet_client->ReplyCancelWorkerLease());
+  ASSERT_EQ(raylet_client->num_leases_canceled, i);
+  ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease());
+  ASSERT_EQ(raylet_client->num_leases_canceled, i);
+  ASSERT_TRUE(raylet_client->GrantWorkerLease("", 0, ClientID::Nil(), /*cancel=*/true));
+  ASSERT_EQ(worker_client->callbacks.size(), 0);
+  // The canceled lease is not returned.
+  ASSERT_EQ(raylet_client->num_workers_returned, 1);
+  ASSERT_EQ(raylet_client->num_workers_disconnected, 0);
+  ASSERT_EQ(task_finisher->num_tasks_complete, 3);
+  ASSERT_EQ(task_finisher->num_tasks_failed, 0);
+}
+
+TEST(DirectTaskTransportTest, TestConcurrentCancellationAndSubmission) {
+  rpc::Address address;
+  auto raylet_client = std::make_shared<MockRayletClient>();
+  auto worker_client = std::make_shared<MockWorkerClient>();
+  auto store = std::make_shared<CoreWorkerMemoryStore>();
+  auto factory = [&](const rpc::Address &addr) { return worker_client; };
+  auto task_finisher = std::make_shared<MockTaskFinisher>();
+  CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store,
+                                          task_finisher, ClientID::Nil(), kLongTimeout);
+  std::unordered_map<std::string, double> empty_resources;
+  ray::FunctionDescriptor empty_descriptor =
+      ray::FunctionDescriptorBuilder::BuildPython("", "", "", "");
+  TaskSpecification task1 = BuildTaskSpec(empty_resources, empty_descriptor);
+  TaskSpecification task2 = BuildTaskSpec(empty_resources, empty_descriptor);
+  TaskSpecification task3 = BuildTaskSpec(empty_resources, empty_descriptor);
+
+  ASSERT_TRUE(submitter.SubmitTask(task1).ok());
+  ASSERT_TRUE(submitter.SubmitTask(task2).ok());
+
+  // Task 1 is pushed.
+  ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1000, ClientID::Nil()));
+  ASSERT_EQ(raylet_client->num_workers_requested, 2);
+  // Task 1 finishes, Task 2 is scheduled on the same worker.
+  ASSERT_TRUE(worker_client->ReplyPushTask());
+
+  // Task 2's lease request gets canceled.
+  ASSERT_EQ(raylet_client->num_leases_canceled, 1);
+
+  // Task 2 finishes, the worker is returned.
+  ASSERT_TRUE(worker_client->ReplyPushTask());
+  ASSERT_EQ(raylet_client->num_workers_returned, 1);
+
+  // Another task is submitted while task 2's lease request is being canceled.
+  ASSERT_TRUE(submitter.SubmitTask(task3).ok());
+  ASSERT_EQ(raylet_client->num_workers_requested, 2);
+
+  // Task 2's lease request is canceled, a new worker is requested for task 3.
+  ASSERT_TRUE(raylet_client->ReplyCancelWorkerLease());
+  ASSERT_EQ(raylet_client->num_workers_requested, 2);
+  ASSERT_TRUE(raylet_client->GrantWorkerLease("", 0, ClientID::Nil(), /*cancel=*/true));
+  ASSERT_EQ(raylet_client->num_workers_requested, 3);
+
+  // Task 3 finishes, all workers returned.
+  ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1000, ClientID::Nil()));
+  ASSERT_TRUE(worker_client->ReplyPushTask());
+  ASSERT_EQ(raylet_client->num_workers_returned, 2);
+  ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease());
+  ASSERT_EQ(raylet_client->num_leases_canceled, 1);
 }
 
 TEST(DirectTaskTransportTest, TestWorkerNotReusedOnError) {
@@ -486,6 +626,8 @@ TEST(DirectTaskTransportTest, TestWorkerNotReusedOnError) {
   ASSERT_EQ(raylet_client->num_workers_disconnected, 1);
   ASSERT_EQ(task_finisher->num_tasks_complete, 1);
   ASSERT_EQ(task_finisher->num_tasks_failed, 1);
+  ASSERT_EQ(raylet_client->num_leases_canceled, 0);
+  ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease());
 }
 
 TEST(DirectTaskTransportTest, TestWorkerNotReturnedOnExit) {
@@ -515,6 +657,8 @@ TEST(DirectTaskTransportTest, TestWorkerNotReturnedOnExit) {
   ASSERT_EQ(raylet_client->num_workers_disconnected, 0);
   ASSERT_EQ(task_finisher->num_tasks_complete, 1);
   ASSERT_EQ(task_finisher->num_tasks_failed, 0);
+  ASSERT_EQ(raylet_client->num_leases_canceled, 0);
+  ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease());
 }
 
 TEST(DirectTaskTransportTest, TestSpillback) {
@@ -565,6 +709,12 @@ TEST(DirectTaskTransportTest, TestSpillback) {
   ASSERT_EQ(remote_lease_clients[7777]->num_workers_disconnected, 0);
   ASSERT_EQ(task_finisher->num_tasks_complete, 1);
   ASSERT_EQ(task_finisher->num_tasks_failed, 0);
+  ASSERT_EQ(raylet_client->num_leases_canceled, 0);
+  ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease());
+  for (const auto &remote_client : remote_lease_clients) {
+    ASSERT_EQ(remote_client.second->num_leases_canceled, 0);
+    ASSERT_FALSE(remote_client.second->ReplyCancelWorkerLease());
+  }
 }
 
 TEST(DirectTaskTransportTest, TestSpillbackRoundTrip) {
@@ -621,6 +771,12 @@ TEST(DirectTaskTransportTest, TestSpillbackRoundTrip) {
   ASSERT_EQ(remote_lease_clients[7777]->num_workers_disconnected, 0);
   ASSERT_EQ(task_finisher->num_tasks_complete, 1);
   ASSERT_EQ(task_finisher->num_tasks_failed, 0);
+  ASSERT_EQ(raylet_client->num_leases_canceled, 0);
+  ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease());
+  for (const auto &remote_client : remote_lease_clients) {
+    ASSERT_EQ(remote_client.second->num_leases_canceled, 0);
+    ASSERT_FALSE(remote_client.second->ReplyCancelWorkerLease());
+  }
 }
 
 // Helper to run a test that checks that 'same1' and 'same2' are treated as the same
@@ -775,6 +931,8 @@ TEST(DirectTaskTransportTest, TestWorkerLeaseTimeout) {
   ASSERT_EQ(worker_client->callbacks.size(), 0);
   ASSERT_EQ(raylet_client->num_workers_returned, 2);
   ASSERT_EQ(raylet_client->num_workers_disconnected, 1);
+  ASSERT_EQ(raylet_client->num_leases_canceled, 0);
+  ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease());
 }
 
 }  // namespace ray

src/ray/core_worker/transport/direct_task_transport.cc

@@ -56,6 +56,7 @@ void CoreWorkerDirectTaskSubmitter::OnWorkerIdle(
     const rpc::WorkerAddress &addr, const SchedulingKey &scheduling_key, bool was_error,
     const google::protobuf::RepeatedPtrField<rpc::ResourceMapEntry> &assigned_resources) {
   auto lease_entry = worker_to_lease_client_[addr];
+  RAY_CHECK(lease_entry.first);
   auto queue_entry = task_queues_.find(scheduling_key);
   // Return the worker if there was an error executing the previous task,
   // the previous task is an actor creation task,
@@ -76,11 +77,46 @@ void CoreWorkerDirectTaskSubmitter::OnWorkerIdle(
     // because this is the only place tasks are removed from it.
     if (queue_entry->second.empty()) {
       task_queues_.erase(queue_entry);
+      RAY_LOG(DEBUG) << "Task queue empty, canceling lease request";
+      CancelWorkerLeaseIfNeeded(scheduling_key);
     }
   }
   RequestNewWorkerIfNeeded(scheduling_key);
 }
 
+void CoreWorkerDirectTaskSubmitter::CancelWorkerLeaseIfNeeded(
+    const SchedulingKey &scheduling_key) {
+  auto queue_entry = task_queues_.find(scheduling_key);
+  if (queue_entry != task_queues_.end()) {
+    // There are still pending tasks, so let the worker lease request succeed.
+    return;
+  }
+
+  auto it = pending_lease_requests_.find(scheduling_key);
+  if (it != pending_lease_requests_.end()) {
+    // There is an in-flight lease request. Cancel it.
+    auto &lease_client = it->second.first;
+    auto &lease_id = it->second.second;
+    RAY_LOG(DEBUG) << "Canceling lease request " << lease_id;
+    RAY_UNUSED(lease_client->CancelWorkerLease(
+        lease_id, [this, scheduling_key](const Status &status,
+                                         const rpc::CancelWorkerLeaseReply &reply) {
+          absl::MutexLock lock(&mu_);
+          if (status.ok() && !reply.success()) {
+            // The cancellation request can fail if the raylet does not have
+            // the request queued. This can happen if: a) due to message
+            // reordering, the raylet has not yet received the worker lease
+            // request, or b) we have already returned the worker lease
+            // request. In the former case, we should try the cancellation
+            // request again. In the latter case, the in-flight lease request
+            // should already have been removed from our local state, so we no
+            // longer need to cancel.
+            CancelWorkerLeaseIfNeeded(scheduling_key);
+          }
+        }));
+  }
+}
+
 std::shared_ptr<WorkerLeaseInterface>
 CoreWorkerDirectTaskSubmitter::GetOrConnectLeaseClient(
     const rpc::Address *raylet_address) {
@@ -120,14 +156,23 @@ void CoreWorkerDirectTaskSubmitter::RequestNewWorkerIfNeeded(
   auto lease_client = GetOrConnectLeaseClient(raylet_address);
   TaskSpecification &resource_spec = it->second.front();
   TaskID task_id = resource_spec.TaskId();
-  auto status = lease_client->RequestWorkerLease(
-      resource_spec,
-      [this, lease_client, task_id, scheduling_key](
-          const Status &status, const rpc::RequestWorkerLeaseReply &reply) mutable {
+  RAY_LOG(DEBUG) << "Lease requested " << task_id;
+  RAY_UNUSED(lease_client->RequestWorkerLease(
+      resource_spec, [this, scheduling_key](const Status &status,
+                                            const rpc::RequestWorkerLeaseReply &reply) {
         absl::MutexLock lock(&mu_);
-        pending_lease_requests_.erase(scheduling_key);
+
+        auto it = pending_lease_requests_.find(scheduling_key);
+        RAY_CHECK(it != pending_lease_requests_.end());
+        auto lease_client = std::move(it->second.first);
+        const auto task_id = it->second.second;
+        pending_lease_requests_.erase(it);
+
         if (status.ok()) {
-          if (!reply.worker_address().raylet_id().empty()) {
+          if (reply.canceled()) {
+            RAY_LOG(DEBUG) << "Lease canceled " << task_id;
+            RequestNewWorkerIfNeeded(scheduling_key);
+          } else if (!reply.worker_address().raylet_id().empty()) {
             // We got a lease for a worker. Add the lease client state and try to
             // assign work to the worker.
             RAY_LOG(DEBUG) << "Lease granted " << task_id;
@@ -140,32 +185,23 @@ void CoreWorkerDirectTaskSubmitter::RequestNewWorkerIfNeeded(
             // The raylet redirected us to a different raylet to retry at.
             RequestNewWorkerIfNeeded(scheduling_key, &reply.retry_at_raylet_address());
           }
+        } else if (lease_client != local_lease_client_) {
+          // A lease request to a remote raylet failed. Retry locally if the lease is
+          // still needed.
+          // TODO(swang): Fail after some number of retries?
+          RAY_LOG(ERROR) << "Retrying attempt to schedule task at remote node. Error: "
+                         << status.ToString();
+          RequestNewWorkerIfNeeded(scheduling_key);
         } else {
-          RetryLeaseRequest(status, lease_client, scheduling_key);
+          // A local request failed. This shouldn't happen if the raylet is still alive
+          // and we don't currently handle raylet failures, so treat it as a fatal
+          // error.
+          RAY_LOG(FATAL) << status.ToString();
         }
-      });
-  if (!status.ok()) {
-    RetryLeaseRequest(status, lease_client, scheduling_key);
-  }
-  pending_lease_requests_.insert(scheduling_key);
-}
-
-void CoreWorkerDirectTaskSubmitter::RetryLeaseRequest(
-    Status status, std::shared_ptr<WorkerLeaseInterface> lease_client,
-    const SchedulingKey &scheduling_key) {
-  if (lease_client != local_lease_client_) {
-    // A lease request to a remote raylet failed. Retry locally if the lease is
-    // still needed.
-    // TODO(swang): Fail after some number of retries?
-    RAY_LOG(ERROR) << "Retrying attempt to schedule task at remote node. Error: "
-                   << status.ToString();
-    RequestNewWorkerIfNeeded(scheduling_key);
-  } else {
-    // A local request failed. This shouldn't happen if the raylet is still alive
-    // and we don't currently handle raylet failures, so treat it as a fatal
-    // error.
-    RAY_LOG(FATAL) << status.ToString();
-  }
+      }));
+  RAY_CHECK(pending_lease_requests_
+                .emplace(scheduling_key, std::make_pair(lease_client, task_id))
+                .second);
 }
 
 void CoreWorkerDirectTaskSubmitter::PushNormalTask(
@@ -185,7 +221,7 @@ void CoreWorkerDirectTaskSubmitter::PushNormalTask(
   request->mutable_task_spec()->CopyFrom(task_spec.GetMessage());
   request->mutable_resource_mapping()->CopyFrom(assigned_resources);
   request->set_intended_worker_id(addr.worker_id.Binary());
-  auto status = client.PushNormalTask(
+  RAY_UNUSED(client.PushNormalTask(
       std::move(request),
       [this, task_id, is_actor, is_actor_creation, scheduling_key, addr,
        assigned_resources](Status status, const rpc::PushTaskReply &reply) {
@@ -212,16 +248,6 @@ void CoreWorkerDirectTaskSubmitter::PushNormalTask(
         } else {
           task_finisher_->CompletePendingTask(task_id, reply, addr.ToProto());
         }
-      });
-  if (!status.ok()) {
-    RAY_LOG(ERROR) << "Error pushing task to worker: " << status.ToString();
-    {
-      absl::MutexLock lock(&mu_);
-      OnWorkerIdle(addr, scheduling_key, /*error=*/true, assigned_resources);
-    }
-    task_finisher_->PendingTaskFailed(
-        task_id, is_actor ? rpc::ErrorType::ACTOR_DIED : rpc::ErrorType::WORKER_DIED,
-        &status);
-  }
+      }));
 }
 };  // namespace ray

src/ray/core_worker/transport/direct_task_transport.h

@@ -84,12 +84,6 @@ class CoreWorkerDirectTaskSubmitter {
       const google::protobuf::RepeatedPtrField<rpc::ResourceMapEntry> &assigned_resources)
       EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
-  /// Retry a failed lease request.
-  void RetryLeaseRequest(Status status,
-                         std::shared_ptr<WorkerLeaseInterface> lease_client,
-                         const SchedulingKey &scheduling_key)
-      EXCLUSIVE_LOCKS_REQUIRED(mu_);
-
   /// Get an existing lease client or connect a new one. If a raylet_address is
   /// provided, this connects to a remote raylet. Else, this connects to the
   /// local raylet.
@@ -104,6 +98,13 @@ class CoreWorkerDirectTaskSubmitter {
                                 const rpc::Address *raylet_address = nullptr)
       EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
+  /// Cancel a pending worker lease and retry until the cancellation succeeds
+  /// (i.e., the raylet drops the request). This should be called when there
+  /// are no more tasks queued with the given scheduling key and there is an
+  /// in-flight lease request for that key.
+  void CancelWorkerLeaseIfNeeded(const SchedulingKey &scheduling_key)
+      EXCLUSIVE_LOCKS_REQUIRED(mu_);
+
   /// Set up client state for newly granted worker lease.
   void AddWorkerLeaseClient(const rpc::WorkerAddress &addr,
                             std::shared_ptr<WorkerLeaseInterface> lease_client)
@@ -161,7 +162,9 @@ class CoreWorkerDirectTaskSubmitter {
       worker_to_lease_client_ GUARDED_BY(mu_);
 
   // Keeps track of pending worker lease requests to the raylet.
-  absl::flat_hash_set<SchedulingKey> pending_lease_requests_ GUARDED_BY(mu_);
+  absl::flat_hash_map<SchedulingKey,
+                      std::pair<std::shared_ptr<WorkerLeaseInterface>, TaskID>>
+      pending_lease_requests_ GUARDED_BY(mu_);
 
   // Tasks that are queued for execution. We keep individual queues per
   // scheduling class to ensure fairness.

src/ray/protobuf/node_manager.proto

@@ -32,6 +32,9 @@ message RequestWorkerLeaseReply {
   Address retry_at_raylet_address = 2;
   // Resource mapping ids acquired by the leased worker.
   repeated ResourceMapEntry resource_mapping = 3;
+  // Whether this lease request was canceled. In this case, the
+  // client should try again if the resources are still required.
+  bool canceled = 4;
 }
 
 // Release a worker back to its raylet.
@@ -48,6 +51,19 @@ message ReturnWorkerRequest {
 message ReturnWorkerReply {
 }
 
+message CancelWorkerLeaseRequest {
+  // The task to cancel.
+  bytes task_id = 1;
+}
+
+message CancelWorkerLeaseReply {
+  // Whether the cancellation request was successful. Cancellation
+  // succeeds if the node manager has the task queued upon receiving
+  // the cancellation request, and the node manager has not yet
+  // granted the lease.
+  bool success = 1;
+}
+
 message ForwardTaskRequest {
   // The ID of the task to be forwarded.
   bytes task_id = 1;
@@ -115,6 +131,9 @@ service NodeManagerService {
   rpc RequestWorkerLease(RequestWorkerLeaseRequest) returns (RequestWorkerLeaseReply);
   // Release a worker back to its raylet.
   rpc ReturnWorker(ReturnWorkerRequest) returns (ReturnWorkerReply);
+  // Cancel a pending lease request. This only returns success if the
+  // lease request was not yet granted.
+  rpc CancelWorkerLease(CancelWorkerLeaseRequest) returns (CancelWorkerLeaseReply);
   // Forward a task and its uncommitted lineage to the remote node manager.
   rpc ForwardTask(ForwardTaskRequest) returns (ForwardTaskReply);
   // Pin the provided object IDs.

src/ray/raylet/node_manager.cc

@@ -1785,6 +1785,11 @@ void NodeManager::HandleRequestWorkerLease(const rpc::RequestWorkerLeaseRequest
         reply->mutable_retry_at_raylet_address()->set_raylet_id(spillback_to.Binary());
         send_reply_callback(Status::OK(), nullptr, nullptr);
       });
+  task.OnCancellationInstead([reply, task_id, send_reply_callback]() {
+    RAY_LOG(DEBUG) << "Task lease request canceled " << task_id;
+    reply->set_canceled(true);
+    send_reply_callback(Status::OK(), nullptr, nullptr);
+  });
   SubmitTask(task, Lineage());
 }
 
@@ -1821,6 +1826,38 @@ void NodeManager::HandleReturnWorker(const rpc::ReturnWorkerRequest &request,
   send_reply_callback(status, nullptr, nullptr);
 }
 
+void NodeManager::HandleCancelWorkerLease(const rpc::CancelWorkerLeaseRequest &request,
+                                          rpc::CancelWorkerLeaseReply *reply,
+                                          rpc::SendReplyCallback send_reply_callback) {
+  const TaskID task_id = TaskID::FromBinary(request.task_id());
+  Task removed_task;
+  TaskState removed_task_state;
+  const auto canceled =
+      local_queues_.RemoveTask(task_id, &removed_task, &removed_task_state);
+  if (!canceled) {
+    // We do not have the task. This could be because we haven't received the
+    // lease request yet, or because we already granted the lease request and
+    // it has already been returned.
+  } else {
+    if (removed_task.OnDispatch()) {
+      // We have not yet granted the worker lease. Cancel it now.
+      removed_task.OnCancellation()();
+      task_dependency_manager_.TaskCanceled(task_id);
+      task_dependency_manager_.UnsubscribeGetDependencies(task_id);
+    } else {
+      // We already granted the worker lease and sent the reply. Re-queue the
+      // task and wait for the requester to return the leased worker.
+      local_queues_.QueueTasks({removed_task}, removed_task_state);
+    }
+  }
+  // The task cancellation failed if we did not have the task queued, since
+  // this means that we may not have received the task request yet. It is
+  // successful if we did have the task queued, since we have now replied to
+  // the client that requested the lease.
+  reply->set_success(canceled);
+  send_reply_callback(Status::OK(), nullptr, nullptr);
+}
+
 void NodeManager::HandleForwardTask(const rpc::ForwardTaskRequest &request,
                                     rpc::ForwardTaskReply *reply,
                                     rpc::SendReplyCallback send_reply_callback) {
@@ -3142,6 +3179,8 @@ void NodeManager::FinishAssignTask(const std::shared_ptr<Worker> &worker,
 
     // Mark the task as running.
     // (See design_docs/task_states.rst for the state transition diagram.)
+    assigned_task.OnDispatchInstead(nullptr);
+    assigned_task.OnSpillbackInstead(nullptr);
     local_queues_.QueueTasks({assigned_task}, TaskState::RUNNING);
     // Notify the task dependency manager that we no longer need this task's
     // object dependencies.

src/ray/raylet/node_manager.h

@@ -576,6 +576,11 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
                           rpc::ReturnWorkerReply *reply,
                           rpc::SendReplyCallback send_reply_callback) override;
 
+  /// Handle a `ReturnWorker` request.
+  void HandleCancelWorkerLease(const rpc::CancelWorkerLeaseRequest &request,
+                               rpc::CancelWorkerLeaseReply *reply,
+                               rpc::SendReplyCallback send_reply_callback) override;
+
   /// Handle a `ForwardTask` request.
   void HandleForwardTask(const rpc::ForwardTaskRequest &request,
                          rpc::ForwardTaskReply *reply,

src/ray/raylet/raylet_client.cc

@@ -367,6 +367,14 @@ Status raylet::RayletClient::ReturnWorker(int worker_port, const WorkerID &worke
       });
 }
 
+ray::Status raylet::RayletClient::CancelWorkerLease(
+    const TaskID &task_id,
+    const rpc::ClientCallback<rpc::CancelWorkerLeaseReply> &callback) {
+  rpc::CancelWorkerLeaseRequest request;
+  request.set_task_id(task_id.Binary());
+  return grpc_client_->CancelWorkerLease(request, callback);
+}
+
 Status raylet::RayletClient::PinObjectIDs(
     const rpc::Address &caller_address, const std::vector<ObjectID> &object_ids,
     const rpc::ClientCallback<rpc::PinObjectIDsReply> &callback) {

src/ray/raylet/raylet_client.h

@@ -66,6 +66,10 @@ class WorkerLeaseInterface {
   virtual ray::Status ReturnWorker(int worker_port, const WorkerID &worker_id,
                                    bool disconnect_worker) = 0;
 
+  virtual ray::Status CancelWorkerLease(
+      const TaskID &task_id,
+      const rpc::ClientCallback<rpc::CancelWorkerLeaseReply> &callback) = 0;
+
   virtual ~WorkerLeaseInterface(){};
 };
 
@@ -277,6 +281,10 @@ class RayletClient : public WorkerLeaseInterface, public DependencyWaiterInterfa
   ray::Status ReturnWorker(int worker_port, const WorkerID &worker_id,
                            bool disconnect_worker) override;
 
+  ray::Status CancelWorkerLease(
+      const TaskID &task_id,
+      const rpc::ClientCallback<rpc::CancelWorkerLeaseReply> &callback) override;
+
   ray::Status PinObjectIDs(
       const rpc::Address &caller_address, const std::vector<ObjectID> &object_ids,
       const ray::rpc::ClientCallback<ray::rpc::PinObjectIDsReply> &callback);

src/ray/rpc/node_manager/node_manager_client.h

@@ -87,6 +87,9 @@ class NodeManagerWorkerClient
   /// Return a worker lease.
   RPC_CLIENT_METHOD(NodeManagerService, ReturnWorker, grpc_client_, )
 
+  /// Cancel a pending worker lease request.
+  RPC_CLIENT_METHOD(NodeManagerService, CancelWorkerLease, grpc_client_, )
+
   /// Notify the raylet to pin the provided object IDs.
   RPC_CLIENT_METHOD(NodeManagerService, PinObjectIDs, grpc_client_, )
 

src/ray/rpc/node_manager/node_manager_server.h

@@ -27,6 +27,7 @@ namespace rpc {
 #define RAY_NODE_MANAGER_RPC_HANDLERS                         \
   RPC_SERVICE_HANDLER(NodeManagerService, RequestWorkerLease) \
   RPC_SERVICE_HANDLER(NodeManagerService, ReturnWorker)       \
+  RPC_SERVICE_HANDLER(NodeManagerService, CancelWorkerLease)  \
   RPC_SERVICE_HANDLER(NodeManagerService, ForwardTask)        \
   RPC_SERVICE_HANDLER(NodeManagerService, PinObjectIDs)       \
   RPC_SERVICE_HANDLER(NodeManagerService, GetNodeStats)       \
@@ -55,6 +56,10 @@ class NodeManagerServiceHandler {
                                   ReturnWorkerReply *reply,
                                   SendReplyCallback send_reply_callback) = 0;
 
+  virtual void HandleCancelWorkerLease(const rpc::CancelWorkerLeaseRequest &request,
+                                       rpc::CancelWorkerLeaseReply *reply,
+                                       rpc::SendReplyCallback send_reply_callback) = 0;
+
   virtual void HandleForwardTask(const ForwardTaskRequest &request,
                                  ForwardTaskReply *reply,
                                  SendReplyCallback send_reply_callback) = 0;