diff --git a/BUILD.bazel b/BUILD.bazel index ca890a77c..6e773e145 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -871,6 +871,16 @@ cc_test( ], ) +cc_test( + name = "scheduling_resources_test", + srcs = ["src/ray/common/task/scheduling_resources_test.cc"], + copts = COPTS, + deps = [ + "ray_common", + "@com_google_googletest//:gtest_main", + ], +) + cc_test( name = "id_test", srcs = ["src/ray/common/id_test.cc"], @@ -1034,6 +1044,20 @@ cc_test( ], ) +cc_test( + name = "gcs_placement_group_scheduler_test", + srcs = [ + "src/ray/gcs/gcs_server/test/gcs_placement_group_scheduler_test.cc", + "src/ray/gcs/gcs_server/test/gcs_server_test_util.h", + ], + copts = COPTS, + deps = [ + ":gcs_server_lib", + ":gcs_test_util_lib", + "@com_google_googletest//:gtest_main", + ], +) + cc_test( name = "gcs_actor_scheduler_test", srcs = [ diff --git a/src/ray/common/bundle_spec.cc b/src/ray/common/bundle_spec.cc new file mode 100644 index 000000000..44c9a961e --- /dev/null +++ b/src/ray/common/bundle_spec.cc @@ -0,0 +1,59 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "bundle_spec.h" +namespace ray { + +void BundleSpecification::ComputeResources() { + auto unit_resource = MapFromProtobuf(message_->unit_resources()); + + if (unit_resource.empty()) { + // A static nil object is used here to avoid allocating the empty object every time. + unit_resource_ = ResourceSet::Nil(); + } else { + unit_resource_.reset(new ResourceSet(unit_resource)); + } +} + +const ResourceSet &BundleSpecification::GetRequiredResources() const { + return *unit_resource_; +} + +std::pair BundleSpecification::BundleId() const { + if (message_->bundle_id() + .placement_group_id() + .empty() /* e.g., empty proto default */) { + int64_t index = message_->bundle_id().bundle_index(); + return std::make_pair(PlacementGroupID::Nil(), index); + } + int64_t index = message_->bundle_id().bundle_index(); + return std::make_pair( + PlacementGroupID::FromBinary(message_->bundle_id().placement_group_id()), index); +} + +std::string BundleSpecification::BundleIdAsString() const { + int64_t index = message_->bundle_id().bundle_index(); + return PlacementGroupID::FromBinary(message_->bundle_id().placement_group_id()).Hex() + + std::to_string(index); +} + +PlacementGroupID BundleSpecification::PlacementGroupId() const { + return PlacementGroupID::FromBinary(message_->bundle_id().placement_group_id()); +} + +int64_t BundleSpecification::Index() const { + return message_->bundle_id().bundle_index(); +} + +} // namespace ray diff --git a/src/ray/common/bundle_spec.h b/src/ray/common/bundle_spec.h new file mode 100644 index 000000000..1176ce02a --- /dev/null +++ b/src/ray/common/bundle_spec.h @@ -0,0 +1,97 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include + +#include "absl/synchronization/mutex.h" +#include "ray/common/function_descriptor.h" +#include "ray/common/grpc_util.h" +#include "ray/common/id.h" +#include "ray/common/task/scheduling_resources.h" +#include "ray/common/task/task_common.h" + +namespace ray { + +typedef std::function ScheduleBundleCallback; +/// Arguments are the raylet ID to spill back to, the raylet's +/// address and the raylet's port. +typedef std::function SpillbackBundleCallback; + +class BundleSpecification : public MessageWrapper { + public: + /// Construct from a protobuf message object. + /// The input message will be **copied** into this object. + /// + /// \param message The protobuf message. + explicit BundleSpecification(rpc::Bundle message) : MessageWrapper(message) { + ComputeResources(); + } + /// Construct from a protobuf message shared_ptr. + /// + /// \param message The protobuf message. + explicit BundleSpecification(std::shared_ptr message) + : MessageWrapper(message) { + ComputeResources(); + } + // Return the bundle_id + std::pair BundleId() const; + + // Return the bundle_id of string. eg: placement_group_id + index. + std::string BundleIdAsString() const; + + // Return the Placement Group id which the Bundle belong to. + PlacementGroupID PlacementGroupId() const; + + // Return the index of the bundle. + int64_t Index() const; + + /// Return the resources that are to be acquired by this bundle. + /// + /// \return The resources that will be acquired by this bundle. + const ResourceSet &GetRequiredResources() const; + + /// Override dispatch behaviour. + void OnScheduleInstead(const ScheduleBundleCallback &callback) { + on_schedule_ = callback; + } + + /// Override spillback behaviour. + void OnSpillbackInstead(const SpillbackBundleCallback &callback) { + on_spillback_ = callback; + } + + /// Returns the schedule bundle callback, or nullptr. + const ScheduleBundleCallback &OnSchedule() const { return on_schedule_; } + + /// Returns the spillback bundle callback, or nullptr. + const SpillbackBundleCallback &OnSpillback() const { return on_spillback_; } + + private: + void ComputeResources(); + + /// Field storing unit resources. Initialized in constructor. + /// TODO(ekl) consider optimizing the representation of ResourceSet for fast copies + /// instead of keeping shared pointers here. + std::shared_ptr unit_resource_; + + mutable ScheduleBundleCallback on_schedule_ = nullptr; + + mutable SpillbackBundleCallback on_spillback_ = nullptr; +}; + +} // namespace ray diff --git a/src/ray/common/id.cc b/src/ray/common/id.cc index 4ee051545..330225a41 100644 --- a/src/ray/common/id.cc +++ b/src/ray/common/id.cc @@ -336,5 +336,5 @@ ID_OSTREAM_OPERATOR(JobID); ID_OSTREAM_OPERATOR(ActorID); ID_OSTREAM_OPERATOR(TaskID); ID_OSTREAM_OPERATOR(ObjectID); - +ID_OSTREAM_OPERATOR(PlacementGroupID); } // namespace ray diff --git a/src/ray/common/id.h b/src/ray/common/id.h index 00475f3b2..39bc0ecd4 100644 --- a/src/ray/common/id.h +++ b/src/ray/common/id.h @@ -53,7 +53,6 @@ enum class ObjectType : uint8_t { using ObjectIDFlagsType = uint16_t; using ObjectIDIndexType = uint32_t; - // Declaration. uint64_t MurmurHash64A(const void *key, int len, unsigned int seed); @@ -345,6 +344,24 @@ class ObjectID : public BaseID { uint8_t id_[kLength]; }; +class PlacementGroupID : public BaseID { + public: + static constexpr size_t kLength = 16; + + /// Size of `PlacementGroupID` in bytes. + /// + /// \return Size of `PlacementGroupID` in bytes. + static size_t Size() { return kLength; } + + /// Constructor of `PlacementGroupID`. + PlacementGroupID() : BaseID() {} + + MSGPACK_DEFINE(id_); + + private: + uint8_t id_[kLength]; +}; + static_assert(sizeof(JobID) == JobID::kLength + sizeof(size_t), "JobID size is not as expected"); static_assert(sizeof(ActorID) == ActorID::kLength + sizeof(size_t), @@ -353,12 +370,15 @@ static_assert(sizeof(TaskID) == TaskID::kLength + sizeof(size_t), "TaskID size is not as expected"); static_assert(sizeof(ObjectID) == ObjectID::kLength + sizeof(size_t), "ObjectID size is not as expected"); +static_assert(sizeof(PlacementGroupID) == PlacementGroupID::kLength + sizeof(size_t), + "PlacementGroupID size is not as expected"); std::ostream &operator<<(std::ostream &os, const UniqueID &id); std::ostream &operator<<(std::ostream &os, const JobID &id); std::ostream &operator<<(std::ostream &os, const ActorID &id); std::ostream &operator<<(std::ostream &os, const TaskID &id); std::ostream &operator<<(std::ostream &os, const ObjectID &id); +std::ostream &operator<<(std::ostream &os, const PlacementGroupID &id); #define DEFINE_UNIQUE_ID(type) \ class RAY_EXPORT type : public UniqueID { \ @@ -489,6 +509,7 @@ DEFINE_UNIQUE_ID(JobID); DEFINE_UNIQUE_ID(ActorID); DEFINE_UNIQUE_ID(TaskID); DEFINE_UNIQUE_ID(ObjectID); +DEFINE_UNIQUE_ID(PlacementGroupID); #include "id_def.h" #undef DEFINE_UNIQUE_ID diff --git a/src/ray/common/placement_group.cc b/src/ray/common/placement_group.cc new file mode 100644 index 000000000..8925aa7fe --- /dev/null +++ b/src/ray/common/placement_group.cc @@ -0,0 +1,51 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "ray/common/placement_group.h" + +namespace ray { +void PlacementGroupSpecification::ConstructBundles() { + for (int i = 0; i < message_->bundles_size(); i++) { + bundles_.push_back(BundleSpecification(message_->bundles(i))); + } +} + +PlacementGroupID PlacementGroupSpecification::PlacementGroupId() const { + if (message_->placement_group_id().empty()) { + return PlacementGroupID::Nil(); + } else { + return PlacementGroupID::FromBinary(message_->placement_group_id()); + } +} + +std::vector PlacementGroupSpecification::GetBundles() const { + return bundles_; +} + +Strategy PlacementGroupSpecification::GetStrategy() const { + if (message_->strategy() == rpc::PlacementStrategy::PACK) { + return Strategy::PACK; + } else { + return Strategy::SPREAD; + } +} + +BundleSpecification PlacementGroupSpecification::GetBundle(int position) const { + return bundles_[position]; +} + +std::string PlacementGroupSpecification::GetName() const { + return std::string(message_->name()); +} +} // namespace ray diff --git a/src/ray/common/placement_group.h b/src/ray/common/placement_group.h new file mode 100644 index 000000000..9ea0220db --- /dev/null +++ b/src/ray/common/placement_group.h @@ -0,0 +1,96 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "ray/common/bundle_spec.h" +#include "ray/common/grpc_util.h" +#include "ray/common/id.h" +#include "ray/protobuf/common.pb.h" + +namespace ray { + +enum class Strategy { + PACK = 0x0, + SPREAD = 0x1, +}; + +class PlacementGroupSpecification : public MessageWrapper { + public: + /// Construct from a protobuf message object. + /// The input message will be **copied** into this object. + /// + /// \param message The protobuf message. + explicit PlacementGroupSpecification(rpc::PlacementGroupSpec message) + : MessageWrapper(message) { + ConstructBundles(); + } + /// Construct from a protobuf message shared_ptr. + /// + /// \param message The protobuf message. + explicit PlacementGroupSpecification(std::shared_ptr message) + : MessageWrapper(message) { + ConstructBundles(); + } + /// Return the placement group id. + PlacementGroupID PlacementGroupId() const; + /// Return the bundles in this placement group. + std::vector GetBundles() const; + /// Return the strategy of the placement group. + Strategy GetStrategy() const; + /// Return the bundle by given index. + BundleSpecification GetBundle(int position) const; + /// Return the name of this placement group. + std::string GetName() const; + + private: + /// Construct bundle vector from protobuf. + void ConstructBundles(); + /// The bundles in this placement group. + std::vector bundles_; +}; + +class PlacementGroupSpecBuilder { + public: + PlacementGroupSpecBuilder() : message_(std::make_shared()) {} + + /// Set the common attributes of the placement group spec. + /// See `common.proto` for meaning of the arguments. + /// + /// \return Reference to the builder object itself. + PlacementGroupSpecBuilder &SetPlacementGroupSpec( + const PlacementGroupID &placement_group_id, std::string name, + const std::vector &bundles, const rpc::PlacementStrategy strategy) { + message_->set_placement_group_id(placement_group_id.Binary()); + message_->set_name(name); + message_->set_strategy(strategy); + for (int i = 0; i < bundles.size(); i++) { + rpc::Bundle bundle = bundles[i]; + auto message_bundle = message_->add_bundles(); + auto mutable_bundle_id = message_bundle->mutable_bundle_id(); + mutable_bundle_id->set_bundle_index(i); + mutable_bundle_id->set_placement_group_id(placement_group_id.Binary()); + message_bundle->mutable_unit_resources()->insert(bundle.unit_resources().begin(), + bundle.unit_resources().end()); + } + return *this; + } + + PlacementGroupSpecification Build() { return PlacementGroupSpecification(message_); } + + private: + std::shared_ptr message_; +}; + +} // namespace ray diff --git a/src/ray/common/task/scheduling_resources.cc b/src/ray/common/task/scheduling_resources.cc index 4d0f4b2ef..fe163eeaa 100644 --- a/src/ray/common/task/scheduling_resources.cc +++ b/src/ray/common/task/scheduling_resources.cc @@ -226,6 +226,30 @@ void ResourceSet::AddResources(const ResourceSet &other) { } } +void ResourceSet::AddBundleResources(const std::string &bundle_id, + const ResourceSet &other) { + for (const auto &resource_pair : other.GetResourceAmountMap()) { + const std::string &resource_label = bundle_id + "_" + resource_pair.first; + const FractionalResourceQuantity &resource_capacity = resource_pair.second; + resource_capacity_[resource_label] += resource_capacity; + } +} + +void ResourceSet::ReturnBundleResources(const std::string &bundle_id) { + for (auto iter = resource_capacity_.begin(); iter != resource_capacity_.end();) { + const std::string &bundle_resource_label = iter->first; + if (bundle_resource_label.find(bundle_id) != std::string::npos) { + const std::string &resource_label = + bundle_resource_label.substr(bundle_resource_label.find_last_of("_") + 1); + const FractionalResourceQuantity &resource_capacity = iter->second; + resource_capacity_[resource_label] += resource_capacity; + iter = resource_capacity_.erase(iter); + } else { + iter++; + } + } +} + FractionalResourceQuantity ResourceSet::GetResource( const std::string &resource_name) const { if (resource_capacity_.count(resource_name) == 0) { @@ -639,6 +663,27 @@ void ResourceIdSet::AddOrUpdateResource(const std::string &resource_name, } } +void ResourceIdSet::AddBundleResource(const std::string &resource_name, + ResourceIds &resource_ids) { + available_resources_[resource_name] = resource_ids; +} + +void ResourceIdSet::CancelResourceReserve(const std::string &resource_name) { + std::string origin_resource_name = resource_name.substr(resource_name.find("_") + 1); + auto iter_orig = available_resources_.find(origin_resource_name); + auto iter_bundle = available_resources_.find(resource_name); + if (iter_bundle == available_resources_.end()) { + return; + } else { + if (iter_orig == available_resources_.end()) { + available_resources_[origin_resource_name] = iter_bundle->second; + } else { + iter_orig->second.Release(iter_bundle->second); + } + available_resources_.erase(iter_bundle); + } +} + void ResourceIdSet::DeleteResource(const std::string &resource_name) { available_resources_.erase(resource_name); } @@ -791,6 +836,19 @@ void SchedulingResources::UpdateResourceCapacity(const std::string &resource_nam } } +void SchedulingResources::UpdateBundleResource(const std::string &bundle_id, + const ResourceSet &resource_set) { + resources_available_.SubtractResourcesStrict(resource_set); + resources_available_.AddBundleResources(bundle_id, resource_set); + resources_total_.SubtractResourcesStrict(resource_set); + resources_total_.AddBundleResources(bundle_id, resource_set); +} + +void SchedulingResources::ReturnBundleResource(const std::string &bundle_id) { + resources_available_.ReturnBundleResources(bundle_id); + resources_total_.ReturnBundleResources(bundle_id); +} + void SchedulingResources::DeleteResource(const std::string &resource_name) { resources_total_.DeleteResource(resource_name); resources_available_.DeleteResource(resource_name); diff --git a/src/ray/common/task/scheduling_resources.h b/src/ray/common/task/scheduling_resources.h index 7f4a0bdcc..b3d1bb998 100644 --- a/src/ray/common/task/scheduling_resources.h +++ b/src/ray/common/task/scheduling_resources.h @@ -147,6 +147,23 @@ class ResourceSet { /// \return Void. void AddResources(const ResourceSet &other); + /// \brief Aggregate resources from the other set into this set, adding any missing + /// resource labels to this set. The resource id will change to bundle_id + "_" + + /// reource_id + /// + /// \param other: The other resource set to add. + /// \param bundle_id: The bundle_id of the bundle. + /// \return Void. + void AddBundleResources(const std::string &bundle_id, const ResourceSet &other); + + /// \brief Return back all the bundle resource. Changing the resource name and adding + /// any missing resource labels to this set. The resource id will remove bundle_id + "_" + /// part. + /// + /// \param bundle_id: The bundle_id of the bundle. + /// \return Void. + void ReturnBundleResources(const std::string &bundle_id); + /// \brief Subtract a set of resources from the current set of resources and /// check that the post-subtraction result nonnegative. Assumes other /// is a subset of the ResourceSet. Deletes any resource if the capacity after @@ -403,6 +420,16 @@ class ResourceIdSet { /// \param capacity capacity of the resource being added void AddOrUpdateResource(const std::string &resource_name, int64_t capacity); + /// \brief Add a Bundle resource in the ResourceIdSet. + /// + /// \param resource_name the name of the resource to create/update + /// \param resource_ids resource_ids of the resource being added + void AddBundleResource(const std::string &resource_name, ResourceIds &resource_ids); + + /// \brief remove a Bundle resource in the ResourceIdSet. + /// + /// \param resource_name the name of the resource to remove. + void CancelResourceReserve(const std::string &resource_name); /// \brief Deletes a resource in the ResourceIdSet. This does not raise an exception, /// just deletes the resource. Tasks with acquired resources keep running. /// @@ -524,6 +551,18 @@ class SchedulingResources { /// \return Void. void UpdateResourceCapacity(const std::string &resource_name, int64_t capacity); + /// \brief Update total, available and load resources with the ResourceIds. + /// Create if not exists. + /// \param resource_name: Name of the resource to be modified + /// \param resource_set: New resource_set of the resource. + void UpdateBundleResource(const std::string &bundle_id, + const ResourceSet &resource_set); + + /// \brief delete total, available and load resources with the ResourceIds. + /// Create if not exists. + /// \param resource_name: Name of the resource to be deleted + void ReturnBundleResource(const std::string &bundle_id); + /// \brief Delete resource from total, available and load resources. /// /// \param resource_name: Name of the resource to be deleted. diff --git a/src/ray/common/task/scheduling_resources_test.cc b/src/ray/common/task/scheduling_resources_test.cc new file mode 100644 index 000000000..6442d459f --- /dev/null +++ b/src/ray/common/task/scheduling_resources_test.cc @@ -0,0 +1,76 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include "gtest/gtest.h" +#include "ray/common/id.h" +#include "scheduling_resources.h" + +namespace ray { +class SchedulingResourcesTest : public ::testing::Test { + public: + void SetUp() override { + resource_set = std::make_shared(); + resource_id_set = std::make_shared(); + } + + protected: + std::shared_ptr resource_set; + std::shared_ptr resource_id_set; +}; + +TEST_F(SchedulingResourcesTest, AddBundleResources) { + UniqueID bundle_id = UniqueID::FromRandom(); + std::vector resource_labels = {"CPU"}; + std::vector resource_capacity = {1.0}; + ResourceSet resource(resource_labels, resource_capacity); + resource_set->AddBundleResources(bundle_id.Binary(), resource); + resource_labels.pop_back(); + resource_labels.push_back(bundle_id.Binary() + "_" + "CPU"); + ResourceSet result_resource(resource_labels, resource_capacity); + ASSERT_EQ(1, resource_set->IsEqual(result_resource)); +} + +TEST_F(SchedulingResourcesTest, AddBundleResource) { + UniqueID bundle_id = UniqueID::FromRandom(); + std::string name = bundle_id.Binary() + "_" + "CPU"; + std::vector whole_ids = {1, 2, 3}; + ResourceIds resource_ids(whole_ids); + resource_id_set->AddBundleResource(name, resource_ids); + ASSERT_EQ(1, resource_id_set->AvailableResources().size()); + ASSERT_EQ(name, resource_id_set->AvailableResources().begin()->first); +} + +TEST_F(SchedulingResourcesTest, ReturnBundleResources) { + UniqueID bundle_id = UniqueID::FromRandom(); + std::vector resource_labels = {"CPU"}; + std::vector resource_capacity = {1.0}; + ResourceSet resource(resource_labels, resource_capacity); + resource_set->AddBundleResources(bundle_id.Binary(), resource); + resource_labels.pop_back(); + resource_labels.push_back(bundle_id.Binary() + "_" + "CPU"); + ResourceSet result_resource(resource_labels, resource_capacity); + ASSERT_EQ(1, resource_set->IsEqual(result_resource)); + + resource_set->ReturnBundleResources(bundle_id.Binary()); + ASSERT_EQ(1, resource_set->IsEqual(resource)); +} + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/src/ray/gcs/gcs_client/service_based_accessor.cc b/src/ray/gcs/gcs_client/service_based_accessor.cc index 651b3cecd..6f0efebca 100644 --- a/src/ray/gcs/gcs_client/service_based_accessor.cc +++ b/src/ray/gcs/gcs_client/service_based_accessor.cc @@ -13,7 +13,6 @@ // limitations under the License. #include "ray/gcs/gcs_client/service_based_accessor.h" - #include "ray/gcs/gcs_client/service_based_gcs_client.h" namespace ray { diff --git a/src/ray/gcs/gcs_server/gcs_placement_group_manager.cc b/src/ray/gcs/gcs_server/gcs_placement_group_manager.cc new file mode 100644 index 000000000..dc25c6a20 --- /dev/null +++ b/src/ray/gcs/gcs_server/gcs_placement_group_manager.cc @@ -0,0 +1,58 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "ray/gcs/gcs_server/gcs_placement_group_manager.h" +#include +#include +#include + +namespace ray { +namespace gcs { + +void GcsPlacementGroup::UpdateState( + rpc::PlacementGroupTableData::PlacementGroupState state) { + placement_group_table_data_.set_state(state); +} + +rpc::PlacementGroupTableData::PlacementGroupState GcsPlacementGroup::GetState() const { + return placement_group_table_data_.state(); +} + +PlacementGroupID GcsPlacementGroup::GetPlacementGroupID() const { + return PlacementGroupID::FromBinary(placement_group_table_data_.placement_group_id()); +} + +std::string GcsPlacementGroup::GetName() const { + return placement_group_table_data_.name(); +} + +std::vector> GcsPlacementGroup::GetBundles() const { + auto bundles = placement_group_table_data_.bundles(); + std::vector> ret_bundles; + for (auto iter = bundles.begin(); iter != bundles.end(); iter++) { + ret_bundles.push_back(std::make_shared(*iter)); + } + return ret_bundles; +} + +rpc::PlacementStrategy GcsPlacementGroup::GetStrategy() const { + return placement_group_table_data_.strategy(); +} + +const rpc::PlacementGroupTableData &GcsPlacementGroup::GetPlacementGroupTableData() { + return placement_group_table_data_; +} + +} // namespace gcs +} // namespace ray diff --git a/src/ray/gcs/gcs_server/gcs_placement_group_manager.h b/src/ray/gcs/gcs_server/gcs_placement_group_manager.h new file mode 100644 index 000000000..19aebf15e --- /dev/null +++ b/src/ray/gcs/gcs_server/gcs_placement_group_manager.h @@ -0,0 +1,82 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once +#include +#include +#include +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "gcs_placement_group_scheduler.h" +#include "gcs_table_storage.h" +#include "ray/gcs/pubsub/gcs_pub_sub.h" + +namespace ray { +namespace gcs { + +/// GcsPlacementGroup just wraps `PlacementGroupTableData` and provides some convenient +/// interfaces to access the fields inside `PlacementGroupTableData`. This class is not +/// thread-safe. +class GcsPlacementGroup { + public: + /// Create a GcsPlacementGroup by placement_group_table_data. + /// + /// \param placement_group_table_data Data of the placement_group (see gcs.proto). + explicit GcsPlacementGroup(rpc::PlacementGroupTableData placement_group_table_data) + : placement_group_table_data_(std::move(placement_group_table_data)) {} + + /// Create a GcsPlacementGroup by CreatePlacementGroupRequest. + /// + /// \param request Contains the placement group creation task specification. + explicit GcsPlacementGroup(const ray::rpc::CreatePlacementGroupRequest &request) { + const auto &placement_group_spec = request.placement_group_spec(); + placement_group_table_data_.set_placement_group_id( + placement_group_spec.placement_group_id()); + + placement_group_table_data_.set_state(rpc::PlacementGroupTableData::PENDING); + placement_group_table_data_.mutable_bundles()->CopyFrom( + placement_group_spec.bundles()); + placement_group_table_data_.set_strategy(placement_group_spec.strategy()); + } + + /// Get the immutable PlacementGroupTableData of this placement group. + const rpc::PlacementGroupTableData &GetPlacementGroupTableData(); + + /// Update the state of this placement_group. + void UpdateState(rpc::PlacementGroupTableData::PlacementGroupState state); + /// Get the state of this gcs placement_group. + rpc::PlacementGroupTableData::PlacementGroupState GetState() const; + + /// Get the id of this placement_group. + PlacementGroupID GetPlacementGroupID() const; + /// Get the name of this placement_group. + std::string GetName() const; + + /// Get the bundles of this placement_group + std::vector> GetBundles() const; + + /// Get the Strategy + rpc::PlacementStrategy GetStrategy() const; + + private: + /// The placement_group meta data which contains the task specification as well as the + /// state of the gcs placement_group and so on (see gcs.proto). + rpc::PlacementGroupTableData placement_group_table_data_; +}; +} // namespace gcs +} // namespace ray diff --git a/src/ray/gcs/gcs_server/gcs_placement_group_scheduler.cc b/src/ray/gcs/gcs_server/gcs_placement_group_scheduler.cc new file mode 100644 index 000000000..af01ad795 --- /dev/null +++ b/src/ray/gcs/gcs_server/gcs_placement_group_scheduler.cc @@ -0,0 +1,246 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "ray/gcs/gcs_server/gcs_placement_group_scheduler.h" +#include +#include "ray/gcs/gcs_server/gcs_placement_group_manager.h" + +namespace ray { +namespace gcs { + +GcsPlacementGroupScheduler::GcsPlacementGroupScheduler( + boost::asio::io_context &io_context, + std::shared_ptr gcs_table_storage, + const gcs::GcsNodeManager &gcs_node_manager, + ReserveResourceClientFactoryFn lease_client_factory) + : return_timer_(io_context), + gcs_table_storage_(gcs_table_storage), + gcs_node_manager_(gcs_node_manager), + lease_client_factory_(std::move(lease_client_factory)) { + scheduler_strategies_.push_back(std::make_shared()); + scheduler_strategies_.push_back(std::make_shared()); +} + +/// This is an initial algorithm to respect pack algorithm. +/// In this algorithm, we try to pack all the bundle in the first node +/// and don't care the real resource. +ScheduleMap GcsPackStrategy::Schedule( + std::vector> &bundles, + const GcsNodeManager &node_manager) { + ScheduleMap schedule_map; + auto &alive_nodes = node_manager.GetAllAliveNodes(); + for (size_t pos = 0; pos < bundles.size(); pos++) { + schedule_map[bundles[pos]->BundleId()] = + ClientID::FromBinary(alive_nodes.begin()->second->node_id()); + } + return schedule_map; +} + +/// This is an initial algorithm to respect spread algorithm. +/// In this algorithm, we try to spread all the bundle in different node +/// and don't care the real resource. +ScheduleMap GcsSpreadStrategy::Schedule( + std::vector> &bundles, + const GcsNodeManager &node_manager) { + ScheduleMap schedule_map; + auto &alive_nodes = node_manager.GetAllAliveNodes(); + auto iter = alive_nodes.begin(); + size_t index = 0; + size_t alive_nodes_size = alive_nodes.size(); + for (; iter != alive_nodes.end(); iter++, index++) { + for (size_t base = 0;; base++) { + if (index + base * alive_nodes_size >= bundles.size()) { + break; + } else { + schedule_map[bundles[index + base * alive_nodes_size]->BundleId()] = + ClientID::FromBinary(iter->second->node_id()); + } + } + } + return schedule_map; +} + +void GcsPlacementGroupScheduler::Schedule( + std::shared_ptr placement_group, + std::function)> schedule_failure_handler, + std::function)> schedule_success_handler) { + auto bundles = placement_group->GetBundles(); + auto strategy = placement_group->GetStrategy(); + auto alive_nodes = gcs_node_manager_.GetAllAliveNodes(); + /// If the placement group don't have bundle, the placement group creates success. + if (bundles.size() == 0) { + schedule_success_handler(placement_group); + return; + } + + // If alive_node is empty, the the placement group creates fail. + if (alive_nodes.size() == 0) { + schedule_failure_handler(placement_group); + return; + } + auto schedule_map = + scheduler_strategies_[strategy]->Schedule(bundles, gcs_node_manager_); + // If schedule success, the decision will be set as schedule_map[bundles[pos]] + // else will be set ClientID::Nil(). + auto decision = std::shared_ptr>( + new std::unordered_map()); + // To count how many scheduler have been return, which include success and failure. + auto finish_count = std::make_shared(); + /// TODO(AlisaWu): Change the strategy when reserve resource failed. + for (size_t pos = 0; pos < bundles.size(); pos++) { + RAY_CHECK(node_to_bundles_when_leasing_[schedule_map.at(bundles[pos]->BundleId())] + .emplace(bundles[pos]->BundleId()) + .second); + ReserveResourceFromNode( + bundles[pos], + gcs_node_manager_.GetNode(schedule_map.at(bundles[pos]->BundleId())), + [this, pos, bundles, schedule_map, placement_group, decision, finish_count, + schedule_failure_handler, schedule_success_handler]( + const Status &status, const rpc::RequestResourceReserveReply &reply) { + (*finish_count)++; + if (status.ok() && reply.success()) { + (*decision)[bundles[pos]->BundleId()] = + schedule_map.at(bundles[pos]->BundleId()); + } else { + (*decision)[bundles[pos]->BundleId()] = ClientID::Nil(); + } + if ((*finish_count) == bundles.size()) { + bool lease_success = true; + for (size_t i = 0; i < (*finish_count); i++) { + if ((*decision)[bundles[i]->BundleId()] == ClientID::Nil()) { + lease_success = false; + break; + } + } + if (lease_success) { + rpc::ScheduleData data; + for (size_t i = 0; i < bundles.size(); i++) { + data.mutable_schedule_plan()->insert( + {bundles[i]->BundleIdAsString(), + (*decision)[bundles[i]->BundleId()].Binary()}); + } + RAY_CHECK_OK(gcs_table_storage_->PlacementGroupScheduleTable().Put( + placement_group->GetPlacementGroupID(), data, [](Status status) {})); + schedule_success_handler(placement_group); + } else { + for (size_t i = 0; i < (*finish_count); i++) { + if ((*decision)[bundles[i]->BundleId()] != ClientID::Nil()) { + CancelResourceReserve(bundles[i], + gcs_node_manager_.GetNode( + schedule_map.at(bundles[pos]->BundleId()))); + } + } + schedule_failure_handler(placement_group); + } + } + }); + } +} + +void GcsPlacementGroupScheduler::ReserveResourceFromNode( + std::shared_ptr bundle, + std::shared_ptr node, ReserveResourceCallback callback) { + RAY_CHECK(node); + + auto node_id = ClientID::FromBinary(node->node_id()); + RAY_LOG(DEBUG) << "Start leasing resource from node " << node_id << " for bundle " + << bundle->BundleId().first << std::to_string(bundle->BundleId().second); + + rpc::Address remote_address; + remote_address.set_raylet_id(node->node_id()); + remote_address.set_ip_address(node->node_manager_address()); + remote_address.set_port(node->node_manager_port()); + auto lease_client = GetOrConnectLeaseClient(remote_address); + RAY_LOG(DEBUG) << "Start leasing resource from node " << node_id << " for bundle " + << bundle->BundleId().first << bundle->BundleId().second; + auto status = lease_client->RequestResourceReserve( + *bundle, [this, node_id, bundle, node, callback]( + const Status &status, const rpc::RequestResourceReserveReply &reply) { + // TODO(AlisaWu): Add placement group cancel. + auto iter = node_to_bundles_when_leasing_.find(node_id); + if (iter != node_to_bundles_when_leasing_.end()) { + auto bundle_iter = iter->second.find(bundle->BundleId()); + RAY_CHECK(bundle_iter != iter->second.end()); + if (status.ok()) { + if (reply.success()) { + RAY_LOG(DEBUG) << "Finished leasing resource from " << node_id + << " for bundle " << bundle->BundleId().first + << bundle->BundleId().second; + } else { + RAY_LOG(DEBUG) << "Failed leasing resource from " << node_id + << " for bundle " << bundle->BundleId().first + << bundle->BundleId().second; + } + // Remove the actor from the leasing map as the reply is returned from the + // remote node. + iter->second.erase(bundle_iter); + if (iter->second.empty()) { + node_to_bundles_when_leasing_.erase(iter); + } + } + callback(status, reply); + } + }); + if (!status.ok()) { + rpc::RequestResourceReserveReply reply; + reply.set_success(false); + callback(status, reply); + } +} + +void GcsPlacementGroupScheduler::CancelResourceReserve( + std::shared_ptr bundle_spec, + std::shared_ptr node) { + RAY_CHECK(node); + + auto node_id = ClientID::FromBinary(node->node_id()); + RAY_LOG(DEBUG) << "Start returning resource for node " << node_id << " for bundle " + << bundle_spec->BundleId().first << bundle_spec->BundleId().second; + + rpc::Address remote_address; + remote_address.set_raylet_id(node->node_id()); + remote_address.set_ip_address(node->node_manager_address()); + remote_address.set_port(node->node_manager_port()); + auto return_client = GetOrConnectLeaseClient(remote_address); + auto status = return_client->CancelResourceReserve( + *bundle_spec, + [this, bundle_spec, node](const Status &status, + const rpc::CancelResourceReserveReply &reply) { + if (!status.ok()) { + return_timer_.expires_from_now(boost::posix_time::milliseconds(5)); + return_timer_.async_wait( + [this, bundle_spec, node](const boost::system::error_code &error) { + if (error == boost::system::errc::operation_canceled) { + return; + } else { + CancelResourceReserve(bundle_spec, node); + } + }); + } + }); +} // namespace gcs + +std::shared_ptr +GcsPlacementGroupScheduler::GetOrConnectLeaseClient(const rpc::Address &raylet_address) { + auto node_id = ClientID::FromBinary(raylet_address.raylet_id()); + auto iter = remote_lease_clients_.find(node_id); + if (iter == remote_lease_clients_.end()) { + auto lease_client = lease_client_factory_(raylet_address); + iter = remote_lease_clients_.emplace(node_id, std::move(lease_client)).first; + } + return iter->second; +} + +} // namespace gcs +} // namespace ray diff --git a/src/ray/gcs/gcs_server/gcs_placement_group_scheduler.h b/src/ray/gcs/gcs_server/gcs_placement_group_scheduler.h new file mode 100644 index 000000000..7fdcf0582 --- /dev/null +++ b/src/ray/gcs/gcs_server/gcs_placement_group_scheduler.h @@ -0,0 +1,147 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "gcs_node_manager.h" +#include "gcs_table_storage.h" + +namespace ray { +namespace gcs { + +using ReserveResourceClientFactoryFn = + std::function(const rpc::Address &address)>; + +using ReserveResourceCallback = + std::function; + +typedef std::pair BundleID; +struct pair_hash { + template + std::size_t operator()(const std::pair &pair) const { + return std::hash()(pair.first) ^ std::hash()(pair.second); + } +}; +using ScheduleMap = std::unordered_map; +class GcsPlacementGroup; + +class GcsPlacementGroupSchedulerInterface { + public: + /// Schedule the specified placement_group. + /// + /// \param placement_group to be scheduled. + virtual void Schedule( + std::shared_ptr placement_group, + std::function)> schedule_failure_handler, + std::function)> + schedule_success_handler) = 0; + + virtual ~GcsPlacementGroupSchedulerInterface() {} +}; + +class GcsScheduleStrategy { + public: + virtual ScheduleMap Schedule( + std::vector> &bundles, + const GcsNodeManager &node_manager) = 0; +}; + +class GcsPackStrategy : public GcsScheduleStrategy { + public: + ScheduleMap Schedule(std::vector> &bundles, + const GcsNodeManager &node_manager) override; +}; + +class GcsSpreadStrategy : public GcsScheduleStrategy { + public: + ScheduleMap Schedule(std::vector> &bundles, + const GcsNodeManager &node_manager) override; +}; + +/// GcsPlacementGroupScheduler is responsible for scheduling placement_groups registered +/// to GcsPlacementGroupManager. This class is not thread-safe. +class GcsPlacementGroupScheduler : public GcsPlacementGroupSchedulerInterface { + public: + /// Create a GcsPlacementGroupScheduler + /// + /// \param io_context The main event loop. + /// \param placement_group_info_accessor Used to flush placement_group info to storage. + /// \param gcs_node_manager The node manager which is used when scheduling. + GcsPlacementGroupScheduler( + boost::asio::io_context &io_context, + std::shared_ptr gcs_table_storage, + const GcsNodeManager &gcs_node_manager, + ReserveResourceClientFactoryFn lease_client_factory = nullptr); + virtual ~GcsPlacementGroupScheduler() = default; + /// Schedule the specified placement_group. + /// If there is no available nodes then the `schedule_failed_handler` will be + /// triggered, otherwise the bundle in placement_group will be add into a queue and + /// schedule all bundle by calling ReserveResourceFromNode(). + /// + /// \param placement_group to be scheduled. + void Schedule( + std::shared_ptr placement_group, + std::function)> schedule_failure_handler, + std::function)> schedule_success_handler) + override; + + protected: + /// Lease resource from the specified node for the specified bundle. + void ReserveResourceFromNode(std::shared_ptr bundle, + std::shared_ptr node, + ReserveResourceCallback callback); + + /// return resource for the specified node for the specified bundle. + /// + /// \param bundle A description of the bundle to return. + /// \param node The node that the worker will be returned for. + void CancelResourceReserve(std::shared_ptr bundle_spec, + std::shared_ptr node); + + /// Get an existing lease client or connect a new one. + std::shared_ptr GetOrConnectLeaseClient( + const rpc::Address &raylet_address); + /// A timer that ticks every cancel resource failure milliseconds. + boost::asio::deadline_timer return_timer_; + /// Used to update placement group information upon creation, deletion, etc. + std::shared_ptr gcs_table_storage_; + /// Reference of GcsNodeManager. + const GcsNodeManager &gcs_node_manager_; + /// The cached node clients which are used to communicate with raylet to lease workers. + absl::flat_hash_map> + remote_lease_clients_; + /// Factory for producing new clients to request leases from remote nodes. + ReserveResourceClientFactoryFn lease_client_factory_; + + /// Map from node ID to the set of bundles for whom we are trying to acquire a lease + /// from that node. This is needed so that we can retry lease requests from the node + /// until we receive a reply or the node is removed. + absl::flat_hash_map> + node_to_bundles_when_leasing_; + /// A vector to store all the schedule strategy. + std::vector> scheduler_strategies_; +}; + +} // namespace gcs +} // namespace ray diff --git a/src/ray/gcs/gcs_server/gcs_table_storage.cc b/src/ray/gcs/gcs_server/gcs_table_storage.cc index a91fd8c48..956235842 100644 --- a/src/ray/gcs/gcs_server/gcs_table_storage.cc +++ b/src/ray/gcs/gcs_server/gcs_table_storage.cc @@ -128,6 +128,7 @@ template class GcsTableWithJobId; template class GcsTableWithJobId; template class GcsTableWithJobId; template class GcsTableWithJobId; +template class GcsTable; } // namespace gcs } // namespace ray diff --git a/src/ray/gcs/gcs_server/gcs_table_storage.h b/src/ray/gcs/gcs_server/gcs_table_storage.h index ddc9a3b81..573157846 100644 --- a/src/ray/gcs/gcs_server/gcs_table_storage.h +++ b/src/ray/gcs/gcs_server/gcs_table_storage.h @@ -36,6 +36,7 @@ using rpc::ObjectTableDataList; using rpc::ProfileTableData; using rpc::ResourceMap; using rpc::ResourceTableData; +using rpc::ScheduleData; using rpc::StoredConfig; using rpc::TaskLeaseData; using rpc::TaskReconstructionData; @@ -242,6 +243,14 @@ class GcsHeartbeatTable : public GcsTable { } }; +class GcsPlacementGroupScheduleTable : public GcsTable { + public: + explicit GcsPlacementGroupScheduleTable(std::shared_ptr &store_client) + : GcsTable(store_client) { + table_name_ = TablePrefix_Name(TablePrefix::PLACEMENT_GROUP_SCHEDULE); + } +}; + class GcsHeartbeatBatchTable : public GcsTable { public: explicit GcsHeartbeatBatchTable(std::shared_ptr &store_client) @@ -338,6 +347,11 @@ class GcsTableStorage { return *node_resource_table_; } + GcsPlacementGroupScheduleTable &PlacementGroupScheduleTable() { + RAY_CHECK(placement_group_schedule_table_ != nullptr); + return *placement_group_schedule_table_; + } + GcsHeartbeatTable &HeartbeatTable() { RAY_CHECK(heartbeat_table_ != nullptr); return *heartbeat_table_; @@ -380,6 +394,7 @@ class GcsTableStorage { std::unique_ptr object_table_; std::unique_ptr node_table_; std::unique_ptr node_resource_table_; + std::unique_ptr placement_group_schedule_table_; std::unique_ptr heartbeat_table_; std::unique_ptr heartbeat_batch_table_; std::unique_ptr error_info_table_; @@ -405,7 +420,11 @@ class RedisGcsTableStorage : public GcsTableStorage { object_table_.reset(new GcsObjectTable(store_client_)); node_table_.reset(new GcsNodeTable(store_client_)); node_resource_table_.reset(new GcsNodeResourceTable(store_client_)); + placement_group_schedule_table_.reset( + new GcsPlacementGroupScheduleTable(store_client_)); heartbeat_table_.reset(new GcsHeartbeatTable(store_client_)); + placement_group_schedule_table_.reset( + new GcsPlacementGroupScheduleTable(store_client_)); heartbeat_batch_table_.reset(new GcsHeartbeatBatchTable(store_client_)); error_info_table_.reset(new GcsErrorInfoTable(store_client_)); profile_table_.reset(new GcsProfileTable(store_client_)); @@ -431,6 +450,8 @@ class InMemoryGcsTableStorage : public GcsTableStorage { object_table_.reset(new GcsObjectTable(store_client_)); node_table_.reset(new GcsNodeTable(store_client_)); node_resource_table_.reset(new GcsNodeResourceTable(store_client_)); + placement_group_schedule_table_.reset( + new GcsPlacementGroupScheduleTable(store_client_)); heartbeat_table_.reset(new GcsHeartbeatTable(store_client_)); heartbeat_batch_table_.reset(new GcsHeartbeatBatchTable(store_client_)); error_info_table_.reset(new GcsErrorInfoTable(store_client_)); diff --git a/src/ray/gcs/gcs_server/test/gcs_actor_manager_test.cc b/src/ray/gcs/gcs_server/test/gcs_actor_manager_test.cc index 6bd3dac49..e5b17b59f 100644 --- a/src/ray/gcs/gcs_server/test/gcs_actor_manager_test.cc +++ b/src/ray/gcs/gcs_server/test/gcs_actor_manager_test.cc @@ -14,11 +14,11 @@ #include #include -#include "ray/common/test_util.h" #include #include "gtest/gtest.h" +#include "ray/common/test_util.h" namespace ray { diff --git a/src/ray/gcs/gcs_server/test/gcs_placement_group_scheduler_test.cc b/src/ray/gcs/gcs_server/test/gcs_placement_group_scheduler_test.cc new file mode 100644 index 000000000..2b773fa0e --- /dev/null +++ b/src/ray/gcs/gcs_server/test/gcs_placement_group_scheduler_test.cc @@ -0,0 +1,178 @@ +// Copyright 2017 The Ray Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include + +#include +#include "gtest/gtest.h" + +namespace ray { + +class GcsPlacementGroupSchedulerTest : public ::testing::Test { + public: + void SetUp() override { + raylet_client_ = std::make_shared(); + gcs_table_storage_ = std::make_shared(redis_client_); + gcs_pub_sub_ = std::make_shared(redis_client_); + gcs_node_manager_ = std::make_shared( + io_service_, error_info_accessor_, gcs_pub_sub_, gcs_table_storage_); + gcs_table_storage_ = std::make_shared(io_service_); + store_client_ = std::make_shared(io_service_); + gcs_placement_group_scheduler_ = + std::make_shared( + io_service_, gcs_table_storage_, *gcs_node_manager_, + /*lease_client_fplacement_groupy=*/ + [this](const rpc::Address &address) { return raylet_client_; }); + } + + protected: + boost::asio::io_service io_service_; + GcsServerMocker::MockedErrorInfoAccessor error_info_accessor_; + std::shared_ptr store_client_; + + std::shared_ptr raylet_client_; + std::shared_ptr gcs_node_manager_; + std::shared_ptr + gcs_placement_group_scheduler_; + std::vector> success_placement_groups_; + std::vector> failure_placement_groups_; + std::shared_ptr gcs_pub_sub_; + std::shared_ptr gcs_table_storage_; + std::shared_ptr redis_client_; +}; + +TEST_F(GcsPlacementGroupSchedulerTest, TestScheduleFailedWithZeroNode) { + ASSERT_EQ(0, gcs_node_manager_->GetAllAliveNodes().size()); + auto create_placement_group_request = Mocker::GenCreatePlacementGroupRequest(); + auto placement_group = + std::make_shared(create_placement_group_request); + + // Schedule the placement_group with zero node. + gcs_placement_group_scheduler_->Schedule( + placement_group, + [this](std::shared_ptr placement_group) { + failure_placement_groups_.emplace_back(std::move(placement_group)); + }, + [this](std::shared_ptr placement_group) { + success_placement_groups_.emplace_back(std::move(placement_group)); + }); + + // The lease request should not be send and the scheduling of placement_group should + // fail as there are no available nodes. + ASSERT_EQ(raylet_client_->num_lease_requested, 0); + ASSERT_EQ(0, success_placement_groups_.size()); + ASSERT_EQ(1, failure_placement_groups_.size()); + ASSERT_EQ(placement_group, failure_placement_groups_.front()); +} + +TEST_F(GcsPlacementGroupSchedulerTest, TestSchedulePlacementGroupSuccess) { + auto node = Mocker::GenNodeInfo(); + gcs_node_manager_->AddNode(node); + ASSERT_EQ(1, gcs_node_manager_->GetAllAliveNodes().size()); + + auto create_placement_group_request = Mocker::GenCreatePlacementGroupRequest(); + auto placement_group = + std::make_shared(create_placement_group_request); + + // Schedule the placement_group with 1 available node, and the lease request should be + // send to the node. + gcs_placement_group_scheduler_->Schedule( + placement_group, + [this](std::shared_ptr placement_group) { + failure_placement_groups_.emplace_back(std::move(placement_group)); + }, + [this](std::shared_ptr placement_group) { + success_placement_groups_.emplace_back(std::move(placement_group)); + }); + + ASSERT_EQ(2, raylet_client_->num_lease_requested); + ASSERT_EQ(2, raylet_client_->lease_callbacks.size()); + ASSERT_TRUE(raylet_client_->GrantResourceReserve()); + ASSERT_TRUE(raylet_client_->GrantResourceReserve()); + ASSERT_EQ(0, failure_placement_groups_.size()); + ASSERT_EQ(1, success_placement_groups_.size()); + ASSERT_EQ(placement_group, success_placement_groups_.front()); +} + +TEST_F(GcsPlacementGroupSchedulerTest, TestSchedulePlacementGroupFailed) { + auto node = Mocker::GenNodeInfo(); + gcs_node_manager_->AddNode(node); + ASSERT_EQ(1, gcs_node_manager_->GetAllAliveNodes().size()); + + auto create_placement_group_request = Mocker::GenCreatePlacementGroupRequest(); + auto placement_group = + std::make_shared(create_placement_group_request); + + // Schedule the placement_group with 1 available node, and the lease request should be + // send to the node. + gcs_placement_group_scheduler_->Schedule( + placement_group, + [this](std::shared_ptr placement_group) { + failure_placement_groups_.emplace_back(std::move(placement_group)); + }, + [this](std::shared_ptr placement_group) { + success_placement_groups_.emplace_back(std::move(placement_group)); + }); + + ASSERT_EQ(2, raylet_client_->num_lease_requested); + ASSERT_EQ(2, raylet_client_->lease_callbacks.size()); + ASSERT_TRUE(raylet_client_->GrantResourceReserve(false)); + ASSERT_TRUE(raylet_client_->GrantResourceReserve(false)); + // // Reply the placement_group creation request, then the placement_group should be + // scheduled successfully. + ASSERT_EQ(1, failure_placement_groups_.size()); + ASSERT_EQ(0, success_placement_groups_.size()); + ASSERT_EQ(placement_group, failure_placement_groups_.front()); +} + +TEST_F(GcsPlacementGroupSchedulerTest, TestSchedulePlacementGroupReturnResource) { + auto node = Mocker::GenNodeInfo(); + gcs_node_manager_->AddNode(node); + ASSERT_EQ(1, gcs_node_manager_->GetAllAliveNodes().size()); + + auto create_placement_group_request = Mocker::GenCreatePlacementGroupRequest(); + auto placement_group = + std::make_shared(create_placement_group_request); + + // Schedule the placement_group with 1 available node, and the lease request should be + // send to the node. + gcs_placement_group_scheduler_->Schedule( + placement_group, + [this](std::shared_ptr placement_group) { + failure_placement_groups_.emplace_back(std::move(placement_group)); + }, + [this](std::shared_ptr placement_group) { + success_placement_groups_.emplace_back(std::move(placement_group)); + }); + + ASSERT_EQ(2, raylet_client_->num_lease_requested); + ASSERT_EQ(2, raylet_client_->lease_callbacks.size()); + // One bundle success and the other failed. + ASSERT_TRUE(raylet_client_->GrantResourceReserve()); + ASSERT_TRUE(raylet_client_->GrantResourceReserve(false)); + ASSERT_EQ(1, raylet_client_->num_return_requested); + // // Reply the placement_group creation request, then the placement_group should be + // scheduled successfully. + ASSERT_EQ(1, failure_placement_groups_.size()); + ASSERT_EQ(0, success_placement_groups_.size()); + ASSERT_EQ(placement_group, failure_placement_groups_.front()); +} + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/src/ray/gcs/gcs_server/test/gcs_server_test_util.h b/src/ray/gcs/gcs_server/test/gcs_server_test_util.h index c9e0a2c1b..79b44ac86 100644 --- a/src/ray/gcs/gcs_server/test/gcs_server_test_util.h +++ b/src/ray/gcs/gcs_server/test/gcs_server_test_util.h @@ -23,6 +23,8 @@ #include "ray/gcs/gcs_server/gcs_actor_manager.h" #include "ray/gcs/gcs_server/gcs_actor_scheduler.h" #include "ray/gcs/gcs_server/gcs_node_manager.h" +#include "ray/gcs/gcs_server/gcs_placement_group_manager.h" +#include "ray/gcs/gcs_server/gcs_placement_group_scheduler.h" #include "ray/util/asio_util.h" namespace ray { @@ -136,6 +138,49 @@ struct GcsServerMocker { std::list> cancel_callbacks = {}; }; + class MockRayletResourceClient : public ResourceReserveInterface { + public: + ray::Status RequestResourceReserve( + const BundleSpecification &bundle_spec, + const ray::rpc::ClientCallback &callback) + override { + num_lease_requested += 1; + lease_callbacks.push_back(callback); + return Status::OK(); + } + + ray::Status CancelResourceReserve( + BundleSpecification &bundle_spec, + const ray::rpc::ClientCallback &callback) + override { + num_return_requested += 1; + return_callbacks.push_back(callback); + return Status::OK(); + } + + // Trigger reply to RequestWorkerLease. + bool GrantResourceReserve(bool success = true) { + Status status = Status::OK(); + rpc::RequestResourceReserveReply reply; + reply.set_success(success); + if (lease_callbacks.size() == 0) { + return false; + } else { + auto callback = lease_callbacks.front(); + callback(status, reply); + lease_callbacks.pop_front(); + return true; + } + } + + ~MockRayletResourceClient() {} + + int num_lease_requested = 0; + int num_return_requested = 0; + ClientID node_id = ClientID::FromRandom(); + std::list> lease_callbacks = {}; + std::list> return_callbacks = {}; + }; class MockedGcsActorScheduler : public gcs::GcsActorScheduler { public: using gcs::GcsActorScheduler::GcsActorScheduler; @@ -166,6 +211,15 @@ struct GcsServerMocker { int num_retry_creating_count_ = 0; }; + class MockedGcsPlacementGroupScheduler : public gcs::GcsPlacementGroupScheduler { + public: + using gcs::GcsPlacementGroupScheduler::GcsPlacementGroupScheduler; + + void ResetLeaseClientFactory( + gcs::ReserveResourceClientFactoryFn lease_client_factory) { + lease_client_factory_ = std::move(lease_client_factory); + } + }; class MockedGcsActorTable : public gcs::GcsActorTable { public: MockedGcsActorTable(std::shared_ptr store_client) diff --git a/src/ray/gcs/test/gcs_test_util.h b/src/ray/gcs/test/gcs_test_util.h index 006815415..f6d8d72c4 100644 --- a/src/ray/gcs/test/gcs_test_util.h +++ b/src/ray/gcs/test/gcs_test_util.h @@ -23,6 +23,7 @@ #include "ray/common/test_util.h" #include "ray/protobuf/gcs_service.grpc.pb.h" #include "ray/util/asio_util.h" +#include "src/ray/common/placement_group.h" namespace ray { @@ -62,6 +63,31 @@ struct Mocker { return request; } + static PlacementGroupSpecification GenPlacementGroupCreation( + const std::string &name, std::vector &bundles, + rpc::PlacementStrategy strategy) { + PlacementGroupSpecBuilder builder; + + auto placement_group_id = PlacementGroupID::FromRandom(); + builder.SetPlacementGroupSpec(placement_group_id, name, bundles, strategy); + return builder.Build(); + } + + static rpc::CreatePlacementGroupRequest GenCreatePlacementGroupRequest( + const std::string name = "") { + rpc::CreatePlacementGroupRequest request; + std::vector bundles; + rpc::PlacementStrategy strategy = rpc::PlacementStrategy::SPREAD; + rpc::Bundle bundle; + bundles.push_back(bundle); + bundles.push_back(bundle); + auto placement_group_creation_spec = + GenPlacementGroupCreation(name, bundles, strategy); + request.mutable_placement_group_spec()->CopyFrom( + placement_group_creation_spec.GetMessage()); + return request; + } + static std::shared_ptr GenNodeInfo( uint16_t port = 0, const std::string address = "127.0.0.1") { auto node = std::make_shared(); diff --git a/src/ray/protobuf/common.proto b/src/ray/protobuf/common.proto index fd36fac2f..ed8d29c8b 100644 --- a/src/ray/protobuf/common.proto +++ b/src/ray/protobuf/common.proto @@ -43,6 +43,14 @@ enum TaskType { DRIVER_TASK = 3; } +// Type of placement group strategy. +enum PlacementStrategy { + // Packs Bundles close together inside processes or nodes as tight as possible + PACK = 0; + // Places Bundles across distinct nodes or processes as even as possible. + SPREAD = 1; +} + // Address of a worker or node manager. message Address { bytes raylet_id = 1; @@ -128,6 +136,26 @@ message TaskSpec { int32 max_retries = 16; } +message Bundle { + message BundleIdentifier { + bytes placement_group_id = 1; + int32 bundle_index = 2; + } + BundleIdentifier bundle_id = 1; + map unit_resources = 2; +} + +message PlacementGroupSpec { + // ID of the PlacementGroup. + bytes placement_group_id = 1; + // The name of the placement group. + string name = 2; + // The array of the bundle in Placement Group. + repeated Bundle bundles = 3; + // The schedule strategy of this Placement Group. + PlacementStrategy strategy = 4; +} + message ObjectReference { // ObjectID that the worker has a reference to. bytes object_id = 1; diff --git a/src/ray/protobuf/gcs.proto b/src/ray/protobuf/gcs.proto index 91940a780..5c4474da4 100644 --- a/src/ray/protobuf/gcs.proto +++ b/src/ray/protobuf/gcs.proto @@ -45,6 +45,7 @@ enum TablePrefix { WORKERS = 19; INTERNAL_CONFIG = 20; TABLE_PREFIX_MAX = 21; + PLACEMENT_GROUP_SCHEDULE = 22; } // The channel that Add operations to the Table should be published on, if any. @@ -156,6 +157,33 @@ message ErrorTableData { double timestamp = 4; } +message PlacementGroupTableData { + // State of a placement group. + enum PlacementGroupState { + // Placement Group is pending or scheduling + PENDING = 0; + // Placement Group is alive. + ALIVE = 1; + // Placement Group is already dead and won't be reschedule. + DEAD = 2; + } + + // ID of the PlacementGroup. + bytes placement_group_id = 1; + // The name of the placement group. + string name = 2; + // The array of the bundle in Placement Group. + repeated Bundle bundles = 3; + // The schedule strategy of this Placement Group. + PlacementStrategy strategy = 4; + // Current state of this placement group. + PlacementGroupState state = 5; +} + +message ScheduleData { + map schedule_plan = 1; +} + message ProfileTableData { // Represents a profile event. message ProfileEvent { diff --git a/src/ray/protobuf/gcs_service.proto b/src/ray/protobuf/gcs_service.proto index 829fb049a..9fac12af7 100644 --- a/src/ray/protobuf/gcs_service.proto +++ b/src/ray/protobuf/gcs_service.proto @@ -486,6 +486,14 @@ message CreateActorReply { GcsStatus status = 1; } +message CreatePlacementGroupRequest { + PlacementGroupSpec placement_group_spec = 1; +} + +message CreatePlacementGroupReply { + GcsStatus status = 1; +} + enum GcsServiceFailureType { RPC_DISCONNECT = 0; GCS_SERVER_RESTART = 1; diff --git a/src/ray/protobuf/node_manager.proto b/src/ray/protobuf/node_manager.proto index 4bdc594c2..f596e8b63 100644 --- a/src/ray/protobuf/node_manager.proto +++ b/src/ray/protobuf/node_manager.proto @@ -37,6 +37,24 @@ message RequestWorkerLeaseReply { bool canceled = 4; } +message RequestResourceReserveRequest { + // Bundle containing the requested resources. + Bundle bundle_spec = 1; +} + +message RequestResourceReserveReply { + // The status if resource reserve success. + bool success = 1; +} + +message CancelResourceReserveRequest { + // Bundle containing the requested resources. + Bundle bundle_spec = 1; +} + +message CancelResourceReserveReply { +} + // Release a worker back to its raylet. message ReturnWorkerRequest { // Port of the leased worker that we are now returning. @@ -120,8 +138,8 @@ message FormatGlobalMemoryInfoRequest { } message FormatGlobalMemoryInfoReply { - // A tabular summary of the memory stats. To get this data in structured form, you - // can instead use GetNodeStats() directly. + // A tabular summary of the memory stats. To get this data in structured form, + // you can instead use GetNodeStats() directly. string memory_summary = 1; } @@ -131,6 +149,12 @@ service NodeManagerService { rpc RequestWorkerLease(RequestWorkerLeaseRequest) returns (RequestWorkerLeaseReply); // Release a worker back to its raylet. rpc ReturnWorker(ReturnWorkerRequest) returns (ReturnWorkerReply); + // Request resource from the raylet for a bundle. + rpc RequestResourceReserve(RequestResourceReserveRequest) + returns (RequestResourceReserveReply); + // Return resource for the raylet. + rpc CancelResourceReserve(CancelResourceReserveRequest) + returns (CancelResourceReserveReply); // Cancel a pending lease request. This only returns success if the // lease request was not yet granted. rpc CancelWorkerLease(CancelWorkerLeaseRequest) returns (CancelWorkerLeaseReply); diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 1e13f6a4c..f8a9b7868 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -1924,6 +1924,38 @@ void NodeManager::HandleRequestWorkerLease(const rpc::RequestWorkerLeaseRequest SubmitTask(task, Lineage()); } +void NodeManager::HandleRequestResourceReserve( + const rpc::RequestResourceReserveRequest &request, + rpc::RequestResourceReserveReply *reply, rpc::SendReplyCallback send_reply_callback) { + auto bundle_spec = BundleSpecification(request.bundle_spec()); + RAY_LOG(DEBUG) << "bundle lease request " << bundle_spec.BundleId().first + << bundle_spec.BundleId().second; + auto resource_ids = ScheduleBundle(cluster_resource_map_, bundle_spec); + if (resource_ids.AvailableResources().size() == 0) { + reply->set_success(false); + send_reply_callback(Status::OK(), nullptr, nullptr); + } else { + reply->set_success(true); + send_reply_callback(Status::OK(), nullptr, nullptr); + } +} + +void NodeManager::HandleCancelResourceReserve( + const rpc::CancelResourceReserveRequest &request, + rpc::CancelResourceReserveReply *reply, rpc::SendReplyCallback send_reply_callback) { + auto bundle_spec = BundleSpecification(request.bundle_spec()); + RAY_LOG(DEBUG) << "bundle return resource request " << bundle_spec.BundleId().first + << bundle_spec.BundleId().second; + auto bundle_id_str = bundle_spec.BundleIdAsString(); + auto resource_set = bundle_spec.GetRequiredResources(); + for (auto resource : resource_set.GetResourceMap()) { + std::string resource_name = bundle_id_str + "_" + resource.first; + local_available_resources_.CancelResourceReserve(resource_name); + } + cluster_resource_map_[self_node_id_].ReturnBundleResource(bundle_id_str); + send_reply_callback(Status::OK(), nullptr, nullptr); +} + void NodeManager::HandleReturnWorker(const rpc::ReturnWorkerRequest &request, rpc::ReturnWorkerReply *reply, rpc::SendReplyCallback send_reply_callback) { @@ -2047,6 +2079,31 @@ void NodeManager::ProcessSetResourceRequest( } } +ResourceIdSet NodeManager::ScheduleBundle( + std::unordered_map &resource_map, + const BundleSpecification &bundle_spec) { + // If the resource map contains the local raylet, update load before calling policy. + if (resource_map.count(self_node_id_) > 0) { + resource_map[self_node_id_].SetLoadResources(local_queues_.GetResourceLoad()); + } + // Invoke the scheduling policy. + auto reserve_resource_success = + scheduling_policy_.ScheduleBundle(resource_map, self_node_id_, bundle_spec); + auto bundle_id_str = bundle_spec.BundleIdAsString(); + ResourceIdSet acquired_resources; + if (reserve_resource_success) { + acquired_resources = + local_available_resources_.Acquire(bundle_spec.GetRequiredResources()); + for (auto resource : acquired_resources.AvailableResources()) { + std::string resource_name = bundle_id_str + "_" + resource.first; + local_available_resources_.AddBundleResource(resource_name, resource.second); + } + cluster_resource_map_[self_node_id_].UpdateBundleResource( + bundle_id_str, bundle_spec.GetRequiredResources()); + } + return acquired_resources; +} + void NodeManager::ScheduleTasks( std::unordered_map &resource_map) { // If the resource map contains the local raylet, update load before calling policy. @@ -2587,7 +2644,6 @@ void NodeManager::AssignTask(const std::shared_ptr &worker, const Task & std::vector> *post_assign_callbacks) { const TaskSpecification &spec = task.GetTaskSpecification(); RAY_CHECK(post_assign_callbacks); - // If this is an actor task, check that the new task has the correct counter. if (spec.IsActorTask()) { // An actor task should only be ready to be assigned if it matches the diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index 2b3d3f73b..8ac626fe2 100644 --- a/src/ray/raylet/node_manager.h +++ b/src/ray/raylet/node_manager.h @@ -36,6 +36,7 @@ #include "ray/raylet/task_dependency_manager.h" #include "ray/raylet/worker_pool.h" #include "ray/util/ordered_set.h" +#include "ray/common/bundle_spec.h" // clang-format on namespace ray { @@ -298,6 +299,17 @@ class NodeManager : public rpc::NodeManagerServiceHandler { /// resource_map argument. /// \return Void. void ScheduleTasks(std::unordered_map &resource_map); + + /// Make a placement decision for the resource_map. + /// + /// \param resource_map A mapping from node manager ID to an estimate of the + /// resources available to that node manager. Scheduling decisions will only + /// consider the local node manager and the node managers in the keys of the + /// resource_map argument. + /// \return ResourceIdSet. + ResourceIdSet ScheduleBundle( + std::unordered_map &resource_map, + const BundleSpecification &bundle_spec); /// Handle a task whose return value(s) must be reconstructed. /// /// \param task_id The relevant task ID. @@ -578,6 +590,16 @@ class NodeManager : public rpc::NodeManagerServiceHandler { /// \return Status indicating whether setup was successful. ray::Status SetupPlasmaSubscription(); + /// Handle a `ResourcesLease` request. + void HandleRequestResourceReserve(const rpc::RequestResourceReserveRequest &request, + rpc::RequestResourceReserveReply *reply, + rpc::SendReplyCallback send_reply_callback) override; + + /// Handle a `ResourcesReturn` request. + void HandleCancelResourceReserve(const rpc::CancelResourceReserveRequest &request, + rpc::CancelResourceReserveReply *reply, + rpc::SendReplyCallback send_reply_callback) override; + /// Handle a `WorkerLease` request. void HandleRequestWorkerLease(const rpc::RequestWorkerLeaseRequest &request, rpc::RequestWorkerLeaseReply *reply, diff --git a/src/ray/raylet/raylet_client.cc b/src/ray/raylet/raylet_client.cc index fbde68211..9b7f023a3 100644 --- a/src/ray/raylet/raylet_client.cc +++ b/src/ray/raylet/raylet_client.cc @@ -421,6 +421,22 @@ ray::Status raylet::RayletClient::CancelWorkerLease( return grpc_client_->CancelWorkerLease(request, callback); } +Status raylet::RayletClient::RequestResourceReserve( + const BundleSpecification &bundle_spec, + const ray::rpc::ClientCallback &callback) { + rpc::RequestResourceReserveRequest request; + request.mutable_bundle_spec()->CopyFrom(bundle_spec.GetMessage()); + return grpc_client_->RequestResourceReserve(request, callback); +} + +Status raylet::RayletClient::CancelResourceReserve( + BundleSpecification &bundle_spec, + const ray::rpc::ClientCallback &callback) { + rpc::CancelResourceReserveRequest request; + request.mutable_bundle_spec()->CopyFrom(bundle_spec.GetMessage()); + return grpc_client_->CancelResourceReserve(request, callback); +} + Status raylet::RayletClient::PinObjectIDs( const rpc::Address &caller_address, const std::vector &object_ids, const rpc::ClientCallback &callback) { diff --git a/src/ray/raylet/raylet_client.h b/src/ray/raylet/raylet_client.h index 37409bf5d..278f8b017 100644 --- a/src/ray/raylet/raylet_client.h +++ b/src/ray/raylet/raylet_client.h @@ -22,6 +22,7 @@ #include #include +#include "ray/common/bundle_spec.h" #include "ray/common/status.h" #include "ray/common/task/task_spec.h" #include "ray/rpc/node_manager/node_manager_client.h" @@ -83,6 +84,24 @@ class WorkerLeaseInterface { virtual ~WorkerLeaseInterface(){}; }; +/// Interface for leasing resource. +class ResourceReserveInterface { + public: + /// Requests a resource from the raylet. The callback will be sent via gRPC. + /// \param resource_spec Resources that should be allocated for the worker. + /// \return ray::Status + virtual ray::Status RequestResourceReserve( + const BundleSpecification &bundle_spec, + const ray::rpc::ClientCallback + &callback) = 0; + + virtual ray::Status CancelResourceReserve( + BundleSpecification &bundle_spec, + const ray::rpc::ClientCallback &callback) = 0; + + virtual ~ResourceReserveInterface(){}; +}; + /// Interface for waiting dependencies. Abstract for testing. class DependencyWaiterInterface { public: @@ -141,7 +160,8 @@ class RayletConnection { class RayletClient : public PinObjectsInterface, public WorkerLeaseInterface, - public DependencyWaiterInterface { + public DependencyWaiterInterface, + public ResourceReserveInterface { public: /// Connect to the raylet. /// @@ -312,6 +332,18 @@ class RayletClient : public PinObjectsInterface, const TaskID &task_id, const rpc::ClientCallback &callback) override; + /// Implements ResourceReserveInterface. + ray::Status RequestResourceReserve( + const BundleSpecification &bundle_spec, + const ray::rpc::ClientCallback &callback) + override; + + /// Implements ResourceReserveInterface. + ray::Status CancelResourceReserve( + BundleSpecification &bundle_spec, + const ray::rpc::ClientCallback &callback) + override; + ray::Status PinObjectIDs( const rpc::Address &caller_address, const std::vector &object_ids, const ray::rpc::ClientCallback &callback) override; diff --git a/src/ray/raylet/scheduling_policy.cc b/src/ray/raylet/scheduling_policy.cc index d4f106998..21e0ab608 100644 --- a/src/ray/raylet/scheduling_policy.cc +++ b/src/ray/raylet/scheduling_policy.cc @@ -133,6 +133,36 @@ std::unordered_map SchedulingPolicy::Schedule( return decision; } +bool SchedulingPolicy::ScheduleBundle( + std::unordered_map &cluster_resources, + const ClientID &local_client_id, const ray::BundleSpecification &bundle_spec) { +#ifndef NDEBUG + RAY_LOG(DEBUG) << "Cluster resource map: "; + for (const auto &client_resource_pair : cluster_resources) { + const ClientID &client_id = client_resource_pair.first; + const SchedulingResources &resources = client_resource_pair.second; + RAY_LOG(DEBUG) << "client_id: " << client_id << " " + << resources.GetAvailableResources().ToString(); + } +#endif + const auto &client_resource_pair = cluster_resources.find(local_client_id); + if (client_resource_pair == cluster_resources.end()) { + return false; + } + const auto &resource_demand = bundle_spec.GetRequiredResources(); + ClientID node_client_id = client_resource_pair->first; + const auto &node_resources = client_resource_pair->second; + ResourceSet available_node_resources = + ResourceSet(node_resources.GetAvailableResources()); + available_node_resources.SubtractResources(node_resources.GetLoadResources()); + RAY_LOG(DEBUG) << "client_id " << node_client_id + << " avail: " << node_resources.GetAvailableResources().ToString() + << " load: " << node_resources.GetLoadResources().ToString(); + /// If the resource_demand is subset of the whole available_node_resources, this bundle + /// can be set in this node, return true. + return resource_demand.IsSubset(available_node_resources); +} + std::vector SchedulingPolicy::SpillOver( SchedulingResources &remote_scheduling_resources) const { // The policy decision to be returned. diff --git a/src/ray/raylet/scheduling_policy.h b/src/ray/raylet/scheduling_policy.h index 62349ca95..d921ad1d7 100644 --- a/src/ray/raylet/scheduling_policy.h +++ b/src/ray/raylet/scheduling_policy.h @@ -17,6 +17,7 @@ #include #include +#include "ray/common/bundle_spec.h" #include "ray/common/task/scheduling_resources.h" #include "ray/raylet/scheduling_queue.h" @@ -49,6 +50,17 @@ class SchedulingPolicy { std::unordered_map &cluster_resources, const ClientID &local_client_id); + /// \param cluster_resources: a set of cluster resources containing resource and load + /// information for some subset of the cluster. + /// \param local_client_id The ID of the node manager that owns this + /// SchedulingPolicy object. + /// \param bundle_spec the description of a bundle which include the resource the bundle + /// need. \return If this bundle can be scheduled in this node, return true; else return + /// false. + bool ScheduleBundle( + std::unordered_map &cluster_resources, + const ClientID &local_client_id, const ray::BundleSpecification &bundle_spec); + /// \brief Given a set of cluster resources perform a spill-over scheduling operation. /// /// \param cluster_resources: a set of cluster resources containing resource and load diff --git a/src/ray/rpc/node_manager/node_manager_client.h b/src/ray/rpc/node_manager/node_manager_client.h index 976990ea6..0373b5c6f 100644 --- a/src/ray/rpc/node_manager/node_manager_client.h +++ b/src/ray/rpc/node_manager/node_manager_client.h @@ -89,6 +89,12 @@ class NodeManagerWorkerClient /// Cancel a pending worker lease request. RPC_CLIENT_METHOD(NodeManagerService, CancelWorkerLease, grpc_client_, ) + /// Request resource lease. + RPC_CLIENT_METHOD(NodeManagerService, RequestResourceReserve, grpc_client_, ) + + /// Return resource lease. + RPC_CLIENT_METHOD(NodeManagerService, CancelResourceReserve, grpc_client_, ) + /// Notify the raylet to pin the provided object IDs. RPC_CLIENT_METHOD(NodeManagerService, PinObjectIDs, grpc_client_, ) diff --git a/src/ray/rpc/node_manager/node_manager_server.h b/src/ray/rpc/node_manager/node_manager_server.h index de80ae260..6daef365c 100644 --- a/src/ray/rpc/node_manager/node_manager_server.h +++ b/src/ray/rpc/node_manager/node_manager_server.h @@ -23,15 +23,17 @@ namespace ray { namespace rpc { /// NOTE: See src/ray/core_worker/core_worker.h on how to add a new grpc handler. -#define RAY_NODE_MANAGER_RPC_HANDLERS \ - RPC_SERVICE_HANDLER(NodeManagerService, RequestWorkerLease) \ - RPC_SERVICE_HANDLER(NodeManagerService, ReturnWorker) \ - RPC_SERVICE_HANDLER(NodeManagerService, CancelWorkerLease) \ - RPC_SERVICE_HANDLER(NodeManagerService, ForwardTask) \ - RPC_SERVICE_HANDLER(NodeManagerService, PinObjectIDs) \ - RPC_SERVICE_HANDLER(NodeManagerService, GetNodeStats) \ - RPC_SERVICE_HANDLER(NodeManagerService, GlobalGC) \ - RPC_SERVICE_HANDLER(NodeManagerService, FormatGlobalMemoryInfo) +#define RAY_NODE_MANAGER_RPC_HANDLERS \ + RPC_SERVICE_HANDLER(NodeManagerService, RequestWorkerLease) \ + RPC_SERVICE_HANDLER(NodeManagerService, ReturnWorker) \ + RPC_SERVICE_HANDLER(NodeManagerService, CancelWorkerLease) \ + RPC_SERVICE_HANDLER(NodeManagerService, ForwardTask) \ + RPC_SERVICE_HANDLER(NodeManagerService, PinObjectIDs) \ + RPC_SERVICE_HANDLER(NodeManagerService, GetNodeStats) \ + RPC_SERVICE_HANDLER(NodeManagerService, GlobalGC) \ + RPC_SERVICE_HANDLER(NodeManagerService, FormatGlobalMemoryInfo) \ + RPC_SERVICE_HANDLER(NodeManagerService, RequestResourceReserve) \ + RPC_SERVICE_HANDLER(NodeManagerService, CancelResourceReserve) /// Interface of the `NodeManagerService`, see `src/ray/protobuf/node_manager.proto`. class NodeManagerServiceHandler { @@ -59,6 +61,16 @@ class NodeManagerServiceHandler { rpc::CancelWorkerLeaseReply *reply, rpc::SendReplyCallback send_reply_callback) = 0; + virtual void HandleRequestResourceReserve( + const rpc::RequestResourceReserveRequest &request, + rpc::RequestResourceReserveReply *reply, + rpc::SendReplyCallback send_reply_callback) = 0; + + virtual void HandleCancelResourceReserve( + const rpc::CancelResourceReserveRequest &request, + rpc::CancelResourceReserveReply *reply, + rpc::SendReplyCallback send_reply_callback) = 0; + virtual void HandleForwardTask(const ForwardTaskRequest &request, ForwardTaskReply *reply, SendReplyCallback send_reply_callback) = 0;