From 05fe6dc278126352e830a21b408027bd7a8cfd78 Mon Sep 17 00:00:00 2001 From: Gabriele Oliaro Date: Sun, 30 Aug 2020 21:49:25 -0400 Subject: [PATCH] Keeping pipelines full (#10225) * requesting new workers only when pipelines to existing ones are full * linting * added unit testing & linting * finished refactoring to consolidate all the fields that belong to a SchedulingKey into a single hashmap * linting * fixed bugs introduced by rebasing from new upstream master * changes as part of the PR review process * Fix typo in src/ray/core_worker/transport/direct_task_transport.cc Co-authored-by: fangfengbin <869218239a@zju.edu.cn> * Fixed comment in src/ray/core_worker/transport/direct_task_transport.cc Co-authored-by: Stephanie Wang * second revision, with linting. all tests are passing locally * Renamed SafeToDeleteEntry method in SchedulingKeyEntry Co-authored-by: Stephanie Wang * all new revisions but the memory leak check. performed linting. * added checks to make sure scheduling_key_entries does not leak memory * linting. all checks passing locally * edited CheckNoSchedulingKeyEntries function * linting * fixed build error on mac * created public version of CheckNoSchedulingKeyEntries to acquire the lock * linting Co-authored-by: fangfengbin <869218239a@zju.edu.cn> Co-authored-by: Stephanie Wang --- .../test/direct_task_transport_test.cc | 261 +++++++++++++++++- .../transport/direct_task_transport.cc | 192 +++++++++---- .../transport/direct_task_transport.h | 95 +++++-- 3 files changed, 473 insertions(+), 75 deletions(-) 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 388eeb36b..ff21188ad 100644 --- a/src/ray/core_worker/test/direct_task_transport_test.cc +++ b/src/ray/core_worker/test/direct_task_transport_test.cc @@ -365,6 +365,10 @@ TEST(DirectTaskTransportTest, TestSubmitOneTask) { ASSERT_EQ(task_finisher->num_tasks_failed, 0); ASSERT_EQ(raylet_client->num_leases_canceled, 0); ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease()); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestHandleTaskFailure) { @@ -395,6 +399,10 @@ TEST(DirectTaskTransportTest, TestHandleTaskFailure) { ASSERT_EQ(task_finisher->num_tasks_failed, 1); ASSERT_EQ(raylet_client->num_leases_canceled, 0); ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease()); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestConcurrentWorkerLeases) { @@ -446,6 +454,10 @@ TEST(DirectTaskTransportTest, TestConcurrentWorkerLeases) { ASSERT_EQ(task_finisher->num_tasks_failed, 0); ASSERT_EQ(raylet_client->num_leases_canceled, 0); ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease()); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestReuseWorkerLease) { @@ -490,7 +502,6 @@ TEST(DirectTaskTransportTest, TestReuseWorkerLease) { 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()); ASSERT_EQ(raylet_client->num_workers_returned, 1); @@ -504,6 +515,10 @@ TEST(DirectTaskTransportTest, TestReuseWorkerLease) { ASSERT_EQ(task_finisher->num_tasks_failed, 0); ASSERT_EQ(raylet_client->num_leases_canceled, 1); ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease()); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestRetryLeaseCancellation) { @@ -560,6 +575,10 @@ TEST(DirectTaskTransportTest, TestRetryLeaseCancellation) { ASSERT_EQ(raylet_client->num_workers_disconnected, 0); ASSERT_EQ(task_finisher->num_tasks_complete, 3); ASSERT_EQ(task_finisher->num_tasks_failed, 0); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestConcurrentCancellationAndSubmission) { @@ -613,6 +632,10 @@ TEST(DirectTaskTransportTest, TestConcurrentCancellationAndSubmission) { ASSERT_EQ(raylet_client->num_workers_returned, 2); ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease()); ASSERT_EQ(raylet_client->num_leases_canceled, 1); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestWorkerNotReusedOnError) { @@ -657,6 +680,10 @@ TEST(DirectTaskTransportTest, TestWorkerNotReusedOnError) { ASSERT_EQ(task_finisher->num_tasks_failed, 1); ASSERT_EQ(raylet_client->num_leases_canceled, 0); ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease()); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestWorkerNotReturnedOnExit) { @@ -691,6 +718,10 @@ TEST(DirectTaskTransportTest, TestWorkerNotReturnedOnExit) { ASSERT_EQ(task_finisher->num_tasks_failed, 0); ASSERT_EQ(raylet_client->num_leases_canceled, 0); ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease()); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestSpillback) { @@ -749,6 +780,10 @@ TEST(DirectTaskTransportTest, TestSpillback) { ASSERT_EQ(remote_client.second->num_leases_canceled, 0); ASSERT_FALSE(remote_client.second->ReplyCancelWorkerLease()); } + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestSpillbackRoundTrip) { @@ -813,6 +848,10 @@ TEST(DirectTaskTransportTest, TestSpillbackRoundTrip) { ASSERT_EQ(remote_client.second->num_leases_canceled, 0); ASSERT_FALSE(remote_client.second->ReplyCancelWorkerLease()); } + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } // Helper to run a test that checks that 'same1' and 'same2' are treated as the same @@ -841,13 +880,16 @@ void TestSchedulingKey(const std::shared_ptr store, ASSERT_EQ(worker_client->callbacks.size(), 1); // Another worker is requested because same2 is pending. ASSERT_EQ(raylet_client->num_workers_requested, 3); + ASSERT_EQ(raylet_client->num_leases_canceled, 0); // same1 runs successfully. Worker isn't returned. ASSERT_TRUE(worker_client->ReplyPushTask()); ASSERT_EQ(raylet_client->num_workers_returned, 0); ASSERT_EQ(raylet_client->num_workers_disconnected, 0); - // taske1_2 is pushed. + // same2 is pushed. ASSERT_EQ(worker_client->callbacks.size(), 1); + ASSERT_EQ(raylet_client->num_leases_canceled, 1); + ASSERT_TRUE(raylet_client->ReplyCancelWorkerLease()); // different is pushed. ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1001, ClientID::Nil())); @@ -863,6 +905,16 @@ void TestSchedulingKey(const std::shared_ptr store, ASSERT_TRUE(worker_client->ReplyPushTask()); ASSERT_EQ(raylet_client->num_workers_returned, 2); ASSERT_EQ(raylet_client->num_workers_disconnected, 0); + + ASSERT_EQ(raylet_client->num_leases_canceled, 1); + + // Trigger reply to RequestWorkerLease to remove the canceled pending lease request + ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1002, ClientID::Nil(), true)); + ASSERT_EQ(raylet_client->num_workers_returned, 2); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestSchedulingKeys) { @@ -983,6 +1035,10 @@ TEST(DirectTaskTransportTest, TestWorkerLeaseTimeout) { ASSERT_EQ(raylet_client->num_workers_disconnected, 1); ASSERT_EQ(raylet_client->num_leases_canceled, 0); ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease()); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestKillExecutingTask) { @@ -1032,6 +1088,10 @@ TEST(DirectTaskTransportTest, TestKillExecutingTask) { ASSERT_EQ(raylet_client->num_workers_disconnected, 0); ASSERT_EQ(task_finisher->num_tasks_complete, 1); ASSERT_EQ(task_finisher->num_tasks_failed, 1); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestKillPendingTask) { @@ -1061,6 +1121,13 @@ TEST(DirectTaskTransportTest, TestKillPendingTask) { ASSERT_EQ(task_finisher->num_tasks_failed, 1); ASSERT_EQ(raylet_client->num_leases_canceled, 1); ASSERT_TRUE(raylet_client->ReplyCancelWorkerLease()); + + // Trigger reply to RequestWorkerLease to remove the canceled pending lease request + ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1000, ClientID::Nil(), true)); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestKillResolvingTask) { @@ -1092,6 +1159,10 @@ TEST(DirectTaskTransportTest, TestKillResolvingTask) { ASSERT_EQ(raylet_client->num_workers_disconnected, 0); ASSERT_EQ(task_finisher->num_tasks_complete, 0); ASSERT_EQ(task_finisher->num_tasks_failed, 1); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestPipeliningConcurrentWorkerLeases) { @@ -1162,6 +1233,10 @@ TEST(DirectTaskTransportTest, TestPipeliningConcurrentWorkerLeases) { ASSERT_EQ(raylet_client->num_leases_canceled, 0); ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease()); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } TEST(DirectTaskTransportTest, TestPipeliningReuseWorkerLease) { @@ -1237,6 +1312,188 @@ TEST(DirectTaskTransportTest, TestPipeliningReuseWorkerLease) { ASSERT_EQ(task_finisher->num_tasks_failed, 0); ASSERT_EQ(raylet_client->num_leases_canceled, 1); ASSERT_FALSE(raylet_client->ReplyCancelWorkerLease()); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); +} + +TEST(DirectTaskTransportTest, TestPipeliningNumberOfWorkersRequested) { + rpc::Address address; + auto raylet_client = std::make_shared(); + auto worker_client = std::make_shared(); + auto store = std::make_shared(); + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); + auto task_finisher = std::make_shared(); + auto actor_creator = std::make_shared(); + + // Set max_tasks_in_flight_per_worker to a value larger than 1 to enable the pipelining + // of task submissions. This is done by passing a max_tasks_in_flight_per_worker + // parameter to the CoreWorkerDirectTaskSubmitter. + uint32_t max_tasks_in_flight_per_worker = 10; + CoreWorkerDirectTaskSubmitter submitter( + address, raylet_client, client_pool, nullptr, store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator, max_tasks_in_flight_per_worker); + + // prepare 30 tasks and save them in a vector + std::unordered_map empty_resources; + ray::FunctionDescriptor empty_descriptor = + ray::FunctionDescriptorBuilder::BuildPython("", "", "", ""); + std::vector tasks; + for (int i = 0; i < 30; i++) { + tasks.push_back(BuildTaskSpec(empty_resources, empty_descriptor)); + } + ASSERT_EQ(tasks.size(), 30); + + // Submit 4 tasks, and check that 1 worker is requested. + for (int i = 1; i <= 4; i++) { + auto task = tasks.front(); + ASSERT_TRUE(submitter.SubmitTask(task).ok()); + tasks.erase(tasks.begin()); + } + ASSERT_EQ(tasks.size(), 26); + ASSERT_EQ(raylet_client->num_workers_requested, 1); + ASSERT_EQ(task_finisher->num_tasks_complete, 0); + ASSERT_EQ(task_finisher->num_tasks_failed, 0); + ASSERT_EQ(raylet_client->num_leases_canceled, 0); + ASSERT_EQ(worker_client->callbacks.size(), 0); + + // Grant a worker lease, and check that still only 1 worker was requested. + ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1000, ClientID::Nil())); + ASSERT_EQ(raylet_client->num_workers_requested, 1); + ASSERT_EQ(raylet_client->num_workers_returned, 0); + ASSERT_EQ(raylet_client->num_workers_disconnected, 0); + ASSERT_EQ(task_finisher->num_tasks_complete, 0); + ASSERT_EQ(task_finisher->num_tasks_failed, 0); + ASSERT_EQ(raylet_client->num_leases_canceled, 0); + ASSERT_EQ(worker_client->callbacks.size(), 4); + + // Submit 6 more tasks, and check that still only 1 worker was requested. + for (int i = 1; i <= 6; i++) { + auto task = tasks.front(); + ASSERT_TRUE(submitter.SubmitTask(task).ok()); + tasks.erase(tasks.begin()); + } + ASSERT_EQ(tasks.size(), 20); + ASSERT_EQ(raylet_client->num_workers_requested, 1); + ASSERT_EQ(raylet_client->num_workers_returned, 0); + ASSERT_EQ(raylet_client->num_workers_disconnected, 0); + ASSERT_EQ(task_finisher->num_tasks_complete, 0); + ASSERT_EQ(task_finisher->num_tasks_failed, 0); + ASSERT_EQ(raylet_client->num_leases_canceled, 0); + ASSERT_EQ(worker_client->callbacks.size(), 10); + + // Submit 1 more task, and check that one more worker is requested, for a total of 2. + auto task = tasks.front(); + ASSERT_TRUE(submitter.SubmitTask(task).ok()); + tasks.erase(tasks.begin()); + ASSERT_EQ(tasks.size(), 19); + ASSERT_EQ(raylet_client->num_workers_requested, 2); + ASSERT_EQ(raylet_client->num_workers_returned, 0); + ASSERT_EQ(raylet_client->num_workers_disconnected, 0); + ASSERT_EQ(task_finisher->num_tasks_complete, 0); + ASSERT_EQ(task_finisher->num_tasks_failed, 0); + ASSERT_EQ(raylet_client->num_leases_canceled, 0); + ASSERT_EQ(worker_client->callbacks.size(), 10); + + // Grant a worker lease, and check that still only 2 workers were requested. + ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1001, ClientID::Nil())); + ASSERT_EQ(raylet_client->num_workers_requested, 2); + ASSERT_EQ(raylet_client->num_workers_returned, 0); + ASSERT_EQ(raylet_client->num_workers_disconnected, 0); + ASSERT_EQ(task_finisher->num_tasks_complete, 0); + ASSERT_EQ(task_finisher->num_tasks_failed, 0); + ASSERT_EQ(raylet_client->num_leases_canceled, 0); + ASSERT_EQ(worker_client->callbacks.size(), 11); + + // Submit 9 more tasks, and check that the total number of workers requested is still 2. + for (int i = 1; i <= 9; i++) { + auto task = tasks.front(); + ASSERT_TRUE(submitter.SubmitTask(task).ok()); + tasks.erase(tasks.begin()); + } + ASSERT_EQ(tasks.size(), 10); + ASSERT_EQ(raylet_client->num_workers_requested, 2); + ASSERT_EQ(raylet_client->num_workers_returned, 0); + ASSERT_EQ(raylet_client->num_workers_disconnected, 0); + ASSERT_EQ(task_finisher->num_tasks_complete, 0); + ASSERT_EQ(task_finisher->num_tasks_failed, 0); + ASSERT_EQ(raylet_client->num_leases_canceled, 0); + ASSERT_EQ(worker_client->callbacks.size(), 20); + + // Call ReplyPushTask on a quarter of the submitted tasks (5), and check that the total + // number of workers requested remains equal to 2. + for (int i = 1; i <= 5; i++) { + ASSERT_TRUE(worker_client->ReplyPushTask()); + } + ASSERT_EQ(raylet_client->num_workers_requested, 2); + ASSERT_EQ(raylet_client->num_workers_returned, 0); + ASSERT_EQ(raylet_client->num_workers_disconnected, 0); + ASSERT_EQ(task_finisher->num_tasks_complete, 5); + ASSERT_EQ(task_finisher->num_tasks_failed, 0); + ASSERT_EQ(raylet_client->num_leases_canceled, 0); + ASSERT_EQ(worker_client->callbacks.size(), 15); + + // Submit 5 new tasks, and check that we still have requested only 2 workers. + for (int i = 1; i <= 5; i++) { + auto task = tasks.front(); + ASSERT_TRUE(submitter.SubmitTask(task).ok()); + tasks.erase(tasks.begin()); + } + ASSERT_EQ(tasks.size(), 5); + ASSERT_EQ(raylet_client->num_workers_requested, 2); + ASSERT_EQ(raylet_client->num_workers_returned, 0); + ASSERT_EQ(raylet_client->num_workers_disconnected, 0); + ASSERT_EQ(task_finisher->num_tasks_complete, 5); + ASSERT_EQ(task_finisher->num_tasks_failed, 0); + ASSERT_EQ(raylet_client->num_leases_canceled, 0); + ASSERT_EQ(worker_client->callbacks.size(), 20); + + // Call ReplyPushTask on a quarter of the submitted tasks (5), and check that the total + // number of workers requested remains equal to 2. + for (int i = 1; i <= 5; i++) { + ASSERT_TRUE(worker_client->ReplyPushTask()); + } + ASSERT_EQ(raylet_client->num_workers_requested, 2); + ASSERT_EQ(raylet_client->num_workers_returned, 0); + ASSERT_EQ(raylet_client->num_workers_disconnected, 0); + ASSERT_EQ(task_finisher->num_tasks_complete, 10); + ASSERT_EQ(task_finisher->num_tasks_failed, 0); + ASSERT_EQ(raylet_client->num_leases_canceled, 0); + ASSERT_EQ(worker_client->callbacks.size(), 15); + + // Submit last 5 tasks, and check that the total number of workers requested is still 2 + for (int i = 1; i <= 5; i++) { + auto task = tasks.front(); + ASSERT_TRUE(submitter.SubmitTask(task).ok()); + tasks.erase(tasks.begin()); + } + ASSERT_EQ(tasks.size(), 0); + ASSERT_EQ(raylet_client->num_workers_requested, 2); + ASSERT_EQ(raylet_client->num_workers_returned, 0); + ASSERT_EQ(raylet_client->num_workers_disconnected, 0); + ASSERT_EQ(task_finisher->num_tasks_complete, 10); + ASSERT_EQ(task_finisher->num_tasks_failed, 0); + ASSERT_EQ(raylet_client->num_leases_canceled, 0); + ASSERT_EQ(worker_client->callbacks.size(), 20); + + // Execute all the resulting 20 tasks, and check that the total number of workers + // requested is 2. + for (int i = 1; i <= 20; i++) { + ASSERT_TRUE(worker_client->ReplyPushTask()); + } + ASSERT_EQ(raylet_client->num_workers_requested, 2); + ASSERT_EQ(raylet_client->num_workers_returned, 2); + ASSERT_EQ(raylet_client->num_workers_disconnected, 0); + ASSERT_EQ(task_finisher->num_tasks_complete, 30); + ASSERT_EQ(task_finisher->num_tasks_failed, 0); + ASSERT_EQ(raylet_client->num_leases_canceled, 0); + ASSERT_EQ(worker_client->callbacks.size(), 0); + + // Check that there are no entries left in the scheduling_key_entries_ hashmap. These + // would otherwise cause a memory leak. + ASSERT_TRUE(submitter.CheckNoSchedulingKeyEntriesPublic()); } } // 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 79f9cd733..3895cfbbe 100644 --- a/src/ray/core_worker/transport/direct_task_transport.cc +++ b/src/ray/core_worker/transport/direct_task_transport.cc @@ -75,12 +75,32 @@ Status CoreWorkerDirectTaskSubmitter::SubmitTask(TaskSpecification task_spec) { task_spec.GetSchedulingClass(), task_spec.GetDependencyIds(), task_spec.IsActorCreationTask() ? task_spec.ActorCreationId() : ActorID::Nil()); - auto it = task_queues_.find(scheduling_key); - if (it == task_queues_.end()) { - it = - task_queues_.emplace(scheduling_key, std::deque()).first; + auto &scheduling_key_entry = scheduling_key_entries_[scheduling_key]; + scheduling_key_entry.task_queue.push_back(task_spec); + if (!scheduling_key_entry.AllPipelinesToWorkersFull( + max_tasks_in_flight_per_worker_)) { + // The pipelines to the current workers are not full yet, so we don't need more + // workers. + + // Find a worker with a number of tasks in flight that is less than the maximum + // value (max_tasks_in_flight_per_worker_) and call OnWorkerIdle to send tasks + // to that worker + for (auto active_worker_addr : scheduling_key_entry.active_workers) { + RAY_CHECK(worker_to_lease_entry_.find(active_worker_addr) != + worker_to_lease_entry_.end()); + auto &lease_entry = worker_to_lease_entry_[active_worker_addr]; + if (!lease_entry.PipelineToWorkerFull(max_tasks_in_flight_per_worker_)) { + OnWorkerIdle(active_worker_addr, scheduling_key, false, + lease_entry.assigned_resources); + // If we find a worker with a non-full pipeline, all we need to do is to + // submit the new task to the worker in question by calling OnWorkerIdle + // once. We don't need to worry about other tasks in the queue because the + // queue cannot have other tasks in it if there are active workers with + // non-full pipelines. + break; + } + } } - it->second.push_back(task_spec); RequestNewWorkerIfNeeded(scheduling_key); } } @@ -93,32 +113,51 @@ Status CoreWorkerDirectTaskSubmitter::SubmitTask(TaskSpecification task_spec) { } void CoreWorkerDirectTaskSubmitter::AddWorkerLeaseClient( - const rpc::WorkerAddress &addr, std::shared_ptr lease_client) { + const rpc::WorkerAddress &addr, std::shared_ptr lease_client, + const google::protobuf::RepeatedPtrField &assigned_resources, + const SchedulingKey &scheduling_key) { client_cache_->GetOrConnect(addr.ToProto()); int64_t expiration = current_time_ms() + lease_timeout_ms_; - LeaseEntry new_lease_entry = LeaseEntry(std::move(lease_client), expiration, 0); + LeaseEntry new_lease_entry = LeaseEntry(std::move(lease_client), expiration, 0, + assigned_resources, scheduling_key); worker_to_lease_entry_.emplace(addr, new_lease_entry); + + auto &scheduling_key_entry = scheduling_key_entries_[scheduling_key]; + RAY_CHECK(scheduling_key_entry.active_workers.emplace(addr).second); + RAY_CHECK(scheduling_key_entry.active_workers.size() >= 1); } 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_entry_[addr]; - if (!lease_entry.lease_client_) { + if (!lease_entry.lease_client) { return; } - RAY_CHECK(lease_entry.lease_client_); + RAY_CHECK(lease_entry.lease_client); - auto queue_entry = task_queues_.find(scheduling_key); + auto &scheduling_key_entry = scheduling_key_entries_[scheduling_key]; + auto ¤t_queue = scheduling_key_entry.task_queue; // Return the worker if there was an error executing the previous task, // the previous task is an actor creation task, // there are no more applicable queued tasks, or the lease is expired. - if (was_error || queue_entry == task_queues_.end() || - current_time_ms() > lease_entry.lease_expiration_time_) { + if (was_error || current_queue.empty() || + current_time_ms() > lease_entry.lease_expiration_time) { + RAY_CHECK(scheduling_key_entry.active_workers.size() >= 1); + // Return the worker only if there are no tasks in flight - if (lease_entry.tasks_in_flight_ == 0) { + if (lease_entry.tasks_in_flight == 0) { + // Decrement the number of active workers consuming tasks from the queue associated + // with the current scheduling_key + scheduling_key_entry.active_workers.erase(addr); + if (scheduling_key_entry.CanDelete()) { + // We can safely remove the entry keyed by scheduling_key from the + // scheduling_key_entries_ hashmap. + scheduling_key_entries_.erase(scheduling_key); + } + auto status = - lease_entry.lease_client_->ReturnWorker(addr.port, addr.worker_id, was_error); + lease_entry.lease_client->ReturnWorker(addr.port, addr.worker_id, was_error); if (!status.ok()) { RAY_LOG(ERROR) << "Error returning worker to raylet: " << status.ToString(); } @@ -128,21 +167,27 @@ void CoreWorkerDirectTaskSubmitter::OnWorkerIdle( } else { auto &client = *client_cache_->GetOrConnect(addr.ToProto()); - while (!queue_entry->second.empty() && - lease_entry.tasks_in_flight_ < max_tasks_in_flight_per_worker_) { - auto task_spec = queue_entry->second.front(); + while (!current_queue.empty() && + !lease_entry.PipelineToWorkerFull(max_tasks_in_flight_per_worker_)) { + auto task_spec = current_queue.front(); lease_entry - .tasks_in_flight_++; // Increment the number of tasks in flight to the worker + .tasks_in_flight++; // Increment the number of tasks in flight to the worker + + // Increment the total number of tasks in flight to any worker associated with the + // current scheduling_key + + RAY_CHECK(scheduling_key_entry.active_workers.size() >= 1); + scheduling_key_entry.total_tasks_in_flight++; + executing_tasks_.emplace(task_spec.TaskId(), addr); PushNormalTask(addr, client, scheduling_key, task_spec, assigned_resources); - queue_entry->second.pop_front(); + current_queue.pop_front(); } // Delete the queue if it's now empty. Note that the queue cannot already be empty // 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"; + if (current_queue.empty()) { + RAY_LOG(INFO) << "Task queue empty, canceling lease request"; CancelWorkerLeaseIfNeeded(scheduling_key); } } @@ -151,17 +196,18 @@ void CoreWorkerDirectTaskSubmitter::OnWorkerIdle( void CoreWorkerDirectTaskSubmitter::CancelWorkerLeaseIfNeeded( const SchedulingKey &scheduling_key) { - auto queue_entry = task_queues_.find(scheduling_key); - if (queue_entry != task_queues_.end()) { + auto &scheduling_key_entry = scheduling_key_entries_[scheduling_key]; + auto &task_queue = scheduling_key_entry.task_queue; + if (!task_queue.empty()) { // 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()) { + auto &pending_lease_request = scheduling_key_entry.pending_lease_request; + if (pending_lease_request.first) { // There is an in-flight lease request. Cancel it. - auto &lease_client = it->second.first; - auto &lease_id = it->second.second; + auto &lease_client = pending_lease_request.first; + auto &lease_id = pending_lease_request.second; RAY_LOG(DEBUG) << "Canceling lease request " << lease_id; lease_client->CancelWorkerLease( lease_id, [this, scheduling_key](const Status &status, @@ -208,30 +254,48 @@ CoreWorkerDirectTaskSubmitter::GetOrConnectLeaseClient( void CoreWorkerDirectTaskSubmitter::RequestNewWorkerIfNeeded( const SchedulingKey &scheduling_key, const rpc::Address *raylet_address) { - if (pending_lease_requests_.find(scheduling_key) != pending_lease_requests_.end()) { + auto &scheduling_key_entry = scheduling_key_entries_[scheduling_key]; + auto &pending_lease_request = scheduling_key_entry.pending_lease_request; + + if (pending_lease_request.first) { // There's already an outstanding lease request for this type of task. return; } - auto it = task_queues_.find(scheduling_key); - if (it == task_queues_.end()) { + + auto &task_queue = scheduling_key_entry.task_queue; + if (task_queue.empty()) { // We don't have any of this type of task to run. + if (scheduling_key_entry.CanDelete()) { + // We can safely remove the entry keyed by scheduling_key from the + // scheduling_key_entries_ hashmap. + scheduling_key_entries_.erase(scheduling_key); + } + return; + } + + // Check whether we really need a new worker or whether we have + // enough room in an existing worker's pipeline to send the new tasks + if (!scheduling_key_entry.AllPipelinesToWorkersFull(max_tasks_in_flight_per_worker_)) { + // The pipelines to the current workers are not full yet, so we don't need more + // workers. + return; } auto lease_client = GetOrConnectLeaseClient(raylet_address); - TaskSpecification &resource_spec = it->second.front(); + TaskSpecification &resource_spec = task_queue.front(); TaskID task_id = resource_spec.TaskId(); - RAY_LOG(DEBUG) << "Lease requested " << task_id; lease_client->RequestWorkerLease( resource_spec, [this, scheduling_key](const Status &status, const rpc::RequestWorkerLeaseReply &reply) { absl::MutexLock lock(&mu_); - 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); + auto &scheduling_key_entry = scheduling_key_entries_[scheduling_key]; + auto &pending_lease_request = scheduling_key_entry.pending_lease_request; + RAY_CHECK(pending_lease_request.first); + auto lease_client = std::move(pending_lease_request.first); + const auto task_id = pending_lease_request.second; + pending_lease_request = std::make_pair(nullptr, TaskID::Nil()); if (status.ok()) { if (reply.canceled()) { @@ -242,8 +306,11 @@ void CoreWorkerDirectTaskSubmitter::RequestNewWorkerIfNeeded( // assign work to the worker. RAY_LOG(DEBUG) << "Lease granted " << task_id; rpc::WorkerAddress addr(reply.worker_address()); - AddWorkerLeaseClient(addr, std::move(lease_client)); auto resources_copy = reply.resource_mapping(); + + AddWorkerLeaseClient(addr, std::move(lease_client), resources_copy, + scheduling_key); + RAY_CHECK(scheduling_key_entry.active_workers.size() >= 1); OnWorkerIdle(addr, scheduling_key, /*error=*/false, resources_copy); } else { @@ -267,9 +334,7 @@ void CoreWorkerDirectTaskSubmitter::RequestNewWorkerIfNeeded( RAY_LOG(FATAL) << status.ToString(); } }); - RAY_CHECK(pending_lease_requests_ - .emplace(scheduling_key, std::make_pair(lease_client, task_id)) - .second); + pending_lease_request = std::make_pair(lease_client, task_id); } void CoreWorkerDirectTaskSubmitter::PushNormalTask( @@ -281,7 +346,6 @@ void CoreWorkerDirectTaskSubmitter::PushNormalTask( bool is_actor = task_spec.IsActorTask(); bool is_actor_creation = task_spec.IsActorCreationTask(); - RAY_LOG(DEBUG) << "Pushing normal task " << task_spec.TaskId(); // NOTE(swang): CopyFrom is needed because if we use Swap here and the task // fails, then the task data will be gone when the TaskManager attempts to // access the task. @@ -298,14 +362,28 @@ void CoreWorkerDirectTaskSubmitter::PushNormalTask( // Decrement the number of tasks in flight to the worker auto &lease_entry = worker_to_lease_entry_[addr]; - RAY_CHECK(lease_entry.tasks_in_flight_ > 0); - lease_entry.tasks_in_flight_--; + RAY_CHECK(lease_entry.tasks_in_flight > 0); + lease_entry.tasks_in_flight--; + + // Decrement the total number of tasks in flight to any worker with the current + // scheduling_key. + auto &scheduling_key_entry = scheduling_key_entries_[scheduling_key]; + RAY_CHECK(scheduling_key_entry.active_workers.size() >= 1); + RAY_CHECK(scheduling_key_entry.total_tasks_in_flight >= 1); + scheduling_key_entry.total_tasks_in_flight--; } if (reply.worker_exiting()) { // The worker is draining and will shutdown after it is done. Don't return // it to the Raylet since that will kill it early. absl::MutexLock lock(&mu_); worker_to_lease_entry_.erase(addr); + auto &scheduling_key_entry = scheduling_key_entries_[scheduling_key]; + scheduling_key_entry.active_workers.erase(addr); + if (scheduling_key_entry.CanDelete()) { + // We can safely remove the entry keyed by scheduling_key from the + // scheduling_key_entries_ hashmap. + scheduling_key_entries_.erase(scheduling_key); + } } else if (!status.ok() || !is_actor_creation) { // Successful actor creation leases the worker indefinitely from the raylet. absl::MutexLock lock(&mu_); @@ -340,17 +418,16 @@ Status CoreWorkerDirectTaskSubmitter::CancelTask(TaskSpecification task_spec, return Status::OK(); } - auto scheduled_tasks = task_queues_.find(scheduling_key); + auto &scheduling_key_entry = scheduling_key_entries_[scheduling_key]; + auto &scheduled_tasks = scheduling_key_entry.task_queue; // This cancels tasks that have completed dependencies and are awaiting // a worker lease. - if (scheduled_tasks != task_queues_.end()) { - for (auto spec = scheduled_tasks->second.begin(); - spec != scheduled_tasks->second.end(); spec++) { + if (!scheduled_tasks.empty()) { + for (auto spec = scheduled_tasks.begin(); spec != scheduled_tasks.end(); spec++) { if (spec->TaskId() == task_spec.TaskId()) { - scheduled_tasks->second.erase(spec); + scheduled_tasks.erase(spec); - if (scheduled_tasks->second.empty()) { - task_queues_.erase(scheduling_key); + if (scheduled_tasks.empty()) { CancelWorkerLeaseIfNeeded(scheduling_key); } RAY_UNUSED(task_finisher_->PendingTaskFailed(task_spec.TaskId(), @@ -359,6 +436,7 @@ Status CoreWorkerDirectTaskSubmitter::CancelTask(TaskSpecification task_spec, } } } + // This will get removed either when the RPC call to cancel is returned // or when all dependencies are resolved. RAY_CHECK(cancelled_tasks_.emplace(task_spec.TaskId()).second); @@ -367,6 +445,11 @@ Status CoreWorkerDirectTaskSubmitter::CancelTask(TaskSpecification task_spec, if (rpc_client == executing_tasks_.end()) { // This case is reached for tasks that have unresolved dependencies. // No executing tasks, so cancelling is a noop. + if (scheduling_key_entry.CanDelete()) { + // We can safely remove the entry keyed by scheduling_key from the + // scheduling_key_entries_ hashmap. + scheduling_key_entries_.erase(scheduling_key); + } return Status::OK(); } // Looks for an RPC handle for the worker executing the task. @@ -385,10 +468,11 @@ Status CoreWorkerDirectTaskSubmitter::CancelTask(TaskSpecification task_spec, request.set_intended_task_id(task_spec.TaskId().Binary()); request.set_force_kill(force_kill); client->CancelTask( - request, [this, task_spec, force_kill](const Status &status, - const rpc::CancelTaskReply &reply) { + request, [this, task_spec, scheduling_key, force_kill]( + const Status &status, const rpc::CancelTaskReply &reply) { absl::MutexLock lock(&mu_); cancelled_tasks_.erase(task_spec.TaskId()); + if (status.ok() && !reply.attempt_succeeded()) { if (cancel_retry_timer_.has_value()) { if (cancel_retry_timer_->expiry().time_since_epoch() <= diff --git a/src/ray/core_worker/transport/direct_task_transport.h b/src/ray/core_worker/transport/direct_task_transport.h index c038bb75c..25985e03f 100644 --- a/src/ray/core_worker/transport/direct_task_transport.h +++ b/src/ray/core_worker/transport/direct_task_transport.h @@ -86,6 +86,13 @@ class CoreWorkerDirectTaskSubmitter { Status CancelRemoteTask(const ObjectID &object_id, const rpc::Address &worker_addr, bool force_kill); + /// Check that the scheduling_key_entries_ hashmap is empty by calling the private + /// CheckNoSchedulingKeyEntries function after acquiring the lock. + bool CheckNoSchedulingKeyEntriesPublic() { + absl::MutexLock lock(&mu_); + return scheduling_key_entries_.empty(); + } + private: /// Schedule more work onto an idle worker or return it back to the raylet if /// no more tasks are queued for submission. If an error was encountered @@ -122,9 +129,10 @@ class CoreWorkerDirectTaskSubmitter { EXCLUSIVE_LOCKS_REQUIRED(mu_); /// Set up client state for newly granted worker lease. - void AddWorkerLeaseClient(const rpc::WorkerAddress &addr, - std::shared_ptr lease_client) - EXCLUSIVE_LOCKS_REQUIRED(mu_); + void AddWorkerLeaseClient( + const rpc::WorkerAddress &addr, std::shared_ptr lease_client, + const google::protobuf::RepeatedPtrField &assigned_resources, + const SchedulingKey &scheduling_key) EXCLUSIVE_LOCKS_REQUIRED(mu_); /// Push a task to a specific worker. void PushNormalTask(const rpc::WorkerAddress &addr, @@ -134,6 +142,11 @@ class CoreWorkerDirectTaskSubmitter { const google::protobuf::RepeatedPtrField &assigned_resources); + /// Check that the scheduling_key_entries_ hashmap is empty. + bool CheckNoSchedulingKeyEntries() const EXCLUSIVE_LOCKS_REQUIRED(mu_) { + return scheduling_key_entries_.empty(); + } + /// Address of our RPC server. rpc::Address rpc_address_; @@ -178,31 +191,75 @@ class CoreWorkerDirectTaskSubmitter { /// (1) The lease client through which the worker should be returned /// (2) The expiration time of a worker's lease. /// (3) The number of tasks that are currently in flight to the worker + /// (4) The resources assigned to the worker + /// (5) The SchedulingKey assigned to tasks that will be sent to the worker struct LeaseEntry { - std::shared_ptr lease_client_; - int64_t lease_expiration_time_; - uint32_t tasks_in_flight_; + std::shared_ptr lease_client; + int64_t lease_expiration_time; + uint32_t tasks_in_flight; + google::protobuf::RepeatedPtrField assigned_resources; + SchedulingKey scheduling_key; - LeaseEntry(std::shared_ptr lease_client = nullptr, - int64_t lease_expiration_time = 0, uint32_t tasks_in_flight = 0) - : lease_client_(lease_client), - lease_expiration_time_(lease_expiration_time), - tasks_in_flight_(tasks_in_flight) {} + LeaseEntry( + std::shared_ptr lease_client = nullptr, + int64_t lease_expiration_time = 0, uint32_t tasks_in_flight = 0, + google::protobuf::RepeatedPtrField assigned_resources = + google::protobuf::RepeatedPtrField(), + SchedulingKey scheduling_key = std::make_tuple(0, std::vector(), + ActorID::Nil())) + : lease_client(lease_client), + lease_expiration_time(lease_expiration_time), + tasks_in_flight(tasks_in_flight), + assigned_resources(assigned_resources), + scheduling_key(scheduling_key) {} + + // Check whether the pipeline to the worker associated with a LeaseEntry is full. + bool PipelineToWorkerFull(uint32_t max_tasks_in_flight_per_worker) const { + return tasks_in_flight == max_tasks_in_flight_per_worker; + } }; // Map from worker address to a LeaseEntry struct containing the lease's metadata. absl::flat_hash_map worker_to_lease_entry_ GUARDED_BY(mu_); - // Keeps track of pending worker lease requests to the raylet. - absl::flat_hash_map, TaskID>> - pending_lease_requests_ GUARDED_BY(mu_); + struct SchedulingKeyEntry { + // Keep track of pending worker lease requests to the raylet. + std::pair, TaskID> pending_lease_request = + std::make_pair(nullptr, TaskID::Nil()); + // Tasks that are queued for execution. We keep an individual queue per + // scheduling class to ensure fairness. + std::deque task_queue = std::deque(); + // Keep track of the active workers, so that we can quickly check if one of them has + // room for more tasks in flight + absl::flat_hash_set active_workers = + absl::flat_hash_set(); + // Keep track of how many tasks with this SchedulingKey are in flight, in total + uint32_t total_tasks_in_flight = 0; - // Tasks that are queued for execution. We keep individual queues per - // scheduling class to ensure fairness. - // Invariant: if a queue is in this map, it has at least one task. - absl::flat_hash_map> task_queues_ + // Check whether it's safe to delete this SchedulingKeyEntry from the + // scheduling_key_entries_ hashmap. + bool CanDelete() const { + if (!pending_lease_request.first && task_queue.empty() && + active_workers.size() == 0 && total_tasks_in_flight == 0) { + return true; + } + + return false; + } + + // Check whether the pipelines to the active workers associated with a + // SchedulingKeyEntry are all full. + bool AllPipelinesToWorkersFull(uint32_t max_tasks_in_flight_per_worker) const { + return total_tasks_in_flight == + (active_workers.size() * max_tasks_in_flight_per_worker); + } + }; + + // For each Scheduling Key, scheduling_key_entries_ contains a SchedulingKeyEntry struct + // with the queue of tasks belonging to that SchedulingKey, together with the other + // fields that are needed to orchestrate the execution of those tasks by the workers. + absl::flat_hash_map scheduling_key_entries_ GUARDED_BY(mu_); // Tasks that were cancelled while being resolved.