diff --git a/src/ray/common/task/task.cc b/src/ray/common/task/task.cc index 61413d767..827408cf4 100644 --- a/src/ray/common/task/task.cc +++ b/src/ray/common/task/task.cc @@ -14,18 +14,7 @@ void Task::IncrementNumForwards() { task_execution_spec_.IncrementNumForwards(); const std::vector &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_; diff --git a/src/ray/common/task/task_spec.cc b/src/ray/common/task/task_spec.cc index 9f2fc0aa3..5d96dae98 100644 --- a/src/ray/common/task/task_spec.cc +++ b/src/ray/common/task/task_spec.cc @@ -112,6 +112,20 @@ const ResourceSet &TaskSpecification::GetRequiredResources() const { return *required_resources_; } +std::vector TaskSpecification::GetDependencies() const { + std::vector 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_; } diff --git a/src/ray/common/task/task_spec.h b/src/ray/common/task/task_spec.h index 51dc52d06..0c14912d0 100644 --- a/src/ray/common/task/task_spec.h +++ b/src/ray/common/task/task_spec.h @@ -113,6 +113,12 @@ class TaskSpecification : public MessageWrapper { /// \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 GetDependencies() const; + bool IsDriverTask() const; Language GetLanguage() const; 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 3f2a05aed..941ea2012 100644 --- a/src/ray/core_worker/test/direct_task_transport_test.cc +++ b/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 &resources, + const std::vector &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(); auto worker_client = std::make_shared(); @@ -224,8 +235,10 @@ TEST(DirectTaskTransportTest, TestSubmitOneTask) { auto task_finisher = std::make_shared(); CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store, task_finisher, kLongTimeout); - TaskSpecification task; - task.GetMutableMessage().set_task_id(TaskID::Nil().Binary()); + + std::unordered_map empty_resources; + std::vector 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(); CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store, task_finisher, kLongTimeout); - TaskSpecification task; - task.GetMutableMessage().set_task_id(TaskID::Nil().Binary()); + std::unordered_map empty_resources; + std::vector 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(); 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 empty_resources; + std::vector 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(); 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 empty_resources; + std::vector 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(); 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 empty_resources; + std::vector 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(); 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 empty_resources; + std::vector 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 store, + const TaskSpecification &same1, const TaskSpecification &same2, + const TaskSpecification &different) { + auto raylet_client = std::make_shared(); + auto worker_client = std::make_shared(); + auto factory = [&](const rpc::WorkerAddress &addr) { return worker_client; }; + auto task_finisher = std::make_shared(); + 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(); + + std::unordered_map resources1({{"a", 1.0}}); + std::unordered_map resources2({{"b", 2.0}}); + std::vector descriptor1({"a"}); + std::vector 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(rpc::ErrorType::OBJECT_IN_PLASMA)); + auto metadata = const_cast(reinterpret_cast(meta.data())); + auto meta_buffer = std::make_shared(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(); auto worker_client = std::make_shared(); @@ -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 empty_resources; + std::vector 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()); diff --git a/src/ray/core_worker/transport/dependency_resolver.cc b/src/ray/core_worker/transport/dependency_resolver.cc index 52930d503..6a0463143 100644 --- a/src/ray/core_worker/transport/dependency_resolver.cc +++ b/src/ray/core_worker/transport/dependency_resolver.cc @@ -3,6 +3,8 @@ namespace ray { struct TaskState { + TaskState(TaskSpecification t, absl::flat_hash_set 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 state = - std::shared_ptr(new TaskState{task, std::move(local_dependencies)}); + std::make_shared(task, std::move(local_dependencies)); num_pending_ += 1; for (const auto &obj_id : state->local_dependencies) { diff --git a/src/ray/core_worker/transport/dependency_resolver.h b/src/ray/core_worker/transport/dependency_resolver.h index ae65dddbd..a2ce057ae 100644 --- a/src/ray/core_worker/transport/dependency_resolver.h +++ b/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 on_complete); /// Return the number of tasks pending dependency resolution. diff --git a/src/ray/core_worker/transport/direct_task_transport.cc b/src/ray/core_worker/transport/direct_task_transport.cc index 7b092c8fd..f275f5403 100644 --- a/src/ray/core_worker/transport/direct_task_transport.cc +++ b/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()).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 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(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(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 @@ -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(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 { diff --git a/src/ray/core_worker/transport/direct_task_transport.h b/src/ray/core_worker/transport/direct_task_transport.h index 39e786abe..e269c26a4 100644 --- a/src/ray/core_worker/transport/direct_task_transport.h +++ b/src/ray/core_worker/transport/direct_task_transport.h @@ -19,6 +19,13 @@ namespace ray { typedef std::function(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>; + // 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 lease_client); + /// Set up client state for newly granted worker lease. + void AddWorkerLeaseClient(const rpc::WorkerAddress &addr, + std::shared_ptr 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 local_lease_client_; @@ -107,11 +114,14 @@ class CoreWorkerDirectTaskSubmitter { std::pair, 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 pending_lease_requests_ GUARDED_BY(mu_); - // Tasks that are queued for execution in this submitter.. - std::deque 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> task_queues_ + GUARDED_BY(mu_); }; }; // namespace ray diff --git a/src/ray/protobuf/node_manager.proto b/src/ray/protobuf/node_manager.proto index f3c429a6d..5661b5296 100644 --- a/src/ray/protobuf/node_manager.proto +++ b/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; }