diff --git a/.travis.yml b/.travis.yml index 6a0e49f1f..cb87bf07f 100644 --- a/.travis.yml +++ b/.travis.yml @@ -162,6 +162,7 @@ install: - ./src/ray/raylet/task_test - ./src/ray/raylet/worker_pool_test - ./src/ray/raylet/lineage_cache_test + - ./src/ray/raylet/task_dependency_manager_test - bash ../../../src/common/test/run_tests.sh - bash ../../../src/plasma/test/run_tests.sh diff --git a/src/ray/object_manager/object_directory.cc b/src/ray/object_manager/object_directory.cc index 990727e40..3a3256e7e 100644 --- a/src/ray/object_manager/object_directory.cc +++ b/src/ray/object_manager/object_directory.cc @@ -9,24 +9,34 @@ ObjectDirectory::ObjectDirectory(std::shared_ptr gcs_client ray::Status ObjectDirectory::ReportObjectAdded(const ObjectID &object_id, const ClientID &client_id, const ObjectInfoT &object_info) { - // TODO(hme): Determine whether we need to do lookup to append. - JobID job_id = JobID::from_random(); + // Append the addition entry to the object table. + JobID job_id = JobID::nil(); auto data = std::make_shared(); data->manager = client_id.binary(); data->is_eviction = false; + data->num_evictions = object_evictions_[object_id]; data->object_size = object_info.data_size; - ray::Status status = gcs_client_->object_table().Append( - job_id, object_id, data, - [](gcs::AsyncGcsClient *client, const UniqueID &id, const ObjectTableDataT &data) { - // Do nothing. - }); + ray::Status status = + gcs_client_->object_table().Append(job_id, object_id, data, nullptr); return status; }; ray::Status ObjectDirectory::ReportObjectRemoved(const ObjectID &object_id, const ClientID &client_id) { - // TODO(hme): Need corresponding remove method in GCS. - return ray::Status::NotImplemented("ObjectTable.Remove is not implemented"); + // Append the eviction entry to the object table. + JobID job_id = JobID::nil(); + auto data = std::make_shared(); + data->manager = client_id.binary(); + data->is_eviction = true; + data->num_evictions = object_evictions_[object_id]; + ray::Status status = + gcs_client_->object_table().Append(job_id, object_id, data, nullptr); + // Increment the number of times we've evicted this object. NOTE(swang): This + // is only necessary because the Ray redis module expects unique entries in a + // log. We track the number of evictions so that the next eviction, if there + // is one, is unique. + object_evictions_[object_id]++; + return status; }; ray::Status ObjectDirectory::GetInformation(const ClientID &client_id, diff --git a/src/ray/object_manager/object_directory.h b/src/ray/object_manager/object_directory.h index 7eca8c550..74fd4ab49 100644 --- a/src/ray/object_manager/object_directory.h +++ b/src/ray/object_manager/object_directory.h @@ -126,6 +126,9 @@ class ObjectDirectory : public ObjectDirectoryInterface { std::unordered_map existing_requests_; /// Reference to the gcs client. std::shared_ptr gcs_client_; + /// Map from object ID to the number of times it's been evicted on this + /// node before. + std::unordered_map object_evictions_; }; } // namespace ray diff --git a/src/ray/object_manager/object_manager.h b/src/ray/object_manager/object_manager.h index d34d50762..98663b52f 100644 --- a/src/ray/object_manager/object_manager.h +++ b/src/ray/object_manager/object_manager.h @@ -43,8 +43,15 @@ struct ObjectManagerConfig { std::string store_socket_name; }; +class ObjectManagerInterface { + public: + virtual ray::Status Pull(const ObjectID &object_id) = 0; + virtual ray::Status Cancel(const ObjectID &object_id) = 0; + virtual ~ObjectManagerInterface(){}; +}; + // TODO(hme): Add success/failure callbacks for push and pull. -class ObjectManager { +class ObjectManager : public ObjectManagerInterface { public: /// Implicitly instantiates Ray implementation of ObjectDirectory. /// diff --git a/src/ray/raylet/CMakeLists.txt b/src/ray/raylet/CMakeLists.txt index 37f61b7c0..8149fb35e 100644 --- a/src/ray/raylet/CMakeLists.txt +++ b/src/ray/raylet/CMakeLists.txt @@ -25,6 +25,7 @@ ADD_RAY_TEST(worker_pool_test STATIC_LINK_LIBS ray_static ${PLASMA_STATIC_LIB} $ ADD_RAY_TEST(task_test STATIC_LINK_LIBS ray_static gtest gtest_main gmock_main pthread ${Boost_SYSTEM_LIBRARY}) ADD_RAY_TEST(lineage_cache_test STATIC_LINK_LIBS ray_static gtest gtest_main gmock_main pthread ${Boost_SYSTEM_LIBRARY}) +ADD_RAY_TEST(task_dependency_manager_test STATIC_LINK_LIBS ray_static gtest gtest_main gmock_main pthread ${Boost_SYSTEM_LIBRARY}) add_library(rayletlib raylet.cc ${NODE_MANAGER_FBS_OUTPUT_FILES}) target_link_libraries(rayletlib ray_static ${Boost_SYSTEM_LIBRARY}) diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 3fc99a2a6..4214b80ad 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -74,10 +74,7 @@ NodeManager::NodeManager(boost::asio::io_service &io_service, local_queues_(SchedulingQueue()), scheduling_policy_(local_queues_), reconstruction_policy_([this](const TaskID &task_id) { ResubmitTask(task_id); }), - task_dependency_manager_( - object_manager, - // reconstruction_policy_, - [this](const TaskID &task_id) { HandleWaitingTaskReady(task_id); }), + task_dependency_manager_(object_manager), lineage_cache_(gcs_client_->client_table().GetLocalClientId(), gcs_client->raylet_task_table(), gcs_client->raylet_task_table()), remote_clients_(), @@ -88,6 +85,13 @@ NodeManager::NodeManager(boost::asio::io_service &io_service, ClientID local_client_id = gcs_client_->client_table().GetLocalClientId(); cluster_resource_map_.emplace(local_client_id, SchedulingResources(config.resource_config)); + + RAY_CHECK_OK(object_manager_.SubscribeObjAdded([this](const ObjectInfoT &object_info) { + ObjectID object_id = ObjectID::from_binary(object_info.object_id); + HandleObjectLocal(object_id); + })); + RAY_CHECK_OK(object_manager_.SubscribeObjDeleted( + [this](const ObjectID &object_id) { HandleObjectMissing(object_id); })); } ray::Status NodeManager::RegisterGcs() { @@ -374,6 +378,9 @@ void NodeManager::ProcessClientMessage( auto message = flatbuffers::GetRoot(message_data); ObjectID object_id = from_flatbuf(*message->object_id()); RAY_LOG(DEBUG) << "reconstructing object " << object_id; + // TODO(swang): Instead of calling Pull on the object directly, record the + // fact that the blocked task is dependent on this object_id in the task + // dependency manager. RAY_CHECK_OK(object_manager_.Pull(object_id)); // If the blocked client is a worker, and the worker isn't already blocked, @@ -468,13 +475,6 @@ void NodeManager::ProcessNodeManagerMessage(TcpClientConnection &node_manager_cl node_manager_client.ProcessMessages(); } -void NodeManager::HandleWaitingTaskReady(const TaskID &task_id) { - auto ready_tasks = local_queues_.RemoveTasks({task_id}); - local_queues_.QueueReadyTasks(std::vector(ready_tasks)); - // Schedule the newly ready tasks if possible. - ScheduleTasks(); -} - void NodeManager::ScheduleTasks() { // This method performs the transition of tasks from PENDING to SCHEDULED. auto policy_decision = scheduling_policy_.Schedule( @@ -500,23 +500,32 @@ void NodeManager::ScheduleTasks() { RAY_CHECK(1 == tasks.size()); Task &task = tasks.front(); // TODO(swang): Handle forward task failure. - // TODO(swang): Unsubscribe this task in the task dependency manager. RAY_CHECK_OK(ForwardTask(task, client_id)); } + // Notify the task dependency manager that we no longer need this task's + // object dependencies. + // NOTE(swang): For local tasks, the scheduled task's dependencies may get + // evicted before it can be assigned to a worker. + task_dependency_manager_.UnsubscribeDependencies(task_id); } // Transition locally scheduled tasks to SCHEDULED and dispatch scheduled tasks. - std::vector tasks = local_queues_.RemoveTasks(local_task_ids); - local_queues_.QueueScheduledTasks(tasks); - DispatchTasks(); + if (local_task_ids.size() > 0) { + std::vector tasks = local_queues_.RemoveTasks(local_task_ids); + local_queues_.QueueScheduledTasks(tasks); + DispatchTasks(); + } } void NodeManager::SubmitTask(const Task &task, const Lineage &uncommitted_lineage) { - const TaskSpecification &spec = task.GetTaskSpecification(); - // Add the task and its uncommitted lineage to the lineage cache. lineage_cache_.AddWaitingTask(task, uncommitted_lineage); + // Mark the task as pending. Once the task has finished execution, or once it + // has been forwarded to another node, the task must be marked as canceled in + // the TaskDependencyManager. + task_dependency_manager_.TaskPending(task); + const TaskSpecification &spec = task.GetTaskSpecification(); if (spec.IsActorTask()) { // Check whether we know the location of the actor. const auto actor_entry = actor_registry_.find(spec.ActorId()); @@ -561,14 +570,31 @@ void NodeManager::SubmitTask(const Task &task, const Lineage &uncommitted_lineag } } +void NodeManager::HandleRemoteDependencyRequired(const ObjectID &dependency_id) { + // Try to fetch the object from the object manager. + RAY_CHECK_OK(object_manager_.Pull(dependency_id)); + // TODO(swang): Request reconstruction of the object, possibly after a + // timeout. +} + +void NodeManager::HandleRemoteDependencyCanceled(const ObjectID &dependency_id) { + // Cancel the fetch request from the object manager. + RAY_CHECK_OK(object_manager_.Cancel(dependency_id)); + // TODO(swang): Cancel reconstruction of the object. +} + void NodeManager::QueueTask(const Task &task) { - // Queue the task depending on the availability of its arguments. - if (task_dependency_manager_.TaskReady(task)) { - local_queues_.QueueReadyTasks(std::vector({task})); + // Subscribe to the task's dependencies. + bool ready = task_dependency_manager_.SubscribeDependencies( + task.GetTaskSpecification().TaskId(), task.GetDependencies()); + // Queue the task. If all dependencies are available, then the task is queued + // in the READY state, else the WAITING. + if (ready) { + local_queues_.QueueReadyTasks({task}); + // Try to schedule the newly ready task. ScheduleTasks(); } else { - local_queues_.QueueWaitingTasks(std::vector({task})); - task_dependency_manager_.SubscribeTaskReady(task); + local_queues_.QueueWaitingTasks({task}); } } @@ -687,9 +713,12 @@ void NodeManager::FinishAssignedTask(Worker &worker) { if (task.GetTaskSpecification().IsActorCreationTask() || task.GetTaskSpecification().IsActorTask()) { auto dummy_object = task.GetTaskSpecification().ActorDummyObject(); - task_dependency_manager_.MarkDependencyReady(dummy_object); + HandleObjectLocal(dummy_object); } + // Notify the task dependency manager that this task has finished execution. + task_dependency_manager_.TaskCanceled(task_id); + // Unset the worker's assigned task. worker.AssignTaskId(TaskID::nil()); } @@ -698,6 +727,36 @@ void NodeManager::ResubmitTask(const TaskID &task_id) { throw std::runtime_error("Method not implemented"); } +void NodeManager::HandleObjectLocal(const ObjectID &object_id) { + // Notify the task dependency manager that this object is local. + const auto ready_task_ids = task_dependency_manager_.HandleObjectLocal(object_id); + // Transition the tasks whose dependencies are now fulfilled to the ready + // state. + if (ready_task_ids.size() > 0) { + std::unordered_set ready_task_id_set(ready_task_ids.begin(), + ready_task_ids.end()); + auto ready_tasks = local_queues_.RemoveTasks(ready_task_id_set); + local_queues_.QueueReadyTasks(std::vector(ready_tasks)); + // Schedule the newly ready tasks. + ScheduleTasks(); + } +} + +void NodeManager::HandleObjectMissing(const ObjectID &object_id) { + // Notify the task dependency manager that this object is no longer local. + const auto waiting_task_ids = task_dependency_manager_.HandleObjectMissing(object_id); + // Transition any tasks that were in the runnable state and are dependent on + // this object to the waiting state. + if (!waiting_task_ids.empty()) { + // Transition the tasks back to the waiting state. They will be made + // runnable once the deleted object becomes available again. + std::unordered_set waiting_task_id_set(waiting_task_ids.begin(), + waiting_task_ids.end()); + auto waiting_tasks = local_queues_.RemoveTasks(waiting_task_id_set); + local_queues_.QueueWaitingTasks(std::vector(waiting_tasks)); + } +} + ray::Status NodeManager::ForwardTask(const Task &task, const ClientID &node_id) { const auto &spec = task.GetTaskSpecification(); auto task_id = spec.TaskId(); @@ -733,7 +792,9 @@ ray::Status NodeManager::ForwardTask(const Task &task, const ClientID &node_id) // lineage cache since the receiving node is now responsible for writing // the task to the GCS. lineage_cache_.RemoveWaitingTask(task_id); - + // Notify the task dependency manager that we are no longer responsible + // for executing this task. + task_dependency_manager_.TaskCanceled(task_id); // Preemptively push any local arguments to the receiving node. For now, we // only do this with actor tasks, since actor tasks must be executed by a // specific process and therefore have affinity to the receiving node. diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index 268e61f23..fd7cbfdab 100644 --- a/src/ray/raylet/node_manager.h +++ b/src/ray/raylet/node_manager.h @@ -57,11 +57,16 @@ class NodeManager { ray::Status RegisterGcs(); private: - // Handler for the addition of a new GCS client. + /// Methods for handling clients. + /// Handler for the addition of a new GCS client. void ClientAdded(const ClientTableDataT &data); - // Handler for the creation of an actor, possibly on a remote node. - void HandleActorCreation(const ActorID &actor_id, - const std::vector &data); + /// Send heartbeats to the GCS. + void Heartbeat(); + /// Handler for a heartbeat notification from the GCS. + void HeartbeatAdded(gcs::AsyncGcsClient *client, const ClientID &id, + const HeartbeatTableDataT &data); + + /// Methods for task scheduling. // Queue a task for local execution. void QueueTask(const Task &task); /// Submit a task to this node. @@ -72,25 +77,35 @@ class NodeManager { void FinishAssignedTask(Worker &worker); /// Schedule tasks. void ScheduleTasks(); - /// Handle a task whose local dependencies were missing and are now available. - void HandleWaitingTaskReady(const TaskID &task_id); /// Resubmit a task whose return value needs to be reconstructed. void ResubmitTask(const TaskID &task_id); /// Forward a task to another node to execute. The task is assumed to not be /// queued in local_queues_. ray::Status ForwardTask(const Task &task, const ClientID &node_id); - /// Send heartbeats to the GCS. - void Heartbeat(); - /// Handler for a notification about a new client from the GCS. - void ClientAdded(gcs::AsyncGcsClient *client, const UniqueID &id, - const ClientTableDataT &data); - /// Handler for a heartbeat notification from the GCS. - void HeartbeatAdded(gcs::AsyncGcsClient *client, const ClientID &id, - const HeartbeatTableDataT &data); /// Dispatch locally scheduled tasks. This attempts the transition from "scheduled" to /// "running" task state. void DispatchTasks(); + /// Methods for actor scheduling. + /// Handler for the creation of an actor, possibly on a remote node. + void HandleActorCreation(const ActorID &actor_id, + const std::vector &data); + + /// Methods for managing object dependencies. + /// Handle a dependency required by a queued task that is missing locally. + /// The dependency is (1) on a remote node, (2) pending creation on a remote + /// node, or (3) missing from all nodes and requires reconstruction. + void HandleRemoteDependencyRequired(const ObjectID &dependency_id); + /// Handle a dependency that was previously required by a queued task that is + /// no longer required. + void HandleRemoteDependencyCanceled(const ObjectID &dependency_id); + /// Handle an object becoming local. This updates any local accounting, but + /// does not write to any global accounting in the GCS. + void HandleObjectLocal(const ObjectID &object_id); + /// Handle an object that is no longer local. This updates any local + /// accounting, but does not write to any global accounting in the GCS. + void HandleObjectMissing(const ObjectID &object_id); + boost::asio::io_service &io_service_; ObjectManager &object_manager_; /// A client connection to the GCS. diff --git a/src/ray/raylet/task_dependency_manager.cc b/src/ray/raylet/task_dependency_manager.cc index 19cf93308..7874f2313 100644 --- a/src/ray/raylet/task_dependency_manager.cc +++ b/src/ray/raylet/task_dependency_manager.cc @@ -4,114 +4,225 @@ namespace ray { namespace raylet { -TaskDependencyManager::TaskDependencyManager( - ObjectManager &object_manager, - // ReconstructionPolicy &reconstruction_policy, - std::function handler) - : object_manager_(object_manager), - // reconstruction_policy_(reconstruction_policy), - task_ready_callback_(handler) { - // TODO(swang): Check return status. - ray::Status status = - object_manager_.SubscribeObjAdded([this](const ObjectInfoT &object_info) { - handleObjectReady(ObjectID::from_binary(object_info.object_id)); - }); - // TODO(swang): Subscribe to object removed notifications. -} +TaskDependencyManager::TaskDependencyManager(ObjectManagerInterface &object_manager) + : object_manager_(object_manager) {} bool TaskDependencyManager::CheckObjectLocal(const ObjectID &object_id) const { return local_objects_.count(object_id) == 1; } -bool TaskDependencyManager::argumentsReady(const std::vector arguments) const { - for (auto &argument : arguments) { - // Check if any argument is missing. - if (local_objects_.count(argument) == 0) { - return false; - } +bool TaskDependencyManager::CheckObjectRequired(const ObjectID &object_id) const { + const TaskID task_id = ComputeTaskId(object_id); + auto task_entry = required_tasks_.find(task_id); + // If there are no subscribed tasks that are dependent on the object, then do + // nothing. + if (task_entry == required_tasks_.end()) { + return false; + } + if (task_entry->second.count(object_id) == 0) { + return false; + } + // If the object is already local, then the dependency is fulfilled. Do + // nothing. + if (local_objects_.count(object_id) == 1) { + return false; + } + // If the task that creates the object is pending execution, then the + // dependency will be fulfilled locally. Do nothing. + if (pending_tasks_.count(task_id) == 1) { + return false; } - // All arguments are ready. return true; } -void TaskDependencyManager::handleObjectReady(const ray::ObjectID &object_id) { - RAY_LOG(DEBUG) << "object ready " << object_id; - // Add the object to the table of locally available objects. - RAY_CHECK(local_objects_.count(object_id) == 0); - local_objects_.insert(object_id); +void TaskDependencyManager::HandleRemoteDependencyRequired(const ObjectID &object_id) { + bool required = CheckObjectRequired(object_id); + // If the object is required, then try to make the object available locally. + if (required) { + auto inserted = required_objects_.insert(object_id); + if (inserted.second) { + // If we haven't already, request the object manager to pull it from a + // remote node. + RAY_CHECK_OK(object_manager_.Pull(object_id)); + } + } +} - // Handle any tasks that were dependent on the newly available object. +void TaskDependencyManager::HandleRemoteDependencyCanceled(const ObjectID &object_id) { + bool required = CheckObjectRequired(object_id); + // If the object is no longer required, then cancel the object. + if (!required) { + auto it = required_objects_.find(object_id); + if (it != required_objects_.end()) { + RAY_CHECK_OK(object_manager_.Cancel(object_id)); + required_objects_.erase(it); + } + } +} + +std::vector TaskDependencyManager::HandleObjectLocal( + const ray::ObjectID &object_id) { + RAY_LOG(DEBUG) << "object ready " << object_id.hex(); + // Add the object to the table of locally available objects. + auto inserted = local_objects_.insert(object_id); + RAY_CHECK(inserted.second); + + // Find any tasks that are dependent on the newly available object. std::vector ready_task_ids; - auto dependent_tasks = remote_object_dependencies_.find(object_id); - if (dependent_tasks != remote_object_dependencies_.end()) { - for (auto &dependent_task_id : dependent_tasks->second) { - // If the dependent task now has all of its arguments ready, it's ready - // to run. - if (argumentsReady(task_dependencies_[dependent_task_id])) { - ready_task_ids.push_back(dependent_task_id); + auto creating_task_entry = required_tasks_.find(ComputeTaskId(object_id)); + if (creating_task_entry != required_tasks_.end()) { + auto object_entry = creating_task_entry->second.find(object_id); + if (object_entry != creating_task_entry->second.end()) { + for (auto &dependent_task_id : object_entry->second) { + auto &task_entry = task_dependencies_[dependent_task_id]; + task_entry.num_missing_dependencies--; + // If the dependent task now has all of its arguments ready, it's ready + // to run. + if (task_entry.num_missing_dependencies == 0) { + ready_task_ids.push_back(dependent_task_id); + } } } - remote_object_dependencies_.erase(dependent_tasks); } + + // The object is now local, so cancel any in-progress operations to make the + // object local. + HandleRemoteDependencyCanceled(object_id); + + return ready_task_ids; +} + +std::vector TaskDependencyManager::HandleObjectMissing( + const ray::ObjectID &object_id) { + // Add the object to the table of locally available objects. + auto erased = local_objects_.erase(object_id); + RAY_CHECK(erased == 1); + + // Find any tasks that are dependent on the missing object. + std::vector waiting_task_ids; + TaskID creating_task_id = ComputeTaskId(object_id); + auto creating_task_entry = required_tasks_.find(creating_task_id); + if (creating_task_entry != required_tasks_.end()) { + auto object_entry = creating_task_entry->second.find(object_id); + if (object_entry != creating_task_entry->second.end()) { + for (auto &dependent_task_id : object_entry->second) { + auto &task_entry = task_dependencies_[dependent_task_id]; + // If the dependent task had all of its arguments ready, it was ready to + // run but must be switched to waiting since one of its arguments is now + // missing. + if (task_entry.num_missing_dependencies == 0) { + waiting_task_ids.push_back(dependent_task_id); + } + task_entry.num_missing_dependencies++; + } + } + } + // The object is no longer local. Try to make the object local if necessary. + HandleRemoteDependencyRequired(object_id); // Process callbacks for all of the tasks dependent on the object that are // now ready to run. - for (auto &ready_task_id : ready_task_ids) { - UnsubscribeTaskReady(ready_task_id); - task_ready_callback_(ready_task_id); - } + return waiting_task_ids; } -bool TaskDependencyManager::TaskReady(const Task &task) const { - const std::vector arguments = task.GetDependencies(); - return argumentsReady(arguments); -} +bool TaskDependencyManager::SubscribeDependencies( + const TaskID &task_id, const std::vector &required_objects) { + auto &task_entry = task_dependencies_[task_id]; -void TaskDependencyManager::SubscribeTaskReady(const Task &task) { - // TODO(swang): Don't pull arguments that are going to be created by a queued - // or running task. - TaskID task_id = task.GetTaskSpecification().TaskId(); - const std::vector arguments = task.GetDependencies(); - // Add the task's arguments to the table of subscribed tasks. - task_dependencies_[task_id] = arguments; - // Add the task's remote arguments to the table of remote objects. - int num_missing_arguments = 0; - for (auto &argument : arguments) { - if (local_objects_.count(argument) == 0) { - remote_object_dependencies_[argument].push_back(task_id); - num_missing_arguments++; - // TODO(swang): Check return status. - // TODO(swang): Handle Pull failure (if object manager does not retry). - // TODO(atumanov): pull return status should be propagated back to the caller. - ray::Status status = object_manager_.Pull(argument); + // Record the task's dependencies. + for (const auto &object_id : required_objects) { + auto inserted = task_entry.object_dependencies.insert(object_id); + if (inserted.second) { + // Get the ID of the task that creates the dependency. + TaskID creating_task_id = ComputeTaskId(object_id); + // Determine whether the dependency can be fulfilled by the local node. + if (local_objects_.count(object_id) == 0) { + // The object is not local. + task_entry.num_missing_dependencies++; + } + // Add the subscribed task to the mapping from object ID to list of + // dependent tasks. + required_tasks_[creating_task_id][object_id].push_back(task_id); } } - // Check that the task has some missing arguments. - RAY_CHECK(num_missing_arguments > 0); + + // These dependencies are required by the given task. Try to make them local + // if necessary. + for (const auto &object_id : required_objects) { + HandleRemoteDependencyRequired(object_id); + } + + // Return whether all dependencies are local. + return (task_entry.num_missing_dependencies == 0); } -void TaskDependencyManager::UnsubscribeTaskReady(const TaskID &task_id) { - const std::vector arguments = task_dependencies_[task_id]; +void TaskDependencyManager::UnsubscribeDependencies(const TaskID &task_id) { // Remove the task from the table of subscribed tasks. - task_dependencies_.erase(task_id); - // Remove the task from the table of remote objects to dependent tasks. - for (auto &argument : arguments) { - if (local_objects_.count(argument) == 1) { - continue; - } - // The argument is not local. Remove the task from the list of tasks that - // are dependent on this object. - std::vector &dependent_tasks = remote_object_dependencies_[task_id]; - for (auto it = dependent_tasks.begin(); it != dependent_tasks.end(); it++) { - if (*it == task_id) { - it = dependent_tasks.erase(it); - break; + auto it = task_dependencies_.find(task_id); + RAY_CHECK(it != task_dependencies_.end()); + const TaskDependencies task_entry = std::move(it->second); + task_dependencies_.erase(it); + + // Remove the task's dependencies. + for (const auto &object_id : task_entry.object_dependencies) { + // Remove the task from the list of tasks that are dependent on this + // object. + // Get the ID of the task that creates the dependency. + TaskID creating_task_id = ComputeTaskId(object_id); + auto creating_task_entry = required_tasks_.find(creating_task_id); + std::vector &dependent_tasks = creating_task_entry->second[object_id]; + auto it = std::find(dependent_tasks.begin(), dependent_tasks.end(), task_id); + RAY_CHECK(it != dependent_tasks.end()); + dependent_tasks.erase(it); + // If the unsubscribed task was the only task dependent on the object, then + // erase the object entry. + if (dependent_tasks.empty()) { + creating_task_entry->second.erase(object_id); + // Remove the task that creates this object if there are no more object + // dependencies created by the task. + if (creating_task_entry->second.empty()) { + required_tasks_.erase(creating_task_entry); } } } + + // These dependencies are no longer required by the given task. Cancel any + // in-progress operations to make them local. + for (const auto &object_id : task_entry.object_dependencies) { + HandleRemoteDependencyCanceled(object_id); + } } -void TaskDependencyManager::MarkDependencyReady(const ObjectID &object) { - handleObjectReady(object); +void TaskDependencyManager::TaskPending(const Task &task) { + TaskID task_id = task.GetTaskSpecification().TaskId(); + + // Record that the task is pending execution. + pending_tasks_.insert(task_id); + // Find any subscribed tasks that are dependent on objects created by the + // pending task. + auto remote_task_entry = required_tasks_.find(task_id); + if (remote_task_entry != required_tasks_.end()) { + for (const auto &object_entry : remote_task_entry->second) { + // This object created by the pending task will appear locally once the + // task completes execution. Cancel any in-progress operations to make + // the object local. + HandleRemoteDependencyCanceled(object_entry.first); + } + } +} + +void TaskDependencyManager::TaskCanceled(const TaskID &task_id) { + // Record that the task is no longer pending execution. + pending_tasks_.erase(task_id); + // Find any subscribed tasks that are dependent on objects created by the + // canceled task. + auto remote_task_entry = required_tasks_.find(task_id); + if (remote_task_entry != required_tasks_.end()) { + for (const auto &object_entry : remote_task_entry->second) { + // This object created by the task will no longer appear locally since + // the task is canceled. Try to make the object local if necessary. + HandleRemoteDependencyRequired(object_entry.first); + } + } } } // namespace raylet diff --git a/src/ray/raylet/task_dependency_manager.h b/src/ray/raylet/task_dependency_manager.h index b336ed4db..1ed7173ef 100644 --- a/src/ray/raylet/task_dependency_manager.h +++ b/src/ray/raylet/task_dependency_manager.h @@ -16,20 +16,18 @@ class ReconstructionPolicy; /// \class TaskDependencyManager /// -/// Responsible for managing task dependencies. Tasks that have object -/// dependencies that are missing locally should contact the manager to try to -/// make their dependencies available. +/// Responsible for managing object dependencies for tasks. The caller can +/// subscribe to object dependencies for a task. The task manager will +/// determine which object dependencies are remote. These are the objects that +/// are neither in the local object store, nor will they be created by a +/// locally queued task. The task manager will request that these objects be +/// made available locally, either by object transfer from a remote node or +/// reconstruction. The task manager will also cancel these objects if they are +/// no longer needed by any task. class TaskDependencyManager { public: /// Create a task dependency manager. - /// - /// \param object_manager A reference to the object manager so that the task - /// dependency manager can issue requests to transfer objects. - /// \param handler The handler to call for subscribed tasks whose - /// dependencies have become available locally. - TaskDependencyManager(ObjectManager &object_manager, - // ReconstructionPolicy &reconstruction_policy, - std::function handler); + TaskDependencyManager(ObjectManagerInterface &object_manager); /// Check whether an object is locally available. /// @@ -37,50 +35,109 @@ class TaskDependencyManager { /// \return Whether the object is local. bool CheckObjectLocal(const ObjectID &object_id) const; - /// Check whether a task's object dependencies are locally available. + /// Subscribe to object depedencies required by the task and check whether + /// all dependencies are fulfilled. This will track this task's dependencies + /// until UnsubscribeDependencies is called on the same task ID. If any + /// dependencies are remote, then they will be requested. When the last + /// remote dependency later appears locally via a call to HandleObjectLocal, + /// the subscribed task will be returned by the HandleObjectLocal call, + /// signifying that it is ready to run. This method may be called multiple + /// times per task. /// - /// \param task The task whose object dependencies will be checked. - /// \return Whether the task's object dependencies are ready. - bool TaskReady(const Task &task) const; + /// \param task_id The ID of the task whose dependencies to subscribe to. + /// \param required_objects The objects required by the task. + /// \return Whether all of the given dependencies for the given task are + /// local. + bool SubscribeDependencies(const TaskID &task_id, + const std::vector &required_objects); - /// Subscribe to a task that has missing dependencies. The manager will - /// attempt to make any missing dependencies available locally by transfer or - /// by reconstruction. The registered handler will be called when the task's - /// dependencies become locally available. + /// Unsubscribe from the object dependencies required by this task. If the + /// objects were remote and are no longer required by any subscribed task, + /// then they will be canceled. /// - /// \param task The task with missing dependencies. - void SubscribeTaskReady(const Task &task); + /// \param task_id The ID of the task whose dependencies to unsubscribe from. + void UnsubscribeDependencies(const TaskID &task_id); - /// Stop waiting for a task's dependencies to become available. + /// Mark that the given task is pending execution. Any objects that it creates + /// are now considered to be pending creation. If there are any subscribed + /// tasks that depend on these objects, then the objects will be canceled. /// - /// \param task_id The task ID of the task with missing dependencies. - void UnsubscribeTaskReady(const TaskID &task_id); + /// \param task The task that is pending execution. + void TaskPending(const Task &task); - /// Mark an object as locally available. This is used for objects that do not - /// have a stored value (e.g., actor execution dependencies). + /// Mark that the given task is no longer pending execution. Any objects that + /// it creates that are not already local are now considered to be remote. If + /// there are any subscribed tasks that depend on these objects, then the + /// objects will be requested. + /// + /// \param task_id The ID of the task to cancel. + void TaskCanceled(const TaskID &task_id); + + /// Handle an object becoming locally available. If there are any subscribed + /// tasks that depend on this object, then the object will be canceled. /// /// \param object_id The object ID of the object to mark as locally /// available. - void MarkDependencyReady(const ObjectID &object_id); + /// \return A list of task IDs. This contains all subscribed tasks that now + /// have all of their dependencies fulfilled, once this object was made + /// local. + std::vector HandleObjectLocal(const ray::ObjectID &object_id); + + /// Handle an object that is no longer locally available. If there are any + /// subscribed tasks that depend on this object, then the object will be + /// requested. + /// + /// \param object_id The object ID of the object that was previously locally + /// available. + /// \return A list of task IDs. This contains all subscribed tasks that + /// previously had all of their dependencies fulfilled, but are now missing + /// this object dependency. + std::vector HandleObjectMissing(const ray::ObjectID &object_id); private: - /// Check whether the given list of objects are ready. - bool argumentsReady(const std::vector arguments) const; - /// Handle an object added to the object store. - void handleObjectReady(const ray::ObjectID &object_id); - /// A reference to the object manager so that we can issue Pull requests of - /// missing objects. - ObjectManager &object_manager_; - /// A mapping from task ID of each subscribed task to its list of + using ObjectDependencyMap = std::unordered_map>; + + /// A struct to represent the object dependencies of a task. + struct TaskDependencies { + /// The objects that the task is dependent on. These must be local before + /// the task is ready to execute. + std::unordered_set object_dependencies; + /// The number of object arguments that are not available locally. This + /// must be zero before the task is ready to execute. + int64_t num_missing_dependencies; + }; + + /// Check whether the given object needs to be made available through object + /// transfer or reconstruction. These are objects for which: (1) there is a + /// subscribed task dependent on it, (2) the object is not local, and (3) the + /// task that creates the object is not pending execution locally. + bool CheckObjectRequired(const ObjectID &object_id) const; + /// If the given object is required, then request that the object be made + /// available through object transfer or reconstruction. + void HandleRemoteDependencyRequired(const ObjectID &object_id); + /// If the given object is no longer required, then cancel any in-progress + /// operations to make the object available through object transfer or + /// reconstruction. + void HandleRemoteDependencyCanceled(const ObjectID &object_id); + + ObjectManagerInterface &object_manager_; + /// A mapping from task ID of each subscribed task to its list of object /// dependencies. - std::unordered_map> task_dependencies_; - // A mapping from object ID of each object that is not locally available to - // the list of subscribed tasks that are dependent on it. - std::unordered_map> remote_object_dependencies_; - // The set of locally available objects. + std::unordered_map task_dependencies_; + /// All tasks whose outputs are required by a subscribed task. This is a + /// mapping from task ID to information about the objects that the task + /// creates, either by return value or by `ray.put`. For each object, we + /// store the IDs of the subscribed tasks that are dependent on the object. + std::unordered_map required_tasks_; + /// Objects that are required by a subscribed task, are not local, and are + /// not created by a pending task. For these objects, there are pending + /// operations to make the object available. + std::unordered_set required_objects_; + /// The set of locally available objects. std::unordered_set local_objects_; - // The callback to call when a subscribed task becomes ready. - std::function task_ready_callback_; + /// The set of tasks that are pending execution. Any objects created by these + /// tasks that are not already local are pending creation. + std::unordered_set pending_tasks_; }; } // namespace raylet diff --git a/src/ray/raylet/task_dependency_manager_test.cc b/src/ray/raylet/task_dependency_manager_test.cc new file mode 100644 index 000000000..936bfb749 --- /dev/null +++ b/src/ray/raylet/task_dependency_manager_test.cc @@ -0,0 +1,347 @@ +#include + +#include "gmock/gmock.h" +#include "gtest/gtest.h" + +#include + +#include "ray/raylet/task_dependency_manager.h" + +namespace ray { + +namespace raylet { + +using ::testing::_; + +class MockObjectManager : public ObjectManagerInterface { + public: + MOCK_METHOD1(Pull, ray::Status(const ObjectID &object_id)); + MOCK_METHOD1(Cancel, ray::Status(const ObjectID &object_id)); +}; + +class TaskDependencyManagerTest : public ::testing::Test { + public: + TaskDependencyManagerTest() + : object_manager_mock_(), task_dependency_manager_(object_manager_mock_) {} + + protected: + MockObjectManager object_manager_mock_; + TaskDependencyManager task_dependency_manager_; +}; + +static inline Task ExampleTask(const std::vector &arguments, + int64_t num_returns) { + std::unordered_map required_resources; + std::vector> task_arguments; + for (auto &argument : arguments) { + std::vector references = {argument}; + task_arguments.emplace_back(std::make_shared(references)); + } + auto spec = TaskSpecification(UniqueID::nil(), UniqueID::from_random(), 0, + UniqueID::from_random(), task_arguments, num_returns, + required_resources); + auto execution_spec = TaskExecutionSpecification(std::vector()); + execution_spec.IncrementNumForwards(); + Task task = Task(execution_spec, spec); + return task; +} + +std::vector MakeTaskChain(int chain_size, + const std::vector &initial_arguments, + int64_t num_returns) { + std::vector task_chain; + std::vector arguments = initial_arguments; + for (int i = 0; i < chain_size; i++) { + auto task = ExampleTask(arguments, num_returns); + task_chain.push_back(task); + arguments.clear(); + for (int j = 0; j < task.GetTaskSpecification().NumReturns(); j++) { + arguments.push_back(task.GetTaskSpecification().ReturnId(j)); + } + } + return task_chain; +} + +TEST_F(TaskDependencyManagerTest, TestSimpleTask) { + // Create a task with 3 arguments. + int num_arguments = 3; + std::vector arguments; + for (int i = 0; i < num_arguments; i++) { + arguments.push_back(ObjectID::from_random()); + } + TaskID task_id = TaskID::from_random(); + // No objects have been registered in the task dependency manager, so all + // arguments should be remote. + for (const auto &argument_id : arguments) { + EXPECT_CALL(object_manager_mock_, Pull(argument_id)); + } + // Subscribe to the task's dependencies. + bool ready = task_dependency_manager_.SubscribeDependencies(task_id, arguments); + ASSERT_FALSE(ready); + + // All arguments should be canceled as they become available locally. + for (const auto &argument_id : arguments) { + EXPECT_CALL(object_manager_mock_, Cancel(argument_id)); + } + // For each argument except the last, tell the task dependency manager that + // the argument is local. + int i = 0; + for (; i < num_arguments - 1; i++) { + auto ready_task_ids = task_dependency_manager_.HandleObjectLocal(arguments[i]); + ASSERT_TRUE(ready_task_ids.empty()); + } + // Tell the task dependency manager that the last argument is local. Now the + // task should be ready to run. + auto ready_task_ids = task_dependency_manager_.HandleObjectLocal(arguments[i]); + ASSERT_EQ(ready_task_ids.size(), 1); + ASSERT_EQ(ready_task_ids.front(), task_id); +} + +TEST_F(TaskDependencyManagerTest, TestDuplicateSubscribe) { + // Create a task with 3 arguments. + TaskID task_id = TaskID::from_random(); + int num_arguments = 3; + std::vector arguments; + for (int i = 0; i < num_arguments; i++) { + // Add the new argument to the list of dependencies to subscribe to. + ObjectID argument_id = ObjectID::from_random(); + arguments.push_back(argument_id); + // Subscribe to the task's dependencies. All arguments except the last are + // duplicates of previous subscription calls. Each argument should only be + // requested from the node manager once. + EXPECT_CALL(object_manager_mock_, Pull(argument_id)); + bool ready = task_dependency_manager_.SubscribeDependencies(task_id, arguments); + ASSERT_FALSE(ready); + } + + // All arguments should be canceled as they become available locally. + for (const auto &argument_id : arguments) { + EXPECT_CALL(object_manager_mock_, Cancel(argument_id)); + } + // For each argument except the last, tell the task dependency manager that + // the argument is local. + int i = 0; + for (; i < num_arguments - 1; i++) { + auto ready_task_ids = task_dependency_manager_.HandleObjectLocal(arguments[i]); + ASSERT_TRUE(ready_task_ids.empty()); + } + // Tell the task dependency manager that the last argument is local. Now the + // task should be ready to run. + auto ready_task_ids = task_dependency_manager_.HandleObjectLocal(arguments[i]); + ASSERT_EQ(ready_task_ids.size(), 1); + ASSERT_EQ(ready_task_ids.front(), task_id); +} + +TEST_F(TaskDependencyManagerTest, TestMultipleTasks) { + // Create 3 tasks that are dependent on the same object. + ObjectID argument_id = ObjectID::from_random(); + std::vector dependent_tasks; + int num_dependent_tasks = 3; + // The object should only be requested from the object manager once for all + // three tasks. + EXPECT_CALL(object_manager_mock_, Pull(argument_id)); + for (int i = 0; i < num_dependent_tasks; i++) { + TaskID task_id = TaskID::from_random(); + dependent_tasks.push_back(task_id); + // Subscribe to each of the task's dependencies. + bool ready = task_dependency_manager_.SubscribeDependencies(task_id, {argument_id}); + ASSERT_FALSE(ready); + } + + // Tell the task dependency manager that the object is local. + EXPECT_CALL(object_manager_mock_, Cancel(argument_id)); + auto ready_task_ids = task_dependency_manager_.HandleObjectLocal(argument_id); + // Check that all tasks are now ready to run. + ASSERT_EQ(ready_task_ids.size(), dependent_tasks.size()); + for (const auto &task_id : ready_task_ids) { + ASSERT_NE(std::find(dependent_tasks.begin(), dependent_tasks.end(), task_id), + dependent_tasks.end()); + } +} + +TEST_F(TaskDependencyManagerTest, TestTaskChain) { + // Create 3 tasks, each dependent on the previous. The first task has no + // arguments. + int num_tasks = 3; + auto tasks = MakeTaskChain(num_tasks, {}, 1); + int num_ready_tasks = 1; + int i = 0; + // No objects should be remote or canceled since each task depends on a + // locally queued task. + EXPECT_CALL(object_manager_mock_, Pull(_)).Times(0); + EXPECT_CALL(object_manager_mock_, Cancel(_)).Times(0); + for (const auto &task : tasks) { + // Subscribe to each of the tasks' arguments. + auto arguments = task.GetDependencies(); + bool ready = task_dependency_manager_.SubscribeDependencies( + task.GetTaskSpecification().TaskId(), arguments); + if (i < num_ready_tasks) { + // The first task should be ready to run since it has no arguments. + ASSERT_TRUE(ready); + } else { + // All remaining tasks depend on the previous task. + ASSERT_FALSE(ready); + } + + // Mark each task as pending. + task_dependency_manager_.TaskPending(task); + + i++; + } + + // Simulate executing each task. Each task's completion should make the next + // task runnable. + while (!tasks.empty()) { + auto task = tasks.front(); + tasks.erase(tasks.begin()); + TaskID task_id = task.GetTaskSpecification().TaskId(); + auto return_id = task.GetTaskSpecification().ReturnId(0); + + task_dependency_manager_.UnsubscribeDependencies(task_id); + // Simulate the object notifications for the task's return values. + auto ready_tasks = task_dependency_manager_.HandleObjectLocal(return_id); + if (tasks.empty()) { + // If there are no more tasks, then there should be no more tasks that + // become ready to run. + ASSERT_TRUE(ready_tasks.empty()); + } else { + // If there are more tasks to run, then the next task in the chain should + // now be ready to run. + ASSERT_EQ(ready_tasks.size(), 1); + ASSERT_EQ(ready_tasks.front(), tasks.front().GetTaskSpecification().TaskId()); + } + // Simulate the task finishing execution. + task_dependency_manager_.TaskCanceled(task_id); + } +} + +TEST_F(TaskDependencyManagerTest, TestDependentPut) { + // Create a task with 3 arguments. + auto task1 = ExampleTask({}, 0); + ObjectID put_id = ComputePutId(task1.GetTaskSpecification().TaskId(), 1); + auto task2 = ExampleTask({put_id}, 0); + + // No objects have been registered in the task dependency manager, so the put + // object should be remote. + EXPECT_CALL(object_manager_mock_, Pull(put_id)); + // Subscribe to the task's dependencies. + bool ready = task_dependency_manager_.SubscribeDependencies( + task2.GetTaskSpecification().TaskId(), {put_id}); + ASSERT_FALSE(ready); + + // The put object should be considered local as soon as the task that creates + // it is pending execution. + EXPECT_CALL(object_manager_mock_, Cancel(put_id)); + task_dependency_manager_.TaskPending(task1); +} + +TEST_F(TaskDependencyManagerTest, TestTaskForwarding) { + // Create 2 tasks, one dependent on the other. The first has no arguments. + int num_tasks = 2; + auto tasks = MakeTaskChain(num_tasks, {}, 1); + for (const auto &task : tasks) { + // Subscribe to each of the tasks' arguments. + auto arguments = task.GetDependencies(); + static_cast(task_dependency_manager_.SubscribeDependencies( + task.GetTaskSpecification().TaskId(), arguments)); + task_dependency_manager_.TaskPending(task); + } + + // Get the first task. + const auto task = tasks.front(); + TaskID task_id = task.GetTaskSpecification().TaskId(); + ObjectID return_id = task.GetTaskSpecification().ReturnId(0); + // Simulate forwarding the first task to a remote node. + task_dependency_manager_.UnsubscribeDependencies(task_id); + // The object returned by the first task should be considered remote once we + // cancel the forwarded task, since the second task depends on it. + EXPECT_CALL(object_manager_mock_, Pull(return_id)); + task_dependency_manager_.TaskCanceled(task_id); + + // Simulate the task executing on a remote node and its return value + // appearing locally. + EXPECT_CALL(object_manager_mock_, Cancel(return_id)); + auto ready_tasks = task_dependency_manager_.HandleObjectLocal(return_id); + // Check that the task that we kept is now ready to run. + ASSERT_EQ(ready_tasks.size(), 1); + ASSERT_EQ(ready_tasks.front(), tasks.back().GetTaskSpecification().TaskId()); +} + +TEST_F(TaskDependencyManagerTest, TestEviction) { + // Create a task with 3 arguments. + int num_arguments = 3; + std::vector arguments; + for (int i = 0; i < num_arguments; i++) { + arguments.push_back(ObjectID::from_random()); + } + TaskID task_id = TaskID::from_random(); + // No objects have been registered in the task dependency manager, so all + // arguments should be remote. + for (const auto &argument_id : arguments) { + EXPECT_CALL(object_manager_mock_, Pull(argument_id)); + } + // Subscribe to the task's dependencies. + bool ready = task_dependency_manager_.SubscribeDependencies(task_id, arguments); + ASSERT_FALSE(ready); + + // Tell the task dependency manager that each of the arguments is now + // available. + for (const auto &argument_id : arguments) { + EXPECT_CALL(object_manager_mock_, Cancel(argument_id)); + } + for (size_t i = 0; i < arguments.size(); i++) { + std::vector ready_tasks; + ready_tasks = task_dependency_manager_.HandleObjectLocal(arguments[i]); + if (i == arguments.size() - 1) { + ASSERT_EQ(ready_tasks.size(), 1); + ASSERT_EQ(ready_tasks.front(), task_id); + } else { + ASSERT_TRUE(ready_tasks.empty()); + } + } + + // Simulate each of the arguments getting evicted. Each object should now be + // considered remote. + for (const auto &argument_id : arguments) { + EXPECT_CALL(object_manager_mock_, Pull(argument_id)); + } + for (size_t i = 0; i < arguments.size(); i++) { + std::vector waiting_tasks; + waiting_tasks = task_dependency_manager_.HandleObjectMissing(arguments[i]); + if (i == 0) { + // The first eviction should cause the task to go back to the waiting + // state. + ASSERT_EQ(waiting_tasks.size(), 1); + ASSERT_EQ(waiting_tasks.front(), task_id); + } else { + // The subsequent evictions shouldn't cause any more tasks to go back to + // the waiting state. + ASSERT_TRUE(waiting_tasks.empty()); + } + } + + // Tell the task dependency manager that each of the arguments is available + // again. + for (const auto &argument_id : arguments) { + EXPECT_CALL(object_manager_mock_, Cancel(argument_id)); + } + for (size_t i = 0; i < arguments.size(); i++) { + std::vector ready_tasks; + ready_tasks = task_dependency_manager_.HandleObjectLocal(arguments[i]); + if (i == arguments.size() - 1) { + ASSERT_EQ(ready_tasks.size(), 1); + ASSERT_EQ(ready_tasks.front(), task_id); + } else { + ASSERT_TRUE(ready_tasks.empty()); + } + } +} + +} // namespace raylet + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +}