diff --git a/src/ray/common/scheduling/cluster_resource_scheduler.cc b/src/ray/common/scheduling/cluster_resource_scheduler.cc index 1e76a4a8e..14472fce3 100644 --- a/src/ray/common/scheduling/cluster_resource_scheduler.cc +++ b/src/ray/common/scheduling/cluster_resource_scheduler.cc @@ -14,6 +14,137 @@ #include "cluster_resource_scheduler.h" +std::string VectorToString(const std::vector &vector) { + std::stringstream buffer; + + buffer << "["; + for (size_t i = 0; i < vector.size(); i++) { + buffer << vector[i]; + if (i < vector.size() - 1) { + buffer << ", "; + } + } + buffer << "]"; + return buffer.str(); +} + +std::string UnorderedMapToString(const std::unordered_map &map) { + std::stringstream buffer; + + buffer << "["; + for (auto it = map.begin(); it != map.end(); ++it) { + buffer << "(" << it->first << ":" << it->second << ")"; + } + buffer << "]"; + return buffer.str(); +} + +/// Convert a map of resources to a TaskRequest data structure. +TaskRequest ResourceMapToTaskRequest( + StringIdMap &string_to_int_map, + const std::unordered_map &resource_map) { + size_t i = 0; + + TaskRequest task_request; + + task_request.predefined_resources.resize(PredefinedResources_MAX); + task_request.custom_resources.resize(resource_map.size()); + for (size_t i = 0; i < PredefinedResources_MAX; i++) { + task_request.predefined_resources[0].demand = 0; + task_request.predefined_resources[0].soft = false; + } + + for (auto const &resource : resource_map) { + if (resource.first == ray::kCPU_ResourceLabel) { + task_request.predefined_resources[CPU].demand = resource.second; + } else if (resource.first == ray::kGPU_ResourceLabel) { + task_request.predefined_resources[GPU].demand = resource.second; + } else if (resource.first == ray::kTPU_ResourceLabel) { + task_request.predefined_resources[TPU].demand = resource.second; + } else if (resource.first == ray::kMemory_ResourceLabel) { + task_request.predefined_resources[MEM].demand = resource.second; + } else { + task_request.custom_resources[i].id = string_to_int_map.Insert(resource.first); + task_request.custom_resources[i].demand = resource.second; + task_request.custom_resources[i].soft = false; + i++; + } + } + task_request.custom_resources.resize(i); + + return task_request; +} + +TaskRequest TaskResourceInstances::ToTaskRequest() const { + TaskRequest task_req; + task_req.predefined_resources.resize(PredefinedResources_MAX); + + for (size_t i = 0; i < PredefinedResources_MAX; i++) { + task_req.predefined_resources[i].demand = 0; + for (auto predefined_resource_instance : this->predefined_resources[i]) { + task_req.predefined_resources[i].demand += predefined_resource_instance; + } + } + + task_req.custom_resources.resize(this->custom_resources.size()); + size_t i = 0; + for (auto it = this->custom_resources.begin(); it != this->custom_resources.end(); + ++it) { + task_req.custom_resources[i].id = it->first; + task_req.custom_resources[i].soft = false; + task_req.custom_resources[i].demand = 0; + for (size_t j = 0; j < it->second.size(); j++) { + task_req.custom_resources[i].demand += it->second[j]; + } + i++; + } + return task_req; +} + +/// Convert a map of resources to a TaskRequest data structure. +/// +/// \param string_to_int_map: Map between names and ids maintained by the +/// \param resource_map_total: Total capacities of resources we want to convert. +/// \param resource_map_available: Available capacities of resources we want to convert. +/// +/// \request Conversion result to a TaskRequest data structure. +NodeResources ResourceMapToNodeResources( + StringIdMap &string_to_int_map, + const std::unordered_map &resource_map_total, + const std::unordered_map &resource_map_available) { + NodeResources node_resources; + node_resources.predefined_resources.resize(PredefinedResources_MAX); + for (size_t i = 0; i < PredefinedResources_MAX; i++) { + node_resources.predefined_resources[i].total = + node_resources.predefined_resources[i].available = 0; + } + + for (auto const &resource : resource_map_total) { + ResourceCapacity resource_capacity; + resource_capacity.total = (int64_t)resource.second; + auto it = resource_map_available.find(resource.first); + if (it == resource_map_available.end()) { + resource_capacity.available = 0; + } else { + resource_capacity.available = (int64_t)it->second; + } + if (resource.first == ray::kCPU_ResourceLabel) { + node_resources.predefined_resources[CPU] = resource_capacity; + } else if (resource.first == ray::kGPU_ResourceLabel) { + node_resources.predefined_resources[GPU] = resource_capacity; + } else if (resource.first == ray::kTPU_ResourceLabel) { + node_resources.predefined_resources[TPU] = resource_capacity; + } else if (resource.first == ray::kMemory_ResourceLabel) { + node_resources.predefined_resources[MEM] = resource_capacity; + } else { + // This is a custom resource. + node_resources.custom_resources.emplace(string_to_int_map.Insert(resource.first), + resource_capacity); + } + } + return node_resources; +} + bool NodeResources::operator==(const NodeResources &other) { for (size_t i = 0; i < PredefinedResources_MAX; i++) { if (this->predefined_resources[i].total != other.predefined_resources[i].total) { @@ -45,39 +176,25 @@ bool NodeResources::operator==(const NodeResources &other) { return true; } -std::string NodeResources::DebugString() { +std::string NodeResources::DebugString(StringIdMap string_to_in_map) const { std::stringstream buffer; - buffer << " node predefined resources {"; + buffer << " {"; for (size_t i = 0; i < static_cast(this->predefined_resources.size()); i++) { buffer << "(" << this->predefined_resources[i].total << ":" << this->predefined_resources[i].available << ") "; } - buffer << "}" << std::endl; + buffer << "}"; - buffer << " node custom resources {"; + buffer << " {"; for (auto it = this->custom_resources.begin(); it != this->custom_resources.end(); ++it) { - buffer << it->first << ":(" << it->second.total << ":" << it->second.available - << ") "; + buffer << string_to_in_map.Get(it->first) << ":(" << it->second.total << ":" + << it->second.available << ") "; } buffer << "}" << std::endl; return buffer.str(); } -std::string VectorToString(std::vector &vector) { - std::stringstream buffer; - - buffer << "["; - for (size_t i = 0; i < vector.size(); i++) { - buffer << vector[i]; - if (i < vector.size() - 1) { - buffer << ", "; - } - } - buffer << "]"; - return buffer.str(); -} - bool NodeResourceInstances::operator==(const NodeResourceInstances &other) { for (size_t i = 0; i < PredefinedResources_MAX; i++) { if (!EqualVectors(this->predefined_resources[i].total, @@ -110,20 +227,20 @@ bool NodeResourceInstances::operator==(const NodeResourceInstances &other) { return true; } -std::string NodeResourceInstances::DebugString() { +std::string NodeResourceInstances::DebugString(StringIdMap string_to_int_map) const { std::stringstream buffer; - buffer << " node predefined resources {"; + buffer << " {"; for (size_t i = 0; i < this->predefined_resources.size(); i++) { buffer << "(" << VectorToString(predefined_resources[i].total) << ":" << VectorToString(this->predefined_resources[i].available) << ") "; } - buffer << "}" << std::endl; + buffer << "}"; - buffer << " node custom resources {"; + buffer << " {"; for (auto it = this->custom_resources.begin(); it != this->custom_resources.end(); ++it) { - buffer << it->first << ":(" << VectorToString(it->second.total) << ":" - << VectorToString(it->second.available) << ") "; + buffer << string_to_int_map.Get(it->first) << ":(" << VectorToString(it->second.total) + << ":" << VectorToString(it->second.available) << ") "; } buffer << "}" << std::endl; return buffer.str(); @@ -144,40 +261,60 @@ TaskResourceInstances NodeResourceInstances::GetAvailableResourceInstances() { return task_resources; }; -std::string TaskRequest::DebugString() { +std::string TaskRequest::DebugString() const { std::stringstream buffer; - buffer << std::endl << " request predefined resources {"; + buffer << " {"; for (size_t i = 0; i < this->predefined_resources.size(); i++) { buffer << "(" << this->predefined_resources[i].demand << ":" << this->predefined_resources[i].soft << ") "; } - buffer << "}" << std::endl; + buffer << "}"; - buffer << " request custom resources {"; + buffer << " ["; for (size_t i = 0; i < this->custom_resources.size(); i++) { buffer << this->custom_resources[i].id << ":" << "(" << this->custom_resources[i].demand << ":" << this->custom_resources[i].soft << ") "; } - buffer << "}" << std::endl; + buffer << "]" << std::endl; return buffer.str(); } -std::string TaskResourceInstances::DebugString() { +bool TaskResourceInstances::IsEmpty() const { + // Check whether all resource instances of a task are zero. + for (const auto &predefined_resource : predefined_resources) { + for (const auto &predefined_resource_instance : predefined_resource) { + if (predefined_resource_instance != 0) { + return false; + } + } + } + + for (const auto custom_resource : custom_resources) { + for (const auto custom_resource_instances : custom_resource.second) { + if (custom_resource_instances != 0) { + return false; + } + } + } + return true; +} + +std::string TaskResourceInstances::DebugString() const { std::stringstream buffer; - buffer << std::endl << " task allocation: P {"; + buffer << std::endl << " Allocation: {"; for (size_t i = 0; i < this->predefined_resources.size(); i++) { buffer << VectorToString(this->predefined_resources[i]); } buffer << "}"; - buffer << " C {"; + buffer << " ["; for (auto it = this->custom_resources.begin(); it != this->custom_resources.end(); ++it) { buffer << it->first << ":" << VectorToString(it->second) << ", "; } - buffer << "}" << std::endl; + buffer << "]" << std::endl; return buffer.str(); } @@ -220,15 +357,19 @@ ClusterResourceScheduler::ClusterResourceScheduler( const std::string &local_node_id, const std::unordered_map &local_node_resources) { local_node_id_ = string_to_int_map_.Insert(local_node_id); - AddOrUpdateNode(local_node_id, local_node_resources, local_node_resources); + NodeResources node_resources = ResourceMapToNodeResources( + string_to_int_map_, local_node_resources, local_node_resources); + + AddOrUpdateNode(local_node_id_, node_resources); + InitLocalResources(node_resources); } void ClusterResourceScheduler::AddOrUpdateNode( const std::string &node_id, const std::unordered_map &resources_total, const std::unordered_map &resources_available) { - NodeResources node_resources; - ResourceMapToNodeResources(resources_total, resources_available, &node_resources); + NodeResources node_resources = ResourceMapToNodeResources( + string_to_int_map_, resources_total, resources_available); AddOrUpdateNode(string_to_int_map_.Insert(node_id), node_resources); } @@ -289,28 +430,35 @@ int64_t ClusterResourceScheduler::IsSchedulable(const TaskRequest &task_req, resources.predefined_resources[i].available) { if (task_req.predefined_resources[i].soft) { // A soft constraint has been violated. + // Just remember this as soft violations do not preclude a task + // from being scheduled. violations++; } else { - // A hard constraint has been violated. + // A hard constraint has been violated, so we cannot schedule + // this task request. return -1; } } } - for (size_t i = 0; i < task_req.custom_resources.size(); i++) { - auto it = resources.custom_resources.find(task_req.custom_resources[i].id); + // No check custom resources. + for (const auto task_req_custom_resource : task_req.custom_resources) { + auto it = resources.custom_resources.find(task_req_custom_resource.id); if (it == resources.custom_resources.end()) { - // Requested resource doesn't exist at this node. - if (task_req.custom_resources[i].soft) { + // Requested resource doesn't exist at this node. However, this + // is a soft constraint, so just increment "violations" and continue. + if (task_req_custom_resource.soft) { violations++; } else { + // This is a hard constraint so cannot schedule this task request. return -1; } } else { - if (task_req.custom_resources[i].demand > it->second.available) { - // Resource constraint is violated. - if (task_req.custom_resources[i].soft) { + if (task_req_custom_resource.demand > it->second.available) { + // Resource constraint is violated, but since it is soft + // just increase the "violations" and continue. + if (task_req_custom_resource.soft) { violations++; } else { return -1; @@ -323,7 +471,7 @@ int64_t ClusterResourceScheduler::IsSchedulable(const TaskRequest &task_req, auto it_p = task_req.placement_hints.find(node_id); if (it_p == task_req.placement_hints.end()) { // Node not found in the placement_hints list, so - // record this a soft constraint violation. + // record this as a soft constraint violation. violations++; } } @@ -333,7 +481,8 @@ int64_t ClusterResourceScheduler::IsSchedulable(const TaskRequest &task_req, int64_t ClusterResourceScheduler::GetBestSchedulableNode(const TaskRequest &task_req, int64_t *total_violations) { - // Min number of violations across all nodes that can schedule the request. + // Minimum number of soft violations across all nodes that can schedule the request. + // We will pick the node with the smallest number of soft violations. int64_t min_violations = INT_MAX; // Node associated to min_violations. int64_t best_node = -1; @@ -350,9 +499,8 @@ int64_t ClusterResourceScheduler::GetBestSchedulableNode(const TaskRequest &task // Check whether any node in the request placement_hints, satisfes // all resource constraints of the request. - for (auto it_p = task_req.placement_hints.begin(); - it_p != task_req.placement_hints.end(); ++it_p) { - auto it = nodes_.find(*it_p); + for (const auto &task_req_placement_hint : task_req.placement_hints) { + auto it = nodes_.find(task_req_placement_hint); if (it != nodes_.end()) { if (IsSchedulable(task_req, it->first, it->second) == 0) { return it->first; @@ -360,19 +508,19 @@ int64_t ClusterResourceScheduler::GetBestSchedulableNode(const TaskRequest &task } } - for (auto it = nodes_.begin(); it != nodes_.end(); ++it) { + for (const auto &node : nodes_) { // Return -1 if node not schedulable. otherwise return the number // of soft constraint violations. int64_t violations; - if ((violations = IsSchedulable(task_req, it->first, it->second)) == -1) { + if ((violations = IsSchedulable(task_req, node.first, node.second)) == -1) { continue; } // Update the node with the smallest number of soft constraints violated. if (min_violations > violations) { min_violations = violations; - best_node = it->first; + best_node = node.first; } if (violations == 0) { *total_violations = 0; @@ -386,14 +534,15 @@ int64_t ClusterResourceScheduler::GetBestSchedulableNode(const TaskRequest &task std::string ClusterResourceScheduler::GetBestSchedulableNode( const std::unordered_map &task_resources, int64_t *total_violations) { - TaskRequest task_request; - ResourceMapToTaskRequest(task_resources, &task_request); + TaskRequest task_request = ResourceMapToTaskRequest(string_to_int_map_, task_resources); int64_t node_id = GetBestSchedulableNode(task_request, total_violations); std::string id_string; if (node_id == -1) { + // This is not a schedulable node, so return empty string. return ""; } + // Return the string name of the node. return string_to_int_map_.Get(node_id); } @@ -416,11 +565,11 @@ bool ClusterResourceScheduler::SubtractNodeAvailableResources( task_req.predefined_resources[i].demand); } - for (size_t i = 0; i < task_req.custom_resources.size(); i++) { - auto it = resources.custom_resources.find(task_req.custom_resources[i].id); + for (const auto &task_req_custom_resource : task_req.custom_resources) { + auto it = resources.custom_resources.find(task_req_custom_resource.id); if (it != resources.custom_resources.end()) { it->second.available = - std::max(0., it->second.available - task_req.custom_resources[i].demand); + std::max(0., it->second.available - task_req_custom_resource.demand); } } return true; @@ -429,8 +578,7 @@ bool ClusterResourceScheduler::SubtractNodeAvailableResources( bool ClusterResourceScheduler::SubtractNodeAvailableResources( const std::string &node_id, const std::unordered_map &resource_map) { - TaskRequest task_request; - ResourceMapToTaskRequest(resource_map, &task_request); + TaskRequest task_request = ResourceMapToTaskRequest(string_to_int_map_, resource_map); return SubtractNodeAvailableResources(string_to_int_map_.Get(node_id), task_request); } @@ -449,11 +597,11 @@ bool ClusterResourceScheduler::AddNodeAvailableResources(int64_t node_id, resources.predefined_resources[i].total); } - for (size_t i = 0; i < task_req.custom_resources.size(); i++) { - auto it = resources.custom_resources.find(task_req.custom_resources[i].id); + for (const auto &task_req_custom_resource : task_req.custom_resources) { + auto it = resources.custom_resources.find(task_req_custom_resource.id); if (it != resources.custom_resources.end()) { it->second.available = std::min( - it->second.available + task_req.custom_resources[i].demand, it->second.total); + it->second.available + task_req_custom_resource.demand, it->second.total); } } return true; @@ -462,8 +610,7 @@ bool ClusterResourceScheduler::AddNodeAvailableResources(int64_t node_id, bool ClusterResourceScheduler::AddNodeAvailableResources( const std::string &node_id, const std::unordered_map &resource_map) { - TaskRequest task_request; - ResourceMapToTaskRequest(resource_map, &task_request); + TaskRequest task_request = ResourceMapToTaskRequest(string_to_int_map_, resource_map); return AddNodeAvailableResources(string_to_int_map_.Get(node_id), task_request); } @@ -480,79 +627,19 @@ bool ClusterResourceScheduler::GetNodeResources(int64_t node_id, int64_t ClusterResourceScheduler::NumNodes() { return nodes_.size(); } -void ClusterResourceScheduler::ResourceMapToNodeResources( - const std::unordered_map &resource_map_total, - const std::unordered_map &resource_map_available, - NodeResources *node_resources) { - node_resources->predefined_resources.resize(PredefinedResources_MAX); - for (size_t i = 0; i < PredefinedResources_MAX; i++) { - node_resources->predefined_resources[i].total = - node_resources->predefined_resources[i].available = 0; - } - - for (auto it = resource_map_total.begin(); it != resource_map_total.end(); ++it) { - ResourceCapacity resource_capacity; - resource_capacity.total = (int64_t)it->second; - auto it2 = resource_map_available.find(it->first); - if (it2 == resource_map_available.end()) { - resource_capacity.available = 0; - } else { - resource_capacity.available = (int64_t)it2->second; - } - if (it->first == ray::kCPU_ResourceLabel) { - node_resources->predefined_resources[CPU] = resource_capacity; - } else if (it->first == ray::kGPU_ResourceLabel) { - node_resources->predefined_resources[GPU] = resource_capacity; - } else if (it->first == ray::kTPU_ResourceLabel) { - node_resources->predefined_resources[TPU] = resource_capacity; - } else if (it->first == ray::kMemory_ResourceLabel) { - node_resources->predefined_resources[MEM] = resource_capacity; - } else { - // This is a custom resource. - node_resources->custom_resources.emplace(string_to_int_map_.Insert(it->first), - resource_capacity); - } - } -} - -void ClusterResourceScheduler::ResourceMapToTaskRequest( - const std::unordered_map &resource_map, - TaskRequest *task_request) { - size_t i = 0; - - task_request->predefined_resources.resize(PredefinedResources_MAX); - task_request->custom_resources.resize(resource_map.size()); - for (size_t i = 0; i < PredefinedResources_MAX; i++) { - task_request->predefined_resources[0].demand = 0; - task_request->predefined_resources[0].soft = false; - } - - for (auto it = resource_map.begin(); it != resource_map.end(); ++it) { - if (it->first == ray::kCPU_ResourceLabel) { - task_request->predefined_resources[CPU].demand = it->second; - } else if (it->first == ray::kGPU_ResourceLabel) { - task_request->predefined_resources[GPU].demand = it->second; - } else if (it->first == ray::kTPU_ResourceLabel) { - task_request->predefined_resources[TPU].demand = it->second; - } else if (it->first == ray::kMemory_ResourceLabel) { - task_request->predefined_resources[MEM].demand = it->second; - } else { - task_request->custom_resources[i].id = string_to_int_map_.Insert(it->first); - task_request->custom_resources[i].demand = it->second; - task_request->custom_resources[i].soft = false; - i++; - } - } - task_request->custom_resources.resize(i); -} - void ClusterResourceScheduler::UpdateResourceCapacity(const std::string &client_id_string, const std::string &resource_name, int64_t resource_total) { int64_t client_id = string_to_int_map_.Get(client_id_string); + auto it = nodes_.find(client_id); if (it == nodes_.end()) { - return; + NodeResources node_resources; + node_resources.predefined_resources.resize(PredefinedResources_MAX); + client_id = string_to_int_map_.Insert(client_id_string); + RAY_CHECK(nodes_.emplace(client_id, node_resources).second); + it = nodes_.find(client_id); + RAY_CHECK(it != nodes_.end()); } int idx = -1; @@ -588,10 +675,11 @@ void ClusterResourceScheduler::UpdateResourceCapacity(const std::string &client_ if (itr->second.total < 0) { itr->second.total = 0; } + } else { + ResourceCapacity resource_capacity; + resource_capacity.total = resource_capacity.available = resource_total; + it->second.custom_resources.emplace(resource_id, resource_capacity); } - ResourceCapacity resource_capacity; - resource_capacity.total = resource_capacity.available = resource_total; - it->second.custom_resources.emplace(resource_id, resource_capacity); } } @@ -625,19 +713,19 @@ void ClusterResourceScheduler::DeleteResource(const std::string &client_id_strin } } -std::string ClusterResourceScheduler::DebugString(void) { +std::string ClusterResourceScheduler::DebugString(void) const { std::stringstream buffer; - buffer << std::endl << "local node id: " << local_node_id_ << std::endl; - for (auto it = nodes_.begin(); it != nodes_.end(); ++it) { - buffer << "node id: " << it->first << std::endl; - buffer << it->second.DebugString(); + buffer << "\n Local id: " << local_node_id_; + buffer << " Local resources: " << local_resources_.DebugString(string_to_int_map_); + for (auto &node : nodes_) { + buffer << " node id: " << node.first; + buffer << node.second.DebugString(string_to_int_map_); } return buffer.str(); } void ClusterResourceScheduler::InitResourceInstances( - double total, bool unit_instances, - ResourceInstanceCapacities *instance_list /* return */) { + double total, bool unit_instances, ResourceInstanceCapacities *instance_list) { if (unit_instances) { size_t num_instances = static_cast(total); instance_list->total.resize(num_instances); @@ -678,29 +766,38 @@ void ClusterResourceScheduler::InitLocalResources(const NodeResources &node_reso } } -void ClusterResourceScheduler::AddAvailableResourceInstances( - std::vector available, - ResourceInstanceCapacities *resource_instances /* return */) { +std::vector ClusterResourceScheduler::AddAvailableResourceInstances( + std::vector available, ResourceInstanceCapacities *resource_instances) { + std::vector overflow(available.size(), 0.); for (size_t i = 0; i < available.size(); i++) { - resource_instances->available[i] = std::min( - resource_instances->available[i] + available[i], resource_instances->total[i]); + resource_instances->available[i] = resource_instances->available[i] + available[i]; + if (resource_instances->available[i] > resource_instances->total[i]) { + overflow[i] = resource_instances->available[i] - resource_instances->total[i]; + resource_instances->available[i] = resource_instances->total[i]; + } } + + return overflow; } -void ClusterResourceScheduler::SubtractAvailableResourceInstances( - std::vector available, - ResourceInstanceCapacities *resource_instances /* return */) { +std::vector ClusterResourceScheduler::SubtractAvailableResourceInstances( + std::vector available, ResourceInstanceCapacities *resource_instances) { RAY_CHECK(available.size() == resource_instances->available.size()); + std::vector underflow(available.size(), 0.); for (size_t i = 0; i < available.size(); i++) { - resource_instances->available[i] = - std::max(resource_instances->available[i] - available[i], 0.); + resource_instances->available[i] = resource_instances->available[i] - available[i]; + if (resource_instances->available[i] < 0) { + underflow[i] = -resource_instances->available[i]; + resource_instances->available[i] = 0; + } } + return underflow; } bool ClusterResourceScheduler::AllocateResourceInstances( double demand, bool soft, std::vector &available, - std::vector *allocation /* return */) { + std::vector *allocation) { allocation->resize(available.size()); double remaining_demand = demand; @@ -790,14 +887,9 @@ bool ClusterResourceScheduler::AllocateResourceInstances( } bool ClusterResourceScheduler::AllocateTaskResourceInstances( - const TaskRequest &task_req, TaskResourceInstances *task_allocation /* return */) { - auto it = nodes_.find(local_node_id_); - if (it == nodes_.end()) { - return false; - } - - // Just double check this node can still schedule the task request. - if (IsSchedulable(task_req, local_node_id_, it->second) == -1) { + const TaskRequest &task_req, std::shared_ptr task_allocation) { + RAY_CHECK(task_allocation != nullptr); + if (nodes_.find(local_node_id_) == nodes_.end()) { return false; } @@ -810,19 +902,19 @@ bool ClusterResourceScheduler::AllocateTaskResourceInstances( &task_allocation->predefined_resources[i])) { // Allocation failed. Restore node's local resources by freeing the resources // of the failed allocation. - FreeTaskResourceInstances(*task_allocation); + FreeTaskResourceInstances(task_allocation); return false; } } } - for (size_t i = 0; i < task_req.custom_resources.size(); i++) { - auto it = local_resources_.custom_resources.find(task_req.custom_resources[i].id); + for (const auto &task_req_custom_resource : task_req.custom_resources) { + auto it = local_resources_.custom_resources.find(task_req_custom_resource.id); if (it != local_resources_.custom_resources.end()) { - if (task_req.custom_resources[i].demand > 0) { + if (task_req_custom_resource.demand > 0) { std::vector allocation; - bool success = AllocateResourceInstances(task_req.custom_resources[i].demand, - task_req.custom_resources[i].soft, + bool success = AllocateResourceInstances(task_req_custom_resource.demand, + task_req_custom_resource.soft, it->second.available, &allocation); // Even if allocation failed we need to remember partial allocations to correctly // free resources. @@ -830,7 +922,7 @@ bool ClusterResourceScheduler::AllocateTaskResourceInstances( if (!success) { // Allocation failed. Restore node's local resources by freeing the resources // of the failed allocation. - FreeTaskResourceInstances(*task_allocation); + FreeTaskResourceInstances(task_allocation); return false; } } @@ -841,30 +933,128 @@ bool ClusterResourceScheduler::AllocateTaskResourceInstances( return true; } -void ClusterResourceScheduler::FreeTaskResourceInstances( - TaskResourceInstances &task_allocation) { +void ClusterResourceScheduler::UpdateLocalAvailableResourcesFromResourceInstances() { + auto it_local_node = nodes_.find(local_node_id_); + RAY_CHECK(it_local_node != nodes_.end()); + for (size_t i = 0; i < PredefinedResources_MAX; i++) { - AddAvailableResourceInstances(task_allocation.predefined_resources[i], - &local_resources_.predefined_resources[i]); + it_local_node->second.predefined_resources[i].available = 0; + for (size_t j = 0; j < local_resources_.predefined_resources[i].available.size(); + j++) { + it_local_node->second.predefined_resources[i].available += + local_resources_.predefined_resources[i].available[j]; + } } - for (auto it = task_allocation.custom_resources.begin(); - it != task_allocation.custom_resources.end(); it++) { - auto it_local = local_resources_.custom_resources.find(it->first); - if (it_local != local_resources_.custom_resources.end()) { - AddAvailableResourceInstances(it->second, &it_local->second); + for (auto &custom_resource : it_local_node->second.custom_resources) { + auto it = local_resources_.custom_resources.find(custom_resource.first); + if (it != local_resources_.custom_resources.end()) { + custom_resource.second.available = 0; + for (const auto available : it->second.available) { + custom_resource.second.available += available; + } } } } -void ClusterResourceScheduler::AddCPUResourceInstances( - std::vector &cpu_instances) { - AddAvailableResourceInstances(cpu_instances, - &local_resources_.predefined_resources[CPU]); +void ClusterResourceScheduler::FreeTaskResourceInstances( + std::shared_ptr task_allocation) { + RAY_CHECK(task_allocation != nullptr); + for (size_t i = 0; i < PredefinedResources_MAX; i++) { + AddAvailableResourceInstances(task_allocation->predefined_resources[i], + &local_resources_.predefined_resources[i]); + } + + for (const auto task_allocation_custom_resource : task_allocation->custom_resources) { + auto it = + local_resources_.custom_resources.find(task_allocation_custom_resource.first); + if (it != local_resources_.custom_resources.end()) { + AddAvailableResourceInstances(task_allocation_custom_resource.second, &it->second); + } + } } -void ClusterResourceScheduler::SubtractCPUResourceInstances( +std::vector ClusterResourceScheduler::AddCPUResourceInstances( std::vector &cpu_instances) { - SubtractAvailableResourceInstances(cpu_instances, - &local_resources_.predefined_resources[CPU]); + if (cpu_instances.size() == 0) { + return cpu_instances; // No oveerflow. + } + RAY_CHECK(nodes_.find(local_node_id_) != nodes_.end()); + + auto overflow = AddAvailableResourceInstances( + cpu_instances, &local_resources_.predefined_resources[CPU]); + UpdateLocalAvailableResourcesFromResourceInstances(); + + return overflow; +} + +std::vector ClusterResourceScheduler::SubtractCPUResourceInstances( + std::vector &cpu_instances) { + if (cpu_instances.size() == 0) { + return cpu_instances; // No underflow. + } + RAY_CHECK(nodes_.find(local_node_id_) != nodes_.end()); + + auto underflow = SubtractAvailableResourceInstances( + cpu_instances, &local_resources_.predefined_resources[CPU]); + UpdateLocalAvailableResourcesFromResourceInstances(); + + return underflow; +} + +bool ClusterResourceScheduler::AllocateTaskResources( + int64_t node_id, const TaskRequest &task_req, + std::shared_ptr task_allocation) { + if (node_id == local_node_id_) { + RAY_CHECK(task_allocation != nullptr); + if (AllocateTaskResourceInstances(task_req, task_allocation)) { + UpdateLocalAvailableResourcesFromResourceInstances(); + return true; + } + } else { + if (SubtractNodeAvailableResources(node_id, task_req)) { + return true; + } + } + return false; +} + +bool ClusterResourceScheduler::AllocateLocalTaskResources( + const std::unordered_map &task_resources, + std::shared_ptr task_allocation) { + RAY_CHECK(task_allocation != nullptr); + TaskRequest task_request = ResourceMapToTaskRequest(string_to_int_map_, task_resources); + return AllocateTaskResources(local_node_id_, task_request, task_allocation); +} + +std::string ClusterResourceScheduler::GetResourceNameFromIndex(int64_t res_idx) { + if (res_idx == CPU) { + return ray::kCPU_ResourceLabel; + } else if (res_idx == GPU) { + return ray::kGPU_ResourceLabel; + } else if (res_idx == TPU) { + return ray::kTPU_ResourceLabel; + } else if (res_idx == MEM) { + return ray::kMemory_ResourceLabel; + } else { + return string_to_int_map_.Get((uint64_t)res_idx); + } +} + +void ClusterResourceScheduler::AllocateRemoteTaskResources( + std::string &node_string, + const std::unordered_map &task_resources) { + TaskRequest task_request = ResourceMapToTaskRequest(string_to_int_map_, task_resources); + auto node_id = string_to_int_map_.Insert(node_string); + RAY_CHECK(node_id != local_node_id_); + AllocateTaskResources(node_id, task_request, nullptr); +} + +void ClusterResourceScheduler::FreeLocalTaskResources( + std::shared_ptr task_allocation) { + if (task_allocation == nullptr || task_allocation->IsEmpty()) { + return; + } + FreeTaskResourceInstances(task_allocation); + UpdateLocalAvailableResourcesFromResourceInstances(); } diff --git a/src/ray/common/scheduling/cluster_resource_scheduler.h b/src/ray/common/scheduling/cluster_resource_scheduler.h index 8ad84064c..fb114b146 100644 --- a/src/ray/common/scheduling/cluster_resource_scheduler.h +++ b/src/ray/common/scheduling/cluster_resource_scheduler.h @@ -60,6 +60,7 @@ struct ResourceRequestWithId : ResourceRequest { int64_t id; }; +// Data structure specifying the capacity of each resource requested by a task. class TaskRequest { public: /// List of predefined resources required by the task. @@ -72,10 +73,11 @@ class TaskRequest { /// nodes in this list can schedule this task. absl::flat_hash_set placement_hints; /// Returns human-readable string for this task request. - std::string DebugString(); + std::string DebugString() const; }; -// Task request specifying instances for each resource. +// Data structure specifying the capacity of each instance of each resource +// allocated to a task. class TaskResourceInstances { public: /// The list of instances of each predifined resource allocated to a task. @@ -83,10 +85,20 @@ class TaskResourceInstances { /// The list of instances of each custom resource allocated to a task. absl::flat_hash_map> custom_resources; bool operator==(const TaskResourceInstances &other); + /// For each resource of this request aggregate its instances. + TaskRequest ToTaskRequest() const; /// Get CPU instances only. - std::vector GetCPUInstances() { return this->predefined_resources[CPU]; }; + std::vector GetCPUInstances() const { + if (!this->predefined_resources.empty()) { + return this->predefined_resources[CPU]; + } else { + return {}; + } + }; + /// Check whether there are no resource instances. + bool IsEmpty() const; /// Returns human-readable string for these resources. - std::string DebugString(); + std::string DebugString() const; }; /// Total and available capacities of each resource of a node. @@ -100,7 +112,7 @@ class NodeResources { /// Returns if this equals another node resources. bool operator==(const NodeResources &other); /// Returns human-readable string for these resources. - std::string DebugString(); + std::string DebugString(StringIdMap string_to_int_map) const; }; /// Total and available capacities of each resource instance. @@ -117,7 +129,7 @@ class NodeResourceInstances { /// Returns if this equals another node resources. bool operator==(const NodeResourceInstances &other); /// Returns human-readable string for these resources. - std::string DebugString(); + std::string DebugString(StringIdMap string_to_int_map) const; }; /// Class encapsulating the cluster resources and the logic to assign @@ -149,6 +161,19 @@ class ClusterResourceScheduler { const absl::flat_hash_map &new_custom_resources, absl::flat_hash_map *old_custom_resources); + /// Subtract the resources required by a given task request (task_req) from + /// a given node (node_id). + /// + /// \param node_id Node whose resources we allocate. Can be the local or a remote node. + /// \param task_req Task for which we allocate resources. + /// \param task_allocation Resources allocated to the task at instance granularity. + /// This is a return parameter. + /// + /// \return True if the node has enough resources to satisfy the task request. + /// False otherwise. + bool AllocateTaskResources(int64_t node_id, const TaskRequest &task_req, + std::shared_ptr task_allocation); + public: ClusterResourceScheduler(void){}; @@ -163,6 +188,9 @@ class ClusterResourceScheduler { const std::string &local_node_id, const std::unordered_map &local_node_resources); + // Mapping from predefined resource indexes to resource strings + std::string GetResourceNameFromIndex(int64_t res_idx); + /// Add a new node or overwrite the resources of an existing node. /// /// \param node_id: Node ID. @@ -264,24 +292,19 @@ class ClusterResourceScheduler { /// Get number of nodes in the cluster. int64_t NumNodes(); - /// Convert a map of resources to a TaskRequest data structure. - void ResourceMapToTaskRequest( - const std::unordered_map &resource_map, - TaskRequest *task_request); - - /// Convert a map of resources to a TaskRequest data structure. - void ResourceMapToNodeResources( - const std::unordered_map &resource_map_total, - const std::unordered_map &resource_map_available, - NodeResources *node_resources); - - /// Update total capacity of resource resource_name at node client_id. - void UpdateResourceCapacity(const std::string &client_id, + /// Update total capacity of a given resource of a given node. + /// + /// \param node_name: Node whose resource we want to update. + /// \param resource_name: Resource which we want to update. + /// \param resource_total: New capacity of the resource. + void UpdateResourceCapacity(const std::string &node_name, const std::string &resource_name, int64_t resource_total); - /// Delete resource resource_name from node cleint_id_string. - void DeleteResource(const std::string &client_id_string, - const std::string &resource_name); + /// Delete a given resource from a given node. + /// + /// \param node_name: Node whose resource we want to delete. + /// \param resource_name: Resource we want to delete + void DeleteResource(const std::string &node_name, const std::string &resource_name); /// Return local resources. NodeResourceInstances GetLocalResources() { return local_resources_; }; @@ -305,25 +328,27 @@ class ClusterResourceScheduler { ResourceInstanceCapacities *instance_list); /// Allocate enough capacity across the instances of a resource to satisfy "demand". - /// If resource has multiple unit-capacity instance, we consider two cases. + /// If resource has multiple unit-capacity instances, we consider two cases. /// /// 1) If the constraint is hard, allocate full unit-capacity instances until - /// demand becomes fractional, and then satisfy the fractional deman using the + /// demand becomes fractional, and then satisfy the fractional demand using the /// instance with the smallest available capacity that can satisfy the fractional /// demand. For example, assume a resource conisting of 4 instances, with available /// capacities: (1., 1., .7, 0.5) and deman of 1.2. Then we allocate one full /// instance and then allocate 0.2 of the 0.5 instance (as this is the instance /// with the smalest available capacity that can satisfy the remaining demand of 0.2). - /// As a result remaining available capacities will be (0., 1., .7, .2). - /// Thus, if the constraint is hard, we will allocate at most a fractional resource. + /// As a result remaining available capacities will be (0., 1., .7, .3). + /// Thus, if the constraint is hard, we will allocate a bunch of full instances and + /// at most a fractional instance. /// /// 2) If the constraint is soft, we can allocate multiple fractional resources, /// and even overallocate the resource. For example, in the previous case, if we /// have a demand of 1.8, we can allocate one full instance, the 0.5 instance, and - /// 0.1 from the 0.7 instance. Furthermore, if the demand is 3.5, then we allocate + /// 0.3 from the 0.7 instance. Furthermore, if the demand is 3.5, then we allocate /// all instances, and return success (true), despite the fact that the total /// available capacity of the rwsource is 3.2 (= 1. + 1. + .7 + .5), which is less - /// than the demand, 3.5. + /// than the demand, 3.5. In this case, the remaining available resource is + /// (0., 0., 0., 0.) /// /// \param demand: The resource amount to be allocated. /// \param soft: Specifies whether this demand has soft or hard constraints. @@ -340,45 +365,94 @@ class ClusterResourceScheduler { /// /// \param task_req: Resources requested by a task. /// \param task_allocation: Local resources allocated to satsify task_req demand. - /// This is an output argument. /// /// \return true, if allocation successful. If false, the caller needs to free the /// allocated resources, i.e., task_allocation. - bool AllocateTaskResourceInstances(const TaskRequest &task_req, - TaskResourceInstances *task_allocation); + bool AllocateTaskResourceInstances( + const TaskRequest &task_req, + std::shared_ptr task_allocation); /// Free resources which were allocated with a task. The freed resources are /// added back to the node's local available resources. /// /// \param task_allocation: Task's resources to be freed. - void FreeTaskResourceInstances(TaskResourceInstances &task_allocation); + void FreeTaskResourceInstances(std::shared_ptr task_allocation); /// Increase the available capacities of the instances of a given resource. /// /// \param available A list of available capacities for resource's instances. /// \param resource_instances List of the resource instances being updated. - void AddAvailableResourceInstances(std::vector available, - ResourceInstanceCapacities *resource_instances); + /// + /// \return Overflow capacities of "resource_instances" after adding instance + /// capacities in "available", i.e., + /// min(available + resource_instances.available, resource_instances.total) + std::vector AddAvailableResourceInstances( + std::vector available, ResourceInstanceCapacities *resource_instances); /// Decrease the available capacities of the instances of a given resource. /// /// \param free A list of capacities for resource's instances to be freed. /// \param resource_instances List of the resource instances being updated. - void SubtractAvailableResourceInstances(std::vector free, - ResourceInstanceCapacities *resource_instances); + /// \return Underflow of "resource_instances" after subtracting instance + /// capacities in "available", i.e.,. + /// max(available - reasource_instances.available, 0) + std::vector SubtractAvailableResourceInstances( + std::vector available, ResourceInstanceCapacities *resource_instances); /// Increase the available CPU instances of this node. /// /// \param cpu_instances CPU instances to be added to available cpus. - void AddCPUResourceInstances(std::vector &cpu_instances); - - /// Decrease the available cpu instances of this node. /// - /// \param cpu_instances Cpu instances to be removed from available cpus. - void SubtractCPUResourceInstances(std::vector &cpu_instances); + /// \return Overflow capacities of CPU instances after adding CPU + /// capacities in cpu_instances. + std::vector AddCPUResourceInstances(std::vector &cpu_instances); + + /// Decrease the available CPU instances of this node. + /// + /// \param cpu_instances CPU instances to be removed from available cpus. + /// + /// \return Underflow capacities of CPU instances after subtracting CPU + /// capacities in cpu_instances. + std::vector SubtractCPUResourceInstances(std::vector &cpu_instances); + + /// Subtract the resources required by a given task request (task_req) from the + /// local node. This function also updates the local node resources + /// at the instance granularity. + /// + /// \param task_req Task for which we allocate resources. + /// \param task_allocation Resources allocated to the task at instance granularity. + /// This is a return parameter. + /// + /// \return True if local node has enough resources to satisfy the task request. + /// False otherwise. + bool AllocateLocalTaskResources( + const std::unordered_map &task_resources, + std::shared_ptr task_allocation); + + /// Subtract the resources required by a given task request (task_req) from a given + /// remote node. + /// + /// \param node_id Remote node whose resources we allocate. + /// \param task_req Task for which we allocate resources. + void AllocateRemoteTaskResources( + std::string &node_id, + const std::unordered_map &task_resources); + + void FreeLocalTaskResources(std::shared_ptr task_allocation); + + /// Update the available resources of the local node given + /// the available instances of each resource of the local node. + /// Basically, this means computing the available resources + /// by adding up the available quantities of each instance of that + /// resources. + /// + /// Example: Assume the local node has four GPU instances with the + /// following availabilities: 0.2, 0.3, 0.1, 1. Then the total GPU + // resources availabile at that node is 0.2 + 0.3 + 0.1 + 1. = 1.6 + void UpdateLocalAvailableResourcesFromResourceInstances(); /// Return human-readable string for this scheduler state. - std::string DebugString(); + std::string DebugString() const; }; #endif // RAY_COMMON_SCHEDULING_SCHEDULING_H diff --git a/src/ray/common/scheduling/scheduling_test.cc b/src/ray/common/scheduling/scheduling_test.cc index 28a97ea88..087af7d1e 100644 --- a/src/ray/common/scheduling/scheduling_test.cc +++ b/src/ray/common/scheduling/scheduling_test.cc @@ -596,9 +596,10 @@ TEST_F(SchedulingTest, TaskResourceInstancesTest) { EmptyBoolVector, EmptyIntVector); NodeResourceInstances old_local_resources = cluster_resources.GetLocalResources(); - TaskResourceInstances task_allocation; + std::shared_ptr task_allocation = + std::make_shared(); bool success = - cluster_resources.AllocateTaskResourceInstances(task_req, &task_allocation); + cluster_resources.AllocateTaskResourceInstances(task_req, task_allocation); ASSERT_EQ(success, true); @@ -621,9 +622,10 @@ TEST_F(SchedulingTest, TaskResourceInstancesTest) { EmptyBoolVector, EmptyIntVector); NodeResourceInstances old_local_resources = cluster_resources.GetLocalResources(); - TaskResourceInstances task_allocation; + std::shared_ptr task_allocation = + std::make_shared(); bool success = - cluster_resources.AllocateTaskResourceInstances(task_req, &task_allocation); + cluster_resources.AllocateTaskResourceInstances(task_req, task_allocation); ASSERT_EQ(success, false); ASSERT_EQ((cluster_resources.GetLocalResources() == old_local_resources), true); @@ -642,9 +644,10 @@ TEST_F(SchedulingTest, TaskResourceInstancesTest) { EmptyBoolVector, EmptyIntVector); NodeResourceInstances old_local_resources = cluster_resources.GetLocalResources(); - TaskResourceInstances task_allocation; + std::shared_ptr task_allocation = + std::make_shared(); bool success = - cluster_resources.AllocateTaskResourceInstances(task_req, &task_allocation); + cluster_resources.AllocateTaskResourceInstances(task_req, task_allocation); ASSERT_EQ(success, true); @@ -677,9 +680,10 @@ TEST_F(SchedulingTest, TaskResourceInstancesTest) { EmptyIntVector); NodeResourceInstances old_local_resources = cluster_resources.GetLocalResources(); - TaskResourceInstances task_allocation; + std::shared_ptr task_allocation = + std::make_shared(); bool success = - cluster_resources.AllocateTaskResourceInstances(task_req, &task_allocation); + cluster_resources.AllocateTaskResourceInstances(task_req, task_allocation); ASSERT_EQ(success, true); @@ -706,9 +710,10 @@ TEST_F(SchedulingTest, TaskResourceInstancesTest) { EmptyIntVector); NodeResourceInstances old_local_resources = cluster_resources.GetLocalResources(); - TaskResourceInstances task_allocation; + std::shared_ptr task_allocation = + std::make_shared(); bool success = - cluster_resources.AllocateTaskResourceInstances(task_req, &task_allocation); + cluster_resources.AllocateTaskResourceInstances(task_req, task_allocation); ASSERT_EQ(success, false); ASSERT_EQ((cluster_resources.GetLocalResources() == old_local_resources), true); @@ -732,9 +737,10 @@ TEST_F(SchedulingTest, TaskResourceInstancesTest) { EmptyIntVector); NodeResourceInstances old_local_resources = cluster_resources.GetLocalResources(); - TaskResourceInstances task_allocation; + std::shared_ptr task_allocation = + std::make_shared(); bool success = - cluster_resources.AllocateTaskResourceInstances(task_req, &task_allocation); + cluster_resources.AllocateTaskResourceInstances(task_req, task_allocation); ASSERT_EQ(success, true); @@ -752,6 +758,138 @@ TEST_F(SchedulingTest, TaskResourceInstancesTest) { } } +TEST_F(SchedulingTest, TaskResourceInstancesTest2) { + { + NodeResources node_resources; + vector pred_capacities{4 /* CPU */, 4 /* MEM */, 5 /* GPU */}; + vector cust_ids{1, 2}; + vector cust_capacities{4, 4}; + initNodeResources(node_resources, pred_capacities, cust_ids, cust_capacities); + ClusterResourceScheduler cluster_resources(0, node_resources); + + TaskRequest task_req; + vector pred_demands = {2. /* CPU */, 2. /* MEM */, 1.5 /* GPU */}; + vector pred_soft = {false}; + vector cust_demands{3, 2}; + vector cust_soft{false, false}; + initTaskRequest(task_req, pred_demands, pred_soft, cust_ids, cust_demands, cust_soft, + EmptyIntVector); + + std::shared_ptr task_allocation = + std::make_shared(); + bool success = + cluster_resources.AllocateTaskResourceInstances(task_req, task_allocation); + + NodeResourceInstances old_local_resources = cluster_resources.GetLocalResources(); + ASSERT_EQ(success, true); + std::vector cpu_instances = task_allocation->GetCPUInstances(); + cluster_resources.AddCPUResourceInstances(cpu_instances); + cluster_resources.SubtractCPUResourceInstances(cpu_instances); + + ASSERT_EQ((cluster_resources.GetLocalResources() == old_local_resources), true); + } +} + +TEST_F(SchedulingTest, TaskCPUResourceInstancesTest) { + { + NodeResources node_resources; + vector pred_capacities{4 /* CPU */, 1 /* MEM */, 1 /* GPU */}; + vector cust_ids{1}; + vector cust_capacities{8}; + initNodeResources(node_resources, pred_capacities, cust_ids, cust_capacities); + ClusterResourceScheduler cluster_resources(0, node_resources); + + std::vector allocate_cpu_instances{0.5, 0.5, 0.5, 0.5}; + cluster_resources.SubtractCPUResourceInstances(allocate_cpu_instances); + std::vector available_cpu_instances = cluster_resources.GetLocalResources() + .GetAvailableResourceInstances() + .GetCPUInstances(); + std::vector expected_available_cpu_instances{0.5, 0.5, 0.5, 0.5}; + ASSERT_TRUE(std::equal(available_cpu_instances.begin(), available_cpu_instances.end(), + expected_available_cpu_instances.begin())); + + cluster_resources.AddCPUResourceInstances(allocate_cpu_instances); + available_cpu_instances = cluster_resources.GetLocalResources() + .GetAvailableResourceInstances() + .GetCPUInstances(); + expected_available_cpu_instances = {1., 1., 1., 1.}; + ASSERT_TRUE(std::equal(available_cpu_instances.begin(), available_cpu_instances.end(), + expected_available_cpu_instances.begin())); + + allocate_cpu_instances = {1.5, 1.5, .5, 1.5}; + std::vector underflow = + cluster_resources.SubtractCPUResourceInstances(allocate_cpu_instances); + std::vector expected_underflow{.5, .5, 0., .5}; + ASSERT_TRUE( + std::equal(underflow.begin(), underflow.end(), expected_underflow.begin())); + available_cpu_instances = cluster_resources.GetLocalResources() + .GetAvailableResourceInstances() + .GetCPUInstances(); + expected_available_cpu_instances = {0., 0., 0.5, 0.}; + ASSERT_TRUE(std::equal(available_cpu_instances.begin(), available_cpu_instances.end(), + expected_available_cpu_instances.begin())); + + allocate_cpu_instances = {1.0, .5, 1., .5}; + std::vector overflow = + cluster_resources.AddCPUResourceInstances(allocate_cpu_instances); + std::vector expected_overflow{.0, .0, .5, 0.}; + ASSERT_TRUE(std::equal(overflow.begin(), overflow.end(), expected_overflow.begin())); + available_cpu_instances = cluster_resources.GetLocalResources() + .GetAvailableResourceInstances() + .GetCPUInstances(); + expected_available_cpu_instances = {1., .5, 1., .5}; + ASSERT_TRUE(std::equal(available_cpu_instances.begin(), available_cpu_instances.end(), + expected_available_cpu_instances.begin())); + } +} + +TEST_F(SchedulingTest, UpdateLocalAvailableResourcesFromResourceInstancesTest) { + { + NodeResources node_resources; + vector pred_capacities{4 /* CPU */, 1 /* MEM */, 1 /* GPU */}; + vector cust_ids{1}; + vector cust_capacities{8}; + initNodeResources(node_resources, pred_capacities, cust_ids, cust_capacities); + ClusterResourceScheduler cluster_resources(0, node_resources); + + { + std::vector allocate_cpu_instances{0.5, 0.5, 2, 0.5}; + // SubtractCPUResourceInstances() calls + // UpdateLocalAvailableResourcesFromResourceInstances() under the hood. + cluster_resources.SubtractCPUResourceInstances(allocate_cpu_instances); + std::vector available_cpu_instances = cluster_resources.GetLocalResources() + .GetAvailableResourceInstances() + .GetCPUInstances(); + std::vector expected_available_cpu_instances{0.5, 0.5, 0., 0.5}; + ASSERT_TRUE(std::equal(available_cpu_instances.begin(), + available_cpu_instances.end(), + expected_available_cpu_instances.begin())); + + NodeResources nr; + cluster_resources.GetNodeResources(0, &nr); + ASSERT_TRUE(nr.predefined_resources[0].available == 1.5); + } + + { + std::vector allocate_cpu_instances{1.5, 0.5, 2, 0.3}; + // SubtractCPUResourceInstances() calls + // UpdateLocalAvailableResourcesFromResourceInstances() under the hood. + cluster_resources.AddCPUResourceInstances(allocate_cpu_instances); + std::vector available_cpu_instances = cluster_resources.GetLocalResources() + .GetAvailableResourceInstances() + .GetCPUInstances(); + std::vector expected_available_cpu_instances{1., 1., 1., 0.8}; + ASSERT_TRUE(std::equal(available_cpu_instances.begin(), + available_cpu_instances.end(), + expected_available_cpu_instances.begin())); + + NodeResources nr; + cluster_resources.GetNodeResources(0, &nr); + ASSERT_TRUE(nr.predefined_resources[0].available == 3.8); + } + } +} + #ifdef UNORDERED_VS_ABSL_MAPS_EVALUATION TEST_F(SchedulingTest, SchedulingMapPerformanceTest) { size_t map_len = 1000000; diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 9a9e46cbd..f48d159c4 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -84,6 +84,41 @@ namespace ray { namespace raylet { +// A helper function to print the leased workers. +std::string LeasedWorkersSring( + const std::unordered_map> &leased_workers) { + std::stringstream buffer; + buffer << " @leased_workers: ("; + for (const auto &pair : leased_workers) { + auto &worker = pair.second; + buffer << worker->WorkerId() << ", "; + } + buffer << ")"; + return buffer.str(); +} + +// A helper function to print the workers in worker_pool_. +std::string WorkerPoolString(const std::vector> &worker_pool) { + std::stringstream buffer; + buffer << " @worker_pool: ("; + for (const auto &worker : worker_pool) { + buffer << worker->WorkerId() << ", "; + } + buffer << ")"; + return buffer.str(); +} + +// Helper function to print the worker's owner worker and and node owner. +std::string WorkerOwnerString(std::shared_ptr &worker) { + std::stringstream buffer; + const auto owner_worker_id = + WorkerID::FromBinary(worker->GetOwnerAddress().worker_id()); + const auto owner_node_id = WorkerID::FromBinary(worker->GetOwnerAddress().raylet_id()); + buffer << "leased_worker Lease " << worker->WorkerId() << " owned by " + << owner_worker_id << " / " << owner_node_id; + return buffer.str(); +} + NodeManager::NodeManager(boost::asio::io_service &io_service, const ClientID &self_node_id, const NodeManagerConfig &config, ObjectManager &object_manager, @@ -667,7 +702,6 @@ void NodeManager::ResourceDeleted(const ClientID &client_id, new_resource_scheduler_->DeleteResource(client_id.Binary(), resource_label); } } - RAY_LOG(DEBUG) << "[ResourceDeleted] Updated cluster_resource_map."; return; } @@ -701,7 +735,6 @@ void NodeManager::HeartbeatAdded(const ClientID &client_id, << client_id; return; } - // Trigger local GC at the next heartbeat interval. if (heartbeat_data.should_global_gc()) { should_local_gc_ = true; @@ -876,6 +909,7 @@ void NodeManager::DispatchTasks( // one class of tasks become stuck behind others in the queue, causing Ray to start // many workers. See #3644 for a more detailed description of this issue. std::vector> *> fair_order; + RAY_CHECK(new_scheduler_enabled_ == false); for (auto &it : tasks_by_class) { fair_order.emplace_back(&it); } @@ -932,6 +966,7 @@ void NodeManager::ProcessClientMessage(const std::shared_ptr & << (registered_worker ? std::to_string(registered_worker->GetProcess().GetId()) : "nil"); + if (registered_worker && registered_worker->IsDead()) { // For a worker that is marked as dead (because the job has died already), // all the messages are ignored except DisconnectClient. @@ -1046,8 +1081,6 @@ void NodeManager::ProcessRegisterClientRequestMessage( static_cast(protocol::MessageType::RegisterClientReply), fbb.GetSize(), fbb.GetBufferPointer(), [this, client](const ray::Status &status) { if (!status.ok()) { - RAY_LOG(WARNING) - << "Failed to send RegisterClientReply to client, so disconnecting"; ProcessDisconnectClientMessage(client); } }); @@ -1145,6 +1178,7 @@ void NodeManager::HandleWorkerAvailable(const std::shared_ptr void NodeManager::HandleWorkerAvailable(const std::shared_ptr &worker) { RAY_CHECK(worker); bool worker_idle = true; + // If the worker was assigned a task, mark it as finished. if (!worker->GetAssignedTaskId().IsNil()) { worker_idle = FinishAssignedTask(*worker); @@ -1155,10 +1189,10 @@ void NodeManager::HandleWorkerAvailable(const std::shared_ptr &worker) { worker_pool_.PushWorker(worker); } + // Local resource availability changed: invoke scheduling policy for local node. if (new_scheduler_enabled_) { - DispatchScheduledTasksToWorkers(); + NewSchedulerSchedulePendingTasks(); } else { - // Local resource availability changed: invoke scheduling policy for local node. cluster_resource_map_[self_node_id_].SetLoadResources( local_queues_.GetResourceLoad()); // Call task dispatch to assign work to the new worker. @@ -1181,10 +1215,10 @@ void NodeManager::ProcessDisconnectClientMessage( } else { RAY_LOG(INFO) << "Ignoring client disconnect because the client has already " << "been disconnected."; + return; } } RAY_CHECK(!(is_worker && is_driver)); - // If the client has any blocked tasks, mark them as unblocked. In // particular, we are no longer waiting for their dependencies. if (worker) { @@ -1203,6 +1237,7 @@ void NodeManager::ProcessDisconnectClientMessage( // Clean up any open ray.wait calls that the worker made. task_dependency_manager_.UnsubscribeWaitDependencies(worker->WorkerId()); } + // Erase any lease metadata. leased_workers_.erase(worker->WorkerId()); @@ -1260,24 +1295,34 @@ void NodeManager::ProcessDisconnectClientMessage( worker_pool_.DisconnectWorker(worker); // Return the resources that were being used by this worker. - auto const &task_resources = worker->GetTaskResourceIds(); - local_available_resources_.ReleaseConstrained( - task_resources, cluster_resource_map_[self_node_id_].GetTotalResources()); - cluster_resource_map_[self_node_id_].Release(task_resources.ToResourceSet()); - worker->ResetTaskResourceIds(); + if (new_scheduler_enabled_) { + new_resource_scheduler_->SubtractCPUResourceInstances( + worker->GetBorrowedCPUInstances()); + new_resource_scheduler_->FreeLocalTaskResources(worker->GetAllocatedInstances()); + worker->ClearAllocatedInstances(); + new_resource_scheduler_->FreeLocalTaskResources( + worker->GetLifetimeAllocatedInstances()); + worker->ClearLifetimeAllocatedInstances(); + } else { + auto const &task_resources = worker->GetTaskResourceIds(); + local_available_resources_.ReleaseConstrained( + task_resources, cluster_resource_map_[self_node_id_].GetTotalResources()); + cluster_resource_map_[self_node_id_].Release(task_resources.ToResourceSet()); + worker->ResetTaskResourceIds(); - auto const &lifetime_resources = worker->GetLifetimeResourceIds(); - local_available_resources_.ReleaseConstrained( - lifetime_resources, cluster_resource_map_[self_node_id_].GetTotalResources()); - cluster_resource_map_[self_node_id_].Release(lifetime_resources.ToResourceSet()); - worker->ResetLifetimeResourceIds(); + auto const &lifetime_resources = worker->GetLifetimeResourceIds(); + local_available_resources_.ReleaseConstrained( + lifetime_resources, cluster_resource_map_[self_node_id_].GetTotalResources()); + cluster_resource_map_[self_node_id_].Release(lifetime_resources.ToResourceSet()); + worker->ResetLifetimeResourceIds(); + } - RAY_LOG(DEBUG) << "Worker (pid=" << worker->GetProcess().GetId() - << ") is disconnected. " - << "job_id: " << worker->GetAssignedJobId(); - - // Since some resources may have been released, we can try to dispatch more tasks. - DispatchTasks(local_queues_.GetReadyTasksByClass()); + // Since some resources may have been released, we can try to dispatch more tasks. YYY + if (new_scheduler_enabled_) { + NewSchedulerSchedulePendingTasks(); + } else { + DispatchTasks(local_queues_.GetReadyTasksByClass()); + } } else if (is_driver) { // The client is a driver. const auto job_id = worker->GetAssignedJobId(); @@ -1289,7 +1334,7 @@ void NodeManager::ProcessDisconnectClientMessage( RAY_LOG(DEBUG) << "Driver (pid=" << worker->GetProcess().GetId() << ") is disconnected. " - << "job_id: " << job_id; + << "job_id: " << worker->GetAssignedJobId(); } client->Close(); @@ -1375,9 +1420,6 @@ void NodeManager::ProcessWaitRequestMessage( } } else { // We failed to write to the client, so disconnect the client. - RAY_LOG(WARNING) - << "Failed to send WaitReply to client, so disconnecting client"; - // We failed to send the reply to the client, so disconnect the worker. ProcessDisconnectClientMessage(client); } }); @@ -1406,7 +1448,7 @@ void NodeManager::ProcessWaitForDirectActorCallArgsRequestMessage( [this, client, tag](std::vector found, std::vector remaining) { RAY_CHECK(remaining.empty()); std::shared_ptr worker = worker_pool_.GetRegisteredWorker(client); - if (worker == nullptr) { + if (!worker) { RAY_LOG(ERROR) << "Lost worker for wait request " << client; } else { worker->DirectActorCallArgWaitComplete(tag); @@ -1492,38 +1534,67 @@ void NodeManager::ProcessSubmitTaskMessage(const uint8_t *message_data) { void NodeManager::DispatchScheduledTasksToWorkers() { RAY_CHECK(new_scheduler_enabled_); - while (!tasks_to_dispatch_.empty()) { + + // Check every task in task_to_dispatch queue to see + // whether it can be dispatched and ran. This avoids head-of-line + // blocking where a task which cannot be dispatched because + // there are not enough available resources blocks other + // tasks from being dispatched. + for (size_t queue_size = tasks_to_dispatch_.size(); queue_size > 0; queue_size--) { auto task = tasks_to_dispatch_.front(); auto reply = task.first; auto spec = task.second.GetTaskSpecification(); + tasks_to_dispatch_.pop_front(); + std::shared_ptr worker = worker_pool_.PopWorker(spec); - if (worker == nullptr) { + if (!worker) { + // No worker available to schedule this task. + // Put the task back in the dispatch queue. + tasks_to_dispatch_.push_front(task); return; } - bool schedulable = new_resource_scheduler_->SubtractNodeAvailableResources( - self_node_id_.Binary(), spec.GetRequiredResources().GetResourceMap()); + std::shared_ptr allocated_instances( + new TaskResourceInstances()); + bool schedulable = new_resource_scheduler_->AllocateLocalTaskResources( + spec.GetRequiredResources().GetResourceMap(), allocated_instances); if (!schedulable) { - return; + // Not enough resources to schedule this task. + // Put it back at the end of the dispatch queue. + tasks_to_dispatch_.push_back(task); + worker_pool_.PushWorker(worker); + // Try next task in the dispatch queue. + continue; } - // Handle the allocation to specific resource IDs. - auto acquired_resources = - local_available_resources_.Acquire(spec.GetRequiredResources()); - cluster_resource_map_[self_node_id_].Acquire(spec.GetRequiredResources()); + worker->SetOwnerAddress(spec.CallerAddress()); if (spec.IsActorCreationTask()) { - worker->SetLifetimeResourceIds(acquired_resources); + worker->SetLifetimeAllocatedInstances(allocated_instances); } else { - worker->SetTaskResourceIds(acquired_resources); + worker->SetAllocatedInstances(allocated_instances); } + worker->AssignTaskId(spec.TaskId()); + worker->AssignJobId(spec.JobId()); + worker->SetAssignedTask(task.second); reply(worker, ClientID::Nil(), "", -1); - tasks_to_dispatch_.pop_front(); } } void NodeManager::NewSchedulerSchedulePendingTasks() { RAY_CHECK(new_scheduler_enabled_); - while (!tasks_to_schedule_.empty()) { + size_t queue_size = tasks_to_schedule_.size(); + + // Check every task in task_to_schedule queue to see + // whether it can be scheduled. This avoids head-of-line + // blocking where a task which cannot be scheduled because + // there are not enough available resources blocks other + // tasks from being scheduled. + while (queue_size > 0) { + if (queue_size == 0) { + return; + } else { + queue_size--; + } auto work = tasks_to_schedule_.front(); auto task = work.second; auto request_resources = @@ -1533,13 +1604,16 @@ void NodeManager::NewSchedulerSchedulePendingTasks() { new_resource_scheduler_->GetBestSchedulableNode(request_resources, &violations); if (node_id_string.empty()) { /// There is no node that has available resources to run the request. - break; + tasks_to_schedule_.pop_front(); + tasks_to_schedule_.push_back(work); + continue; } else { if (node_id_string == self_node_id_.Binary()) { WaitForTaskArgsRequests(work); } else { - new_resource_scheduler_->SubtractNodeAvailableResources(node_id_string, - request_resources); + new_resource_scheduler_->AllocateRemoteTaskResources(node_id_string, + request_resources); + ClientID node_id = ClientID::FromBinary(node_id_string); auto node_info_opt = gcs_client_->Nodes().Get(node_id); RAY_CHECK(node_info_opt) @@ -1556,17 +1630,19 @@ void NodeManager::NewSchedulerSchedulePendingTasks() { void NodeManager::WaitForTaskArgsRequests(std::pair &work) { RAY_CHECK(new_scheduler_enabled_); - std::vector object_ids = work.second.GetTaskSpecification().GetDependencies(); + const Task &task = work.second; + std::vector object_ids = task.GetTaskSpecification().GetDependencies(); if (object_ids.size() > 0) { - ray::Status status = object_manager_.Wait( - object_ids, -1, object_ids.size(), false, - [this, work](std::vector found, std::vector remaining) { - RAY_CHECK(remaining.empty()); - tasks_to_dispatch_.push_back(work); - DispatchScheduledTasksToWorkers(); - }); - RAY_CHECK_OK(status); + bool args_ready = task_dependency_manager_.SubscribeGetDependencies( + task.GetTaskSpecification().TaskId(), task.GetDependencies()); + if (args_ready) { + task_dependency_manager_.UnsubscribeGetDependencies( + task.GetTaskSpecification().TaskId()); + tasks_to_dispatch_.push_back(work); + } else { + waiting_tasks_[task.GetTaskSpecification().TaskId()] = work; + } } else { tasks_to_dispatch_.push_back(work); } @@ -1580,6 +1656,7 @@ void NodeManager::HandleRequestWorkerLease(const rpc::RequestWorkerLeaseRequest Task task(task_message); bool is_actor_creation_task = task.GetTaskSpecification().IsActorCreationTask(); ActorID actor_id = ActorID::Nil(); + if (is_actor_creation_task) { actor_id = task.GetTaskSpecification().ActorCreationId(); @@ -1592,11 +1669,11 @@ void NodeManager::HandleRequestWorkerLease(const rpc::RequestWorkerLeaseRequest } if (new_scheduler_enabled_) { - auto request_resources = task.GetTaskSpecification().GetRequiredResources(); + auto task_spec = task.GetTaskSpecification(); auto work = std::make_pair( - [this, request_resources, reply, send_reply_callback]( - std::shared_ptr worker, ClientID spillback_to, std::string address, - int port) { + [this, task_spec, reply, send_reply_callback](std::shared_ptr worker, + ClientID spillback_to, + std::string address, int port) { if (worker != nullptr) { reply->mutable_worker_address()->set_ip_address( initial_config_.node_manager_address); @@ -1605,7 +1682,52 @@ void NodeManager::HandleRequestWorkerLease(const rpc::RequestWorkerLeaseRequest reply->mutable_worker_address()->set_raylet_id(self_node_id_.Binary()); RAY_CHECK(leased_workers_.find(worker->WorkerId()) == leased_workers_.end()); leased_workers_[worker->WorkerId()] = worker; - leased_worker_resources_[worker->WorkerId()] = request_resources; +// TODO (Ion): Fix handling floating point errors, maybe by moving to integers. +#define ZERO_CAPACITY 1.0e-5 + std::shared_ptr allocated_resources; + if (task_spec.IsActorCreationTask()) { + allocated_resources = worker->GetLifetimeAllocatedInstances(); + } else { + allocated_resources = worker->GetAllocatedInstances(); + } + auto predefined_resources = allocated_resources->predefined_resources; + ::ray::rpc::ResourceMapEntry *resource; + for (size_t res_idx = 0; res_idx < predefined_resources.size(); res_idx++) { + bool first = true; // Set resource name only if at least one of its + // instances has available capacity. + for (size_t inst_idx = 0; inst_idx < predefined_resources[res_idx].size(); + inst_idx++) { + if (std::abs(predefined_resources[res_idx][inst_idx]) > ZERO_CAPACITY) { + if (first) { + resource = reply->add_resource_mapping(); + resource->set_name( + new_resource_scheduler_->GetResourceNameFromIndex(res_idx)); + first = false; + } + auto rid = resource->add_resource_ids(); + rid->set_index(inst_idx); + rid->set_quantity(predefined_resources[res_idx][inst_idx]); + } + } + } + auto custom_resources = allocated_resources->custom_resources; + for (auto it = custom_resources.begin(); it != custom_resources.end(); ++it) { + bool first = true; // Set resource name only if at least one of its + // instances has available capacity. + for (size_t inst_idx = 0; inst_idx < it->second.size(); inst_idx++) { + if (std::abs(it->second[inst_idx]) > ZERO_CAPACITY) { + if (first) { + resource = reply->add_resource_mapping(); + resource->set_name( + new_resource_scheduler_->GetResourceNameFromIndex(it->first)); + first = false; + } + auto rid = resource->add_resource_ids(); + rid->set_index(inst_idx); + rid->set_quantity(it->second[inst_idx]); + } + } + } } else { reply->mutable_retry_at_raylet_address()->set_ip_address(address); reply->mutable_retry_at_raylet_address()->set_port(port); @@ -1648,7 +1770,6 @@ void NodeManager::HandleRequestWorkerLease(const rpc::RequestWorkerLeaseRequest } } send_reply_callback(Status::OK(), nullptr, nullptr); - RAY_CHECK(leased_workers_.find(worker_id) == leased_workers_.end()) << "Worker is already leased out " << worker_id; @@ -1672,46 +1793,11 @@ void NodeManager::HandleReturnWorker(const rpc::ReturnWorkerRequest &request, rpc::SendReplyCallback send_reply_callback) { // Read the resource spec submitted by the client. auto worker_id = WorkerID::FromBinary(request.worker_id()); - RAY_LOG(DEBUG) << "Return worker " << worker_id; std::shared_ptr worker = leased_workers_[worker_id]; - if (new_scheduler_enabled_) { - if (worker->IsBlocked()) { - // If worker blocked, unblock it to return the cpu resources back to the worker. - HandleDirectCallTaskUnblocked(worker); - } - auto it = leased_worker_resources_.find(worker_id); - RAY_CHECK(it != leased_worker_resources_.end()); - - new_resource_scheduler_->AddNodeAvailableResources(self_node_id_.Binary(), - it->second.GetResourceMap()); - - if (worker->borrowed_cpu_resources_.GetResourceMap().size()) { - // This machine is oversubscribed, so the worker didn't get back cpus when - // unblocked. Thus we need to substract these cpus, as the previous - // "AddNodeAvailableResources" call assumed they were allocated to this worker. - new_resource_scheduler_->SubtractNodeAvailableResources( - self_node_id_.Binary(), worker->borrowed_cpu_resources_.GetResourceMap()); - worker->borrowed_cpu_resources_ = ResourceSet(); - } - leased_worker_resources_.erase(it); - - // Update resource ids. - auto const &task_resources = worker->GetTaskResourceIds(); - local_available_resources_.ReleaseConstrained( - task_resources, cluster_resource_map_[self_node_id_].GetTotalResources()); - cluster_resource_map_[self_node_id_].Release(task_resources.ToResourceSet()); - worker->ResetTaskResourceIds(); - - // TODO (ion): Handle ProcessDisconnectClientMessage() - HandleWorkerAvailable(worker); - leased_workers_.erase(worker_id); - send_reply_callback(Status::OK(), nullptr, nullptr); - return; - } - - leased_workers_.erase(worker_id); Status status; + leased_workers_.erase(worker_id); + if (worker) { if (request.disconnect_worker()) { ProcessDisconnectClientMessage(worker->Connection()); @@ -1721,6 +1807,12 @@ void NodeManager::HandleReturnWorker(const rpc::ReturnWorkerRequest &request, if (worker->IsBlocked()) { HandleDirectCallTaskUnblocked(worker); } + if (new_scheduler_enabled_) { + new_resource_scheduler_->SubtractCPUResourceInstances( + worker->GetBorrowedCPUInstances()); + new_resource_scheduler_->FreeLocalTaskResources(worker->GetAllocatedInstances()); + worker->ClearAllocatedInstances(); + } HandleWorkerAvailable(worker); } } else { @@ -2099,14 +2191,16 @@ void NodeManager::HandleDirectCallTaskBlocked(const std::shared_ptr &wor if (!worker) { return; } - auto const cpu_resource_ids = worker->ReleaseTaskCpuResources(); - local_available_resources_.Release(cpu_resource_ids); - cluster_resource_map_[self_node_id_].Release(cpu_resource_ids.ToResourceSet()); - new_resource_scheduler_->AddNodeAvailableResources( - self_node_id_.Binary(), // A - cpu_resource_ids.ToResourceSet().GetResourceMap()); - - worker->MarkBlocked(); + std::vector cpu_instances; + if (worker->GetAllocatedInstances() != nullptr) { + cpu_instances = worker->GetAllocatedInstances()->GetCPUInstances(); + } + if (cpu_instances.size() > 0) { + std::vector borrowed_cpu_instances = + new_resource_scheduler_->AddCPUResourceInstances(cpu_instances); + worker->SetBorrowedCPUInstances(borrowed_cpu_instances); + worker->MarkBlocked(); + } NewSchedulerSchedulePendingTasks(); return; } @@ -2126,26 +2220,15 @@ void NodeManager::HandleDirectCallTaskUnblocked(const std::shared_ptr &w if (!worker) { return; } - auto it = leased_worker_resources_.find(worker->WorkerId()); - RAY_CHECK(it != leased_worker_resources_.end()); - const auto cpu_resources = it->second.GetNumCpus(); - bool oversubscribed = !local_available_resources_.Contains(cpu_resources); - if (!oversubscribed) { - // Reacquire the CPU resources for the worker. Note that care needs to be - // taken if the user is using the specific CPU IDs since the IDs that we - // reacquire here may be different from the ones that the task started with. - auto const resource_ids = local_available_resources_.Acquire(cpu_resources); - worker->AcquireTaskCpuResources(resource_ids); - cluster_resource_map_[self_node_id_].Acquire(cpu_resources); - new_resource_scheduler_->SubtractNodeAvailableResources( - self_node_id_.Binary(), cpu_resources.GetResourceMap()); - worker->borrowed_cpu_resources_ = ResourceSet(); - } else { - // Remember these are borrowed cpus resources, i.e., we did not return then to the - // worker. - worker->borrowed_cpu_resources_ = cpu_resources; + std::vector cpu_instances; + if (worker->GetAllocatedInstances() != nullptr) { + cpu_instances = worker->GetAllocatedInstances()->GetCPUInstances(); + } + if (cpu_instances.size() > 0) { + new_resource_scheduler_->SubtractCPUResourceInstances(cpu_instances); + new_resource_scheduler_->AddCPUResourceInstances(worker->GetBorrowedCPUInstances()); + worker->MarkUnblocked(); } - worker->MarkUnblocked(); NewSchedulerSchedulePendingTasks(); return; } @@ -2412,18 +2495,29 @@ bool NodeManager::FinishAssignedTask(Worker &worker) { TaskID task_id = worker.GetAssignedTaskId(); RAY_LOG(DEBUG) << "Finished task " << task_id; - // (See design_docs/task_states.rst for the state transition diagram.) Task task; - RAY_CHECK(local_queues_.RemoveTask(task_id, &task)); + if (new_scheduler_enabled_) { + task = worker.GetAssignedTask(); + // leased_workers_.erase(worker.WorkerId()); // Maybe RAY_CHECK ??? + if (worker.GetAllocatedInstances() != nullptr) { + new_resource_scheduler_->SubtractCPUResourceInstances( + worker.GetBorrowedCPUInstances()); + new_resource_scheduler_->FreeLocalTaskResources(worker.GetAllocatedInstances()); + worker.ClearAllocatedInstances(); + } + } else { + // (See design_docs/task_states.rst for the state transition diagram.) + RAY_CHECK(local_queues_.RemoveTask(task_id, &task)); - // Release task's resources. The worker's lifetime resources are still held. - auto const &task_resources = worker.GetTaskResourceIds(); - local_available_resources_.ReleaseConstrained( - task_resources, cluster_resource_map_[self_node_id_].GetTotalResources()); - cluster_resource_map_[self_node_id_].Release(task_resources.ToResourceSet()); - worker.ResetTaskResourceIds(); + // Release task's resources. The worker's lifetime resources are still held. + auto const &task_resources = worker.GetTaskResourceIds(); + local_available_resources_.ReleaseConstrained( + task_resources, cluster_resource_map_[self_node_id_].GetTotalResources()); + cluster_resource_map_[self_node_id_].Release(task_resources.ToResourceSet()); + worker.ResetTaskResourceIds(); + } - const auto &spec = task.GetTaskSpecification(); + const auto &spec = task.GetTaskSpecification(); // if ((spec.IsActorCreationTask() || spec.IsActorTask())) { // If this was an actor or actor creation task, handle the actor's new // state. @@ -2756,31 +2850,43 @@ void NodeManager::HandleObjectLocal(const ObjectID &object_id) { << " on " << self_node_id_ << ", " << ready_task_ids.size() << " tasks ready"; // 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()); - - // First filter out the tasks that should not be moved to READY. - local_queues_.FilterState(ready_task_id_set, TaskState::BLOCKED); - local_queues_.FilterState(ready_task_id_set, TaskState::RUNNING); - local_queues_.FilterState(ready_task_id_set, TaskState::DRIVER); - local_queues_.FilterState(ready_task_id_set, TaskState::WAITING_FOR_ACTOR_CREATION); - - // Make sure that the remaining tasks are all WAITING or direct call - // actors. - auto ready_task_id_set_copy = ready_task_id_set; - local_queues_.FilterState(ready_task_id_set_copy, TaskState::WAITING); - // Filter out direct call actors. These are not tracked by the raylet and - // their assigned task ID is the actor ID. - for (const auto &id : ready_task_id_set_copy) { - RAY_CHECK(actor_registry_.count(id.ActorId()) > 0); - ready_task_id_set.erase(id); + if (new_scheduler_enabled_) { + for (auto task_id : ready_task_ids) { + auto it = waiting_tasks_.find(task_id); + if (it != waiting_tasks_.end()) { + task_dependency_manager_.UnsubscribeGetDependencies(task_id); + tasks_to_dispatch_.push_back(it->second); + waiting_tasks_.erase(it); + } } + NewSchedulerSchedulePendingTasks(); + } else { + if (ready_task_ids.size() > 0) { + std::unordered_set ready_task_id_set(ready_task_ids.begin(), + ready_task_ids.end()); - // Queue and dispatch the tasks that are ready to run (i.e., WAITING). - auto ready_tasks = local_queues_.RemoveTasks(ready_task_id_set); - local_queues_.QueueTasks(ready_tasks, TaskState::READY); - DispatchTasks(MakeTasksByClass(ready_tasks)); + // First filter out the tasks that should not be moved to READY. + local_queues_.FilterState(ready_task_id_set, TaskState::BLOCKED); + local_queues_.FilterState(ready_task_id_set, TaskState::RUNNING); + local_queues_.FilterState(ready_task_id_set, TaskState::DRIVER); + local_queues_.FilterState(ready_task_id_set, TaskState::WAITING_FOR_ACTOR_CREATION); + + // Make sure that the remaining tasks are all WAITING or direct call + // actors. + auto ready_task_id_set_copy = ready_task_id_set; + local_queues_.FilterState(ready_task_id_set_copy, TaskState::WAITING); + // Filter out direct call actors. These are not tracked by the raylet and + // their assigned task ID is the actor ID. + for (const auto &id : ready_task_id_set_copy) { + RAY_CHECK(actor_registry_.count(id.ActorId()) > 0); + ready_task_id_set.erase(id); + } + + // Queue and dispatch the tasks that are ready to run (i.e., WAITING). + auto ready_tasks = local_queues_.RemoveTasks(ready_task_id_set); + local_queues_.QueueTasks(ready_tasks, TaskState::READY); + DispatchTasks(MakeTasksByClass(ready_tasks)); + } } } diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index 13c710362..153aafadb 100644 --- a/src/ray/raylet/node_manager.h +++ b/src/ray/raylet/node_manager.h @@ -719,9 +719,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler { /// The new resource scheduler for direct task calls. std::shared_ptr new_resource_scheduler_; - /// Map of leased workers to their current resource usage. - /// TODO(ion): Check whether we can track these resources in the worker. - std::unordered_map leased_worker_resources_; typedef std::function, ClientID spillback_to, std::string address, int port)> @@ -732,6 +729,8 @@ class NodeManager : public rpc::NodeManagerServiceHandler { std::deque> tasks_to_schedule_; /// Queue of lease requests that should be scheduled onto workers. std::deque> tasks_to_dispatch_; + /// Queue tasks waiting for arguments to be transferred locally. + absl::flat_hash_map> waiting_tasks_; /// Cache of gRPC clients to workers (not necessarily running on this node). /// Also includes the number of inflight requests to each worker - when this diff --git a/src/ray/raylet/worker.h b/src/ray/raylet/worker.h index 9700cd9d9..7b6b4e084 100644 --- a/src/ray/raylet/worker.h +++ b/src/ray/raylet/worker.h @@ -19,6 +19,8 @@ #include "ray/common/client_connection.h" #include "ray/common/id.h" +#include "ray/common/scheduling/cluster_resource_scheduler.h" +#include "ray/common/scheduling/scheduling_ids.h" #include "ray/common/task/scheduling_resources.h" #include "ray/common/task/task.h" #include "ray/common/task/task_common.h" @@ -85,11 +87,40 @@ class Worker { void DirectActorCallArgWaitComplete(int64_t tag); void WorkerLeaseGranted(const std::string &address, int port); - /// Cpus borrowed by the worker. This happens when the machine is oversubscribed - /// and the worker does not get back the cpu resources when unblocked. - /// TODO (ion): Add methods to access this variable. - /// TODO (ion): Investigate a more intuitive alternative to track these Cpus. - ResourceSet borrowed_cpu_resources_; + // Setter, geter, and clear methods for allocated_instances_. + void SetAllocatedInstances( + std::shared_ptr &allocated_instances) { + allocated_instances_ = allocated_instances; + }; + + std::shared_ptr GetAllocatedInstances() { + return allocated_instances_; + }; + + void ClearAllocatedInstances() { allocated_instances_ = nullptr; }; + + void SetLifetimeAllocatedInstances( + std::shared_ptr &allocated_instances) { + lifetime_allocated_instances_ = allocated_instances; + }; + + std::shared_ptr GetLifetimeAllocatedInstances() { + return lifetime_allocated_instances_; + }; + + void ClearLifetimeAllocatedInstances() { lifetime_allocated_instances_ = nullptr; }; + + void SetBorrowedCPUInstances(std::vector &cpu_instances) { + borrowed_cpu_instances_ = cpu_instances; + }; + + std::vector &GetBorrowedCPUInstances() { return borrowed_cpu_instances_; }; + + void ClearBorrowedCPUInstances() { return borrowed_cpu_instances_.clear(); }; + + Task &GetAssignedTask() { return assigned_task_; }; + + void SetAssignedTask(Task &assigned_task) { assigned_task_ = assigned_task; }; rpc::CoreWorkerClient *rpc_client() { return rpc_client_.get(); } @@ -134,6 +165,22 @@ class Worker { /// The address of this worker's owner. The owner is the worker that /// currently holds the lease on this worker, if any. rpc::Address owner_address_; + /// The capacity of each resource instance allocated to this worker in order + /// to satisfy the resource requests of the task is currently running. + std::shared_ptr allocated_instances_; + /// The capacity of each resource instance allocated to this worker + /// when running as an actor. + std::shared_ptr lifetime_allocated_instances_; + /// CPUs borrowed by the worker. This happens in the following scenario: + /// 1) Worker A is blocked, so it donates its CPUs back to the node. + /// 2) Other workers are scheduled and are allocated some of the CPUs donated by A. + /// 3) Task A is unblocked, but it cannot get all CPUs back. At this point, + /// the node is oversubscribed. borrowed_cpu_instances_ represents the number + /// of CPUs this node is oversubscribed by. + /// TODO (Ion): Investigate a more intuitive alternative to track these Cpus. + std::vector borrowed_cpu_instances_; + /// Task being assigned to this worker. + Task assigned_task_; }; } // namespace raylet