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 <swang@cs.berkeley.edu> * second revision, with linting. all tests are passing locally * Renamed SafeToDeleteEntry method in SchedulingKeyEntry Co-authored-by: Stephanie Wang <swang@cs.berkeley.edu> * 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 <swang@cs.berkeley.edu>
2026-07-21 12:50:45 +08:00 · 2020-08-30 18:49:25 -07:00
parent 3e5cac80d8
commit 05fe6dc278
3 changed files with 473 additions and 75 deletions
@@ -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<CoreWorkerMemoryStore> 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<CoreWorkerMemoryStore> 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<MockRayletClient>();
+  auto worker_client = std::make_shared<MockWorkerClient>();
+  auto store = std::make_shared<CoreWorkerMemoryStore>();
+  auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
+      [&](const rpc::Address &addr) { return worker_client; });
+  auto task_finisher = std::make_shared<MockTaskFinisher>();
+  auto actor_creator = std::make_shared<MockActorCreator>();
+
+  // 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<std::string, double> empty_resources;
+  ray::FunctionDescriptor empty_descriptor =
+      ray::FunctionDescriptorBuilder::BuildPython("", "", "", "");
+  std::vector<TaskSpecification> 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
@@ -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<TaskSpecification>()).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<WorkerLeaseInterface> lease_client) {
+    const rpc::WorkerAddress &addr, std::shared_ptr<WorkerLeaseInterface> lease_client,
+    const google::protobuf::RepeatedPtrField<rpc::ResourceMapEntry> &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<rpc::ResourceMapEntry> &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 &current_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() <=
@@ -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<WorkerLeaseInterface> lease_client)
-      EXCLUSIVE_LOCKS_REQUIRED(mu_);
+  void AddWorkerLeaseClient(
+      const rpc::WorkerAddress &addr, std::shared_ptr<WorkerLeaseInterface> lease_client,
+      const google::protobuf::RepeatedPtrField<rpc::ResourceMapEntry> &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<rpc::ResourceMapEntry>
                          &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<WorkerLeaseInterface> lease_client_;
-    int64_t lease_expiration_time_;
-    uint32_t tasks_in_flight_;
+    std::shared_ptr<WorkerLeaseInterface> lease_client;
+    int64_t lease_expiration_time;
+    uint32_t tasks_in_flight;
+    google::protobuf::RepeatedPtrField<rpc::ResourceMapEntry> assigned_resources;
+    SchedulingKey scheduling_key;

-    LeaseEntry(std::shared_ptr<WorkerLeaseInterface> 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<WorkerLeaseInterface> lease_client = nullptr,
+        int64_t lease_expiration_time = 0, uint32_t tasks_in_flight = 0,
+        google::protobuf::RepeatedPtrField<rpc::ResourceMapEntry> assigned_resources =
+            google::protobuf::RepeatedPtrField<rpc::ResourceMapEntry>(),
+        SchedulingKey scheduling_key = std::make_tuple(0, std::vector<ObjectID>(),
+                                                       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<rpc::WorkerAddress, LeaseEntry> worker_to_lease_entry_
      GUARDED_BY(mu_);

-  // Keeps track of pending worker lease requests to the raylet.
-  absl::flat_hash_map<SchedulingKey,
-                      std::pair<std::shared_ptr<WorkerLeaseInterface>, TaskID>>
-      pending_lease_requests_ GUARDED_BY(mu_);
+  struct SchedulingKeyEntry {
+    // Keep track of pending worker lease requests to the raylet.
+    std::pair<std::shared_ptr<WorkerLeaseInterface>, 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<TaskSpecification> task_queue = std::deque<TaskSpecification>();
+    // 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<rpc::WorkerAddress> active_workers =
+        absl::flat_hash_set<rpc::WorkerAddress>();
+    // 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<SchedulingKey, std::deque<TaskSpecification>> 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<SchedulingKey, SchedulingKeyEntry> scheduling_key_entries_
      GUARDED_BY(mu_);

  // Tasks that were cancelled while being resolved.