diff --git a/python/ray/tests/test_advanced_3.py b/python/ray/tests/test_advanced_3.py index e4bb178bc..451ca1492 100644 --- a/python/ray/tests/test_advanced_3.py +++ b/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__])) diff --git a/src/ray/common/task/task.h b/src/ray/common/task/task.h index a606badec..99396fb42 100644 --- a/src/ray/common/task/task.h +++ b/src/ray/common/task/task.h @@ -17,6 +17,8 @@ typedef std::function, const std::string &, int typedef std::function SpillbackTaskCallback; +typedef std::function 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 diff --git a/src/ray/core_worker/test/direct_task_transport_test.cc b/src/ray/core_worker/test/direct_task_transport_test.cc index 0b19f2772..8146630d0 100644 --- a/src/ray/core_worker/test/direct_task_transport_test.cc +++ b/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 &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> callbacks = {}; + std::list> 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(); + auto worker_client = std::make_shared(); + auto store = std::make_shared(); + auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto task_finisher = std::make_shared(); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, + task_finisher, ClientID::Nil(), kLongTimeout); + std::unordered_map 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(); + auto worker_client = std::make_shared(); + auto store = std::make_shared(); + auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto task_finisher = std::make_shared(); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, + task_finisher, ClientID::Nil(), kLongTimeout); + std::unordered_map 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 diff --git a/src/ray/core_worker/transport/direct_task_transport.cc b/src/ray/core_worker/transport/direct_task_transport.cc index a266be6fb..043e68bbc 100644 --- a/src/ray/core_worker/transport/direct_task_transport.cc +++ b/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 &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 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 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 diff --git a/src/ray/core_worker/transport/direct_task_transport.h b/src/ray/core_worker/transport/direct_task_transport.h index 130a91feb..163e9279c 100644 --- a/src/ray/core_worker/transport/direct_task_transport.h +++ b/src/ray/core_worker/transport/direct_task_transport.h @@ -84,12 +84,6 @@ class CoreWorkerDirectTaskSubmitter { const google::protobuf::RepeatedPtrField &assigned_resources) EXCLUSIVE_LOCKS_REQUIRED(mu_); - /// Retry a failed lease request. - void RetryLeaseRequest(Status status, - std::shared_ptr 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 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 pending_lease_requests_ GUARDED_BY(mu_); + absl::flat_hash_map, TaskID>> + pending_lease_requests_ GUARDED_BY(mu_); // Tasks that are queued for execution. We keep individual queues per // scheduling class to ensure fairness. diff --git a/src/ray/protobuf/node_manager.proto b/src/ray/protobuf/node_manager.proto index b05ee78ac..4bdc594c2 100644 --- a/src/ray/protobuf/node_manager.proto +++ b/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. diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index f48d159c4..e0be0a8a1 100644 --- a/src/ray/raylet/node_manager.cc +++ b/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, // 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. diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index 153aafadb..8adb3fbc5 100644 --- a/src/ray/raylet/node_manager.h +++ b/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, diff --git a/src/ray/raylet/raylet_client.cc b/src/ray/raylet/raylet_client.cc index 3a66cc852..a6aa423d4 100644 --- a/src/ray/raylet/raylet_client.cc +++ b/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 &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 &object_ids, const rpc::ClientCallback &callback) { diff --git a/src/ray/raylet/raylet_client.h b/src/ray/raylet/raylet_client.h index 822956f7d..0c1acb5c8 100644 --- a/src/ray/raylet/raylet_client.h +++ b/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 &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 &callback) override; + ray::Status PinObjectIDs( const rpc::Address &caller_address, const std::vector &object_ids, const ray::rpc::ClientCallback &callback); diff --git a/src/ray/rpc/node_manager/node_manager_client.h b/src/ray/rpc/node_manager/node_manager_client.h index 22e494692..595f4b5a3 100644 --- a/src/ray/rpc/node_manager/node_manager_client.h +++ b/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_, ) diff --git a/src/ray/rpc/node_manager/node_manager_server.h b/src/ray/rpc/node_manager/node_manager_server.h index 3fd91fb44..8e4e9e170 100644 --- a/src/ray/rpc/node_manager/node_manager_server.h +++ b/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;