Use one queue per resource shape in direct task transport (#6277)

src/ray/common/task/task.cc

@@ -14,18 +14,7 @@ void Task::IncrementNumForwards() { task_execution_spec_.IncrementNumForwards();
 
 const std::vector<ObjectID> &Task::GetDependencies() const { return dependencies_; }
 
-void Task::ComputeDependencies() {
-  dependencies_.clear();
-  for (size_t i = 0; i < task_spec_.NumArgs(); ++i) {
-    int count = task_spec_.ArgIdCount(i);
-    for (int j = 0; j < count; j++) {
-      dependencies_.push_back(task_spec_.ArgId(i, j));
-    }
-  }
-  if (task_spec_.IsActorTask()) {
-    dependencies_.push_back(task_spec_.PreviousActorTaskDummyObjectId());
-  }
-}
+void Task::ComputeDependencies() { dependencies_ = task_spec_.GetDependencies(); }
 
 void Task::CopyTaskExecutionSpec(const Task &task) {
   task_execution_spec_ = task.task_execution_spec_;

src/ray/common/task/task_spec.cc

@@ -112,6 +112,20 @@ const ResourceSet &TaskSpecification::GetRequiredResources() const {
   return *required_resources_;
 }
 
+std::vector<ObjectID> TaskSpecification::GetDependencies() const {
+  std::vector<ObjectID> dependencies;
+  for (size_t i = 0; i < NumArgs(); ++i) {
+    int count = ArgIdCount(i);
+    for (int j = 0; j < count; j++) {
+      dependencies.push_back(ArgId(i, j));
+    }
+  }
+  if (IsActorTask()) {
+    dependencies.push_back(PreviousActorTaskDummyObjectId());
+  }
+  return dependencies;
+}
+
 const ResourceSet &TaskSpecification::GetRequiredPlacementResources() const {
   return *required_placement_resources_;
 }

src/ray/common/task/task_spec.h

@@ -113,6 +113,12 @@ class TaskSpecification : public MessageWrapper<rpc::TaskSpec> {
   /// \return The resources that are required to place a task on a node.
   const ResourceSet &GetRequiredPlacementResources() const;
 
+  /// Return the dependencies of this task. This is recomputed each time, so it can
+  /// be used if the task spec is mutated.
+  ///
+  /// \return The recomputed dependencies for the task.
+  std::vector<ObjectID> GetDependencies() const;
+
   bool IsDriverTask() const;
 
   Language GetLanguage() const;

src/ray/core_worker/test/direct_task_transport_test.cc

@@ -1,6 +1,7 @@
 #include "gtest/gtest.h"
 
 #include "ray/common/task/task_spec.h"
+#include "ray/common/task/task_util.h"
 #include "ray/core_worker/store_provider/memory_store/memory_store.h"
 #include "ray/core_worker/transport/direct_task_transport.h"
 #include "ray/raylet/raylet_client.h"
@@ -216,6 +217,16 @@ TEST(LocalDependencyResolverTest, TestInlinePendingDependencies) {
   ASSERT_EQ(resolver.NumPendingTasks(), 0);
 }
 
+TaskSpecification BuildTaskSpec(const std::unordered_map<std::string, double> &resources,
+                                const std::vector<std::string> &function_descriptor) {
+  TaskSpecBuilder builder;
+  rpc::Address empty_address;
+  builder.SetCommonTaskSpec(TaskID::Nil(), Language::PYTHON, function_descriptor,
+                            JobID::Nil(), TaskID::Nil(), 0, TaskID::Nil(), empty_address,
+                            1, true, resources, resources);
+  return builder.Build();
+}
+
 TEST(DirectTaskTransportTest, TestSubmitOneTask) {
   auto raylet_client = std::make_shared<MockRayletClient>();
   auto worker_client = std::make_shared<MockWorkerClient>();
@@ -224,8 +235,10 @@ TEST(DirectTaskTransportTest, TestSubmitOneTask) {
   auto task_finisher = std::make_shared<MockTaskFinisher>();
   CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store,
                                           task_finisher, kLongTimeout);
-  TaskSpecification task;
-  task.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
+
+  std::unordered_map<std::string, double> empty_resources;
+  std::vector<std::string> empty_descriptor;
+  TaskSpecification task = BuildTaskSpec(empty_resources, empty_descriptor);
 
   ASSERT_TRUE(submitter.SubmitTask(task).ok());
   ASSERT_EQ(raylet_client->num_workers_requested, 1);
@@ -252,8 +265,9 @@ TEST(DirectTaskTransportTest, TestHandleTaskFailure) {
   auto task_finisher = std::make_shared<MockTaskFinisher>();
   CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store,
                                           task_finisher, kLongTimeout);
-  TaskSpecification task;
-  task.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
+  std::unordered_map<std::string, double> empty_resources;
+  std::vector<std::string> empty_descriptor;
+  TaskSpecification task = BuildTaskSpec(empty_resources, empty_descriptor);
 
   ASSERT_TRUE(submitter.SubmitTask(task).ok());
   ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1234, ClientID::Nil()));
@@ -274,12 +288,11 @@ TEST(DirectTaskTransportTest, TestConcurrentWorkerLeases) {
   auto task_finisher = std::make_shared<MockTaskFinisher>();
   CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store,
                                           task_finisher, kLongTimeout);
-  TaskSpecification task1;
-  TaskSpecification task2;
-  TaskSpecification task3;
-  task1.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
-  task2.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
-  task3.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
+  std::unordered_map<std::string, double> empty_resources;
+  std::vector<std::string> empty_descriptor;
+  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());
@@ -319,12 +332,11 @@ TEST(DirectTaskTransportTest, TestReuseWorkerLease) {
   auto task_finisher = std::make_shared<MockTaskFinisher>();
   CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store,
                                           task_finisher, kLongTimeout);
-  TaskSpecification task1;
-  TaskSpecification task2;
-  TaskSpecification task3;
-  task1.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
-  task2.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
-  task3.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
+  std::unordered_map<std::string, double> empty_resources;
+  std::vector<std::string> empty_descriptor;
+  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());
@@ -367,10 +379,10 @@ TEST(DirectTaskTransportTest, TestWorkerNotReusedOnError) {
   auto task_finisher = std::make_shared<MockTaskFinisher>();
   CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store,
                                           task_finisher, kLongTimeout);
-  TaskSpecification task1;
-  TaskSpecification task2;
-  task1.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
-  task2.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
+  std::unordered_map<std::string, double> empty_resources;
+  std::vector<std::string> empty_descriptor;
+  TaskSpecification task1 = BuildTaskSpec(empty_resources, empty_descriptor);
+  TaskSpecification task2 = BuildTaskSpec(empty_resources, empty_descriptor);
 
   ASSERT_TRUE(submitter.SubmitTask(task1).ok());
   ASSERT_TRUE(submitter.SubmitTask(task2).ok());
@@ -414,8 +426,9 @@ TEST(DirectTaskTransportTest, TestSpillback) {
   auto task_finisher = std::make_shared<MockTaskFinisher>();
   CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, lease_client_factory,
                                           store, task_finisher, kLongTimeout);
-  TaskSpecification task;
-  task.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
+  std::unordered_map<std::string, double> empty_resources;
+  std::vector<std::string> empty_descriptor;
+  TaskSpecification task = BuildTaskSpec(empty_resources, empty_descriptor);
 
   ASSERT_TRUE(submitter.SubmitTask(task).ok());
   ASSERT_EQ(raylet_client->num_workers_requested, 1);
@@ -443,6 +456,108 @@ TEST(DirectTaskTransportTest, TestSpillback) {
   ASSERT_EQ(task_finisher->num_tasks_failed, 0);
 }
 
+// Helper to run a test that checks that 'same1' and 'same2' are treated as the same
+// resource shape, while 'different' is treated as a separate shape.
+void TestSchedulingKey(const std::shared_ptr<CoreWorkerMemoryStore> store,
+                       const TaskSpecification &same1, const TaskSpecification &same2,
+                       const TaskSpecification &different) {
+  auto raylet_client = std::make_shared<MockRayletClient>();
+  auto worker_client = std::make_shared<MockWorkerClient>();
+  auto factory = [&](const rpc::WorkerAddress &addr) { return worker_client; };
+  auto task_finisher = std::make_shared<MockTaskFinisher>();
+  CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store,
+                                          task_finisher, kLongTimeout);
+
+  ASSERT_TRUE(submitter.SubmitTask(same1).ok());
+  ASSERT_TRUE(submitter.SubmitTask(same2).ok());
+  ASSERT_TRUE(submitter.SubmitTask(different).ok());
+  ASSERT_EQ(raylet_client->num_workers_requested, 2);
+
+  // same1 is pushed.
+  ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1000, ClientID::Nil()));
+  ASSERT_EQ(worker_client->callbacks.size(), 1);
+  // Another worker is requested because same2 is pending.
+  ASSERT_EQ(raylet_client->num_workers_requested, 3);
+
+  // 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.
+  ASSERT_EQ(worker_client->callbacks.size(), 1);
+
+  // different is pushed.
+  ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1001, ClientID::Nil()));
+  ASSERT_EQ(worker_client->callbacks.size(), 2);
+  ASSERT_EQ(raylet_client->num_workers_requested, 3);
+
+  // same2 runs successfully. Worker is returned.
+  ASSERT_TRUE(worker_client->ReplyPushTask());
+  ASSERT_EQ(raylet_client->num_workers_returned, 1);
+  ASSERT_EQ(raylet_client->num_workers_disconnected, 0);
+
+  // different runs successfully. Worker is returned.
+  ASSERT_TRUE(worker_client->ReplyPushTask());
+  ASSERT_EQ(raylet_client->num_workers_returned, 2);
+  ASSERT_EQ(raylet_client->num_workers_disconnected, 0);
+}
+
+TEST(DirectTaskTransportTest, TestSchedulingKeys) {
+  auto store = std::make_shared<CoreWorkerMemoryStore>();
+
+  std::unordered_map<std::string, double> resources1({{"a", 1.0}});
+  std::unordered_map<std::string, double> resources2({{"b", 2.0}});
+  std::vector<std::string> descriptor1({"a"});
+  std::vector<std::string> descriptor2({"b"});
+
+  // Tasks with different resources should request different worker leases.
+  RAY_LOG(INFO) << "Test different resources";
+  TestSchedulingKey(store, BuildTaskSpec(resources1, descriptor1),
+                    BuildTaskSpec(resources1, descriptor1),
+                    BuildTaskSpec(resources2, descriptor1));
+
+  // Tasks with different function descriptors should request different worker leases.
+  RAY_LOG(INFO) << "Test different descriptors";
+  TestSchedulingKey(store, BuildTaskSpec(resources1, descriptor1),
+                    BuildTaskSpec(resources1, descriptor1),
+                    BuildTaskSpec(resources1, descriptor2));
+
+  ObjectID direct1 = ObjectID::FromRandom().WithTransportType(TaskTransportType::DIRECT);
+  ObjectID direct2 = ObjectID::FromRandom().WithTransportType(TaskTransportType::DIRECT);
+  ObjectID plasma1 = ObjectID::FromRandom().WithTransportType(TaskTransportType::DIRECT);
+  ObjectID plasma2 = ObjectID::FromRandom().WithTransportType(TaskTransportType::DIRECT);
+  // Ensure the data is already present in the local store for direct call objects.
+  auto data = GenerateRandomObject();
+  ASSERT_TRUE(store->Put(*data, direct1).ok());
+  ASSERT_TRUE(store->Put(*data, direct2).ok());
+
+  // Force plasma objects to be promoted.
+  std::string meta = std::to_string(static_cast<int>(rpc::ErrorType::OBJECT_IN_PLASMA));
+  auto metadata = const_cast<uint8_t *>(reinterpret_cast<const uint8_t *>(meta.data()));
+  auto meta_buffer = std::make_shared<LocalMemoryBuffer>(metadata, meta.size());
+  auto plasma_data = RayObject(nullptr, meta_buffer);
+  ASSERT_TRUE(store->Put(plasma_data, plasma1).ok());
+  ASSERT_TRUE(store->Put(plasma_data, plasma2).ok());
+
+  TaskSpecification same_deps_1 = BuildTaskSpec(resources1, descriptor1);
+  same_deps_1.GetMutableMessage().add_args()->add_object_ids(direct1.Binary());
+  same_deps_1.GetMutableMessage().add_args()->add_object_ids(plasma1.Binary());
+  TaskSpecification same_deps_2 = BuildTaskSpec(resources1, descriptor1);
+  same_deps_2.GetMutableMessage().add_args()->add_object_ids(direct1.Binary());
+  same_deps_2.GetMutableMessage().add_args()->add_object_ids(direct2.Binary());
+  same_deps_2.GetMutableMessage().add_args()->add_object_ids(plasma1.Binary());
+
+  TaskSpecification different_deps = BuildTaskSpec(resources1, descriptor1);
+  different_deps.GetMutableMessage().add_args()->add_object_ids(direct1.Binary());
+  different_deps.GetMutableMessage().add_args()->add_object_ids(direct2.Binary());
+  different_deps.GetMutableMessage().add_args()->add_object_ids(plasma2.Binary());
+
+  // Tasks with different plasma dependencies should request different worker leases,
+  // but direct call dependencies shouldn't be considered.
+  RAY_LOG(INFO) << "Test different dependencies";
+  TestSchedulingKey(store, same_deps_1, same_deps_2, different_deps);
+}
+
 TEST(DirectTaskTransportTest, TestWorkerLeaseTimeout) {
   auto raylet_client = std::make_shared<MockRayletClient>();
   auto worker_client = std::make_shared<MockWorkerClient>();
@@ -452,12 +567,11 @@ TEST(DirectTaskTransportTest, TestWorkerLeaseTimeout) {
   CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store,
                                           task_finisher,
                                           /*lease_timeout_ms=*/5);
-  TaskSpecification task1;
-  TaskSpecification task2;
-  TaskSpecification task3;
-  task1.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
-  task2.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
-  task3.GetMutableMessage().set_task_id(TaskID::Nil().Binary());
+  std::unordered_map<std::string, double> empty_resources;
+  std::vector<std::string> empty_descriptor;
+  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());

src/ray/core_worker/transport/dependency_resolver.cc

@@ -3,6 +3,8 @@
 namespace ray {
 
 struct TaskState {
+  TaskState(TaskSpecification t, absl::flat_hash_set<ObjectID> deps)
+      : task(t), local_dependencies(deps) {}
   /// The task to be run.
   TaskSpecification task;
   /// The remaining dependencies to resolve for this task.
@@ -62,7 +64,7 @@ void LocalDependencyResolver::ResolveDependencies(TaskSpecification &task,
 
   // This is deleted when the last dependency fetch callback finishes.
   std::shared_ptr<TaskState> state =
-      std::shared_ptr<TaskState>(new TaskState{task, std::move(local_dependencies)});
+      std::make_shared<TaskState>(task, std::move(local_dependencies));
   num_pending_ += 1;
 
   for (const auto &obj_id : state->local_dependencies) {

src/ray/core_worker/transport/dependency_resolver.h

@@ -21,7 +21,8 @@ class LocalDependencyResolver {
   //
   /// Note: This method **will mutate** the given TaskSpecification.
   ///
-  /// Postcondition: all direct call ids in arguments are converted to values.
+  /// Postcondition: all direct call ids in arguments are converted to values and all
+  /// remaining by-reference arguments are TaskTransportType::RAYLET.
   void ResolveDependencies(TaskSpecification &task, std::function<void()> on_complete);
 
   /// Return the number of tasks pending dependency resolution.

src/ray/core_worker/transport/direct_task_transport.cc

@@ -6,53 +6,56 @@ namespace ray {
 
 Status CoreWorkerDirectTaskSubmitter::SubmitTask(TaskSpecification task_spec) {
   resolver_.ResolveDependencies(task_spec, [this, task_spec]() {
-    // TODO(ekl) should have a queue per distinct resource type required
     absl::MutexLock lock(&mu_);
-    queued_tasks_.push_back(task_spec);
-    RequestNewWorkerIfNeeded(task_spec);
+    // Note that the dependencies in the task spec are mutated to only contain
+    // plasma dependencies after ResolveDependencies finishes.
+    const SchedulingKey scheduling_key(task_spec.GetSchedulingClass(),
+                                       task_spec.GetDependencies());
+    auto it = task_queues_.find(scheduling_key);
+    if (it == task_queues_.end()) {
+      it = task_queues_.emplace(scheduling_key, std::deque<TaskSpecification>()).first;
+    }
+    it->second.push_back(task_spec);
+    RequestNewWorkerIfNeeded(scheduling_key);
   });
   return Status::OK();
 }
 
-void CoreWorkerDirectTaskSubmitter::HandleWorkerLeaseGranted(
+void CoreWorkerDirectTaskSubmitter::AddWorkerLeaseClient(
     const rpc::WorkerAddress &addr, std::shared_ptr<WorkerLeaseInterface> lease_client) {
-  // Setup client state for this worker.
-  {
-    absl::MutexLock lock(&mu_);
-    worker_request_pending_ = false;
-
-    auto it = client_cache_.find(addr);
-    if (it == client_cache_.end()) {
-      client_cache_[addr] =
-          std::shared_ptr<rpc::CoreWorkerClientInterface>(client_factory_(addr));
-      RAY_LOG(INFO) << "Connected to " << addr.first << ":" << addr.second;
-    }
-    int64_t expiration = current_time_ms() + lease_timeout_ms_;
-    worker_to_lease_client_.emplace(addr,
-                                    std::make_pair(std::move(lease_client), expiration));
+  auto it = client_cache_.find(addr);
+  if (it == client_cache_.end()) {
+    client_cache_[addr] =
+        std::shared_ptr<rpc::CoreWorkerClientInterface>(client_factory_(addr));
+    RAY_LOG(INFO) << "Connected to " << addr.first << ":" << addr.second;
   }
-
-  // Try to assign it work.
-  OnWorkerIdle(addr, /*error=*/false);
+  int64_t expiration = current_time_ms() + lease_timeout_ms_;
+  worker_to_lease_client_.emplace(addr,
+                                  std::make_pair(std::move(lease_client), expiration));
 }
 
 void CoreWorkerDirectTaskSubmitter::OnWorkerIdle(const rpc::WorkerAddress &addr,
+                                                 const SchedulingKey &scheduling_key,
                                                  bool was_error) {
-  absl::MutexLock lock(&mu_);
-  auto entry = worker_to_lease_client_[addr];
-  if (was_error || queued_tasks_.empty() || current_time_ms() > entry.second) {
-    RAY_CHECK_OK(entry.first->ReturnWorker(addr.second, was_error));
+  auto lease_entry = worker_to_lease_client_[addr];
+  auto queue_entry = task_queues_.find(scheduling_key);
+  // Return the worker if there was an error executing the previous 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.second) {
+    RAY_CHECK_OK(lease_entry.first->ReturnWorker(addr.second, was_error));
     worker_to_lease_client_.erase(addr);
-  } else if (!queued_tasks_.empty()) {
+  } else {
     auto &client = *client_cache_[addr];
-    PushNormalTask(addr, client, queued_tasks_.front());
-    queued_tasks_.pop_front();
-  }
-
-  // There are more tasks to run, so try to get another worker.
-  if (!queued_tasks_.empty()) {
-    RequestNewWorkerIfNeeded(queued_tasks_.front());
+    PushNormalTask(addr, client, scheduling_key, queue_entry->second.front());
+    queue_entry->second.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);
+    }
   }
+  RequestNewWorkerIfNeeded(scheduling_key);
 }
 
 std::shared_ptr<WorkerLeaseInterface>
@@ -78,67 +81,74 @@ CoreWorkerDirectTaskSubmitter::GetOrConnectLeaseClient(
 }
 
 void CoreWorkerDirectTaskSubmitter::RequestNewWorkerIfNeeded(
-    const TaskSpecification &resource_spec, const rpc::Address *raylet_address) {
-  if (worker_request_pending_) {
+    const SchedulingKey &scheduling_key, const rpc::Address *raylet_address) {
+  if (pending_lease_requests_.find(scheduling_key) != pending_lease_requests_.end()) {
+    // There's already an outstanding lease request for this type of task.
     return;
   }
-  if (queued_tasks_.empty()) {
-    // We don't have any tasks to run, so no need to request a worker.
+  auto it = task_queues_.find(scheduling_key);
+  if (it == task_queues_.end()) {
+    // We don't have any of this type of task to run.
     return;
   }
 
-  // NOTE(swang): We must copy the resource spec here because the resource spec
-  // may get swapped out by the time the callback fires. If we change this so
-  // that we associate the granted worker with the requested resource spec,
-  // then we can just pass the ref instead of copying.
-  TaskSpecification resource_spec_copy(resource_spec.GetMessage());
   auto lease_client = GetOrConnectLeaseClient(raylet_address);
+  TaskSpecification &resource_spec = it->second.front();
+  TaskID task_id = resource_spec.TaskId();
   RAY_CHECK_OK(lease_client->RequestWorkerLease(
-      resource_spec_copy,
-      [this, resource_spec_copy, lease_client](
+      resource_spec,
+      [this, lease_client, task_id, scheduling_key](
           const Status &status, const rpc::WorkerLeaseReply &reply) mutable {
+        absl::MutexLock lock(&mu_);
+        pending_lease_requests_.erase(scheduling_key);
         if (status.ok()) {
           if (!reply.worker_address().raylet_id().empty()) {
-            RAY_LOG(DEBUG) << "Lease granted " << resource_spec_copy.TaskId();
-            HandleWorkerLeaseGranted(
-                {reply.worker_address().ip_address(), reply.worker_address().port()},
-                std::move(lease_client));
+            // 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;
+            rpc::WorkerAddress addr(reply.worker_address().ip_address(),
+                                    reply.worker_address().port());
+            AddWorkerLeaseClient(addr, std::move(lease_client));
+            OnWorkerIdle(addr, scheduling_key, /*error=*/false);
           } else {
-            absl::MutexLock lock(&mu_);
-            worker_request_pending_ = false;
-            RequestNewWorkerIfNeeded(resource_spec_copy,
-                                     &reply.retry_at_raylet_address());
+            // The raylet redirected us to a different raylet to retry at.
+            RequestNewWorkerIfNeeded(scheduling_key, &reply.retry_at_raylet_address());
           }
         } else {
-          RAY_LOG(DEBUG) << "Retrying lease request " << resource_spec_copy.TaskId();
-          absl::MutexLock lock(&mu_);
-          worker_request_pending_ = false;
+          RAY_LOG(DEBUG) << "Retrying lease request " << task_id;
           if (lease_client != local_lease_client_) {
-            // A remote request failed. Retry the worker lease request locally
-            // if it's still in the queue.
+            // 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(resource_spec_copy);
+            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) << "Lost connection with local raylet. Error: "
                            << status.ToString();
           }
         }
       }));
-  worker_request_pending_ = true;
+  pending_lease_requests_.insert(scheduling_key);
 }
 
 void CoreWorkerDirectTaskSubmitter::PushNormalTask(const rpc::WorkerAddress &addr,
                                                    rpc::CoreWorkerClientInterface &client,
+                                                   const SchedulingKey &scheduling_key,
                                                    TaskSpecification &task_spec) {
   auto task_id = task_spec.TaskId();
   auto request = std::unique_ptr<rpc::PushTaskRequest>(new rpc::PushTaskRequest);
   request->mutable_task_spec()->Swap(&task_spec.GetMutableMessage());
   auto status = client.PushNormalTask(
-      std::move(request),
-      [this, task_id, addr](Status status, const rpc::PushTaskReply &reply) {
-        OnWorkerIdle(addr, /*error=*/!status.ok());
+      std::move(request), [this, task_id, scheduling_key, addr](
+                              Status status, const rpc::PushTaskReply &reply) {
+        {
+          absl::MutexLock lock(&mu_);
+          OnWorkerIdle(addr, scheduling_key, /*error=*/!status.ok());
+        }
         if (!status.ok()) {
           task_finisher_->FailPendingTask(task_id, rpc::ErrorType::WORKER_DIED);
         } else {

src/ray/core_worker/transport/direct_task_transport.h

@@ -19,6 +19,13 @@ namespace ray {
 typedef std::function<std::shared_ptr<WorkerLeaseInterface>(const rpc::Address &)>
     LeaseClientFactoryFn;
 
+// The task queues are keyed on resource shape & function descriptor
+// (encapsulated in SchedulingClass) to defer resource allocation decisions to the raylet
+// and ensure fairness between different tasks, as well as plasma task dependencies as
+// a performance optimization because the raylet will fetch plasma dependencies to the
+// scheduled worker.
+using SchedulingKey = std::pair<SchedulingClass, std::vector<ObjectID>>;
+
 // This class is thread-safe.
 class CoreWorkerDirectTaskSubmitter {
  public:
@@ -44,7 +51,8 @@ class CoreWorkerDirectTaskSubmitter {
   /// 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
   /// processing the worker, we don't attempt to re-use the worker.
-  void OnWorkerIdle(const rpc::WorkerAddress &addr, bool was_error);
+  void OnWorkerIdle(const rpc::WorkerAddress &addr, const SchedulingKey &task_queue_key,
+                    bool was_error) 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
@@ -56,20 +64,19 @@ class CoreWorkerDirectTaskSubmitter {
   /// flight and there are tasks queued. If a raylet address is provided, then
   /// the worker should be requested from the raylet at that address. Else, the
   /// worker should be requested from the local raylet.
-  void RequestNewWorkerIfNeeded(const TaskSpecification &resource_spec,
+  void RequestNewWorkerIfNeeded(const SchedulingKey &task_queue_key,
                                 const rpc::Address *raylet_address = nullptr)
       EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
-  /// Callback for when the raylet grants us a worker lease. The worker is returned
-  /// to the raylet via the given lease client once the task queue is empty.
-  /// TODO: Implement a lease term by which we need to return the worker.
-  void HandleWorkerLeaseGranted(const rpc::WorkerAddress &addr,
-                                std::shared_ptr<WorkerLeaseInterface> lease_client);
+  /// 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_);
 
   /// Push a task to a specific worker.
   void PushNormalTask(const rpc::WorkerAddress &addr,
                       rpc::CoreWorkerClientInterface &client,
-                      TaskSpecification &task_spec);
+                      const SchedulingKey &task_queue_key, TaskSpecification &task_spec);
 
   // Client that can be used to lease and return workers from the local raylet.
   std::shared_ptr<WorkerLeaseInterface> local_lease_client_;
@@ -107,11 +114,14 @@ class CoreWorkerDirectTaskSubmitter {
                       std::pair<std::shared_ptr<WorkerLeaseInterface>, int64_t>>
       worker_to_lease_client_ GUARDED_BY(mu_);
 
-  // Whether we have a request to the Raylet to acquire a new worker in flight.
-  bool worker_request_pending_ GUARDED_BY(mu_) = false;
+  // Keeps track of pending worker lease requests to the raylet.
+  absl::flat_hash_set<SchedulingKey> pending_lease_requests_ GUARDED_BY(mu_);
 
-  // Tasks that are queued for execution in this submitter..
-  std::deque<TaskSpecification> queued_tasks_ GUARDED_BY(mu_);
+  // 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_
+      GUARDED_BY(mu_);
 };
 
 };  // namespace ray

src/ray/protobuf/node_manager.proto

@@ -14,7 +14,7 @@ message SubmitTaskReply {
 
 // Request a worker from the raylet with the specified resources.
 message WorkerLeaseRequest {
-  // Task containing the requested resources.
+  // TaskSpec containing the requested resources.
   TaskSpec resource_spec = 1;
 }