wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-04 13:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Object Manager] Pull Manager refactor (#12335)

Browse Source
This commit is contained in:
Alex Wu
2020-12-11 11:56:23 -08:00
committed by GitHub
parent 3d8c1cbae6
commit 676ec363f6
10 changed files with 539 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files
BUILD.bazel
+12
View File
@@ -813,6 +813,18 @@ cc_test(
],
)
cc_test(
name = "pull_manager_test",
srcs = [
"src/ray/object_manager/test/pull_manager_test.cc",
],
copts = COPTS,
deps = [
":raylet_lib",
"@com_google_googletest//:gtest_main",
],
)
cc_test(
name = "push_manager_test",
srcs = [
src/ray/common/ray_config_def.h
+4
View File
@@ -177,6 +177,10 @@ RAY_CONFIG(int64_t, worker_register_timeout_seconds, 30)
RAY_CONFIG(int64_t, redis_db_connect_retries, 50)
RAY_CONFIG(int64_t, redis_db_connect_wait_milliseconds, 100)
/// Timeout, in milliseconds, to wait before retrying a failed pull in the
/// ObjectManager.
RAY_CONFIG(int, object_manager_timer_freq_ms, 100)
/// Timeout, in milliseconds, to wait before retrying a failed pull in the
/// ObjectManager.
RAY_CONFIG(int, object_manager_pull_timeout_ms, 10000)
src/ray/object_manager/common.h
+8
View File
@@ -1,8 +1,12 @@
#pragma once
#include <boost/asio.hpp>
#include <functional>
#include "ray/common/id.h"
#include "ray/object_manager/format/object_manager_generated.h"
namespace ray {
/// A callback to asynchronously spill objects when space is needed.
@@ -18,4 +22,8 @@ using SpillObjectsCallback =
/// A callback to call when space has been released.
using SpaceReleasedCallback = std::function<void()>;
/// A callback to call when a spilled object needs to be returned to the object store.
using RestoreSpilledObjectCallback = std::function<void(
const ObjectID &, const std::string &, std::function<void(const ray::Status &)>)>;
} // namespace ray
src/ray/object_manager/object_directory.h
+6 -6
View File
@@ -40,6 +40,11 @@ struct RemoteConnectionInfo {
uint16_t port;
};
/// Callback for object location notifications.
using OnLocationsFound =
std::function<void(const ray::ObjectID &object_id,
const std::unordered_set<ray::NodeID> &, const std::string &)>;
class ObjectDirectoryInterface {
public:
virtual ~ObjectDirectoryInterface() {}
@@ -58,12 +63,7 @@ class ObjectDirectoryInterface {
/// \return A vector of information for all connected remote object managers.
virtual std::vector<RemoteConnectionInfo> LookupAllRemoteConnections() const = 0;
/// Callback for object location notifications.
using OnLocationsFound =
std::function<void(const ray::ObjectID &object_id,
const std::unordered_set<ray::NodeID> &, const std::string &)>;
/// Lookup object locations. Callback may be invoked with empty list of node ids.
/// Lookup object locations. Callback may be invoked with empty list of client ids.
///
/// \param object_id The object's ObjectID.
/// \param callback Invoked with (possibly empty) list of node ids and object_id.
src/ray/object_manager/object_manager.cc
+60 -160
View File
@@ -57,25 +57,47 @@ ObjectManager::ObjectManager(asio::io_service &main_service, const NodeID &self_
RestoreSpilledObjectCallback restore_spilled_object,
SpillObjectsCallback spill_objects_callback,
std::function<void()> object_store_full_callback)
: self_node_id_(self_node_id),
: main_service_(&main_service),
self_node_id_(self_node_id),
config_(config),
object_directory_(std::move(object_directory)),
object_store_internal_(config, spill_objects_callback, object_store_full_callback),
buffer_pool_(config_.store_socket_name, config_.object_chunk_size),
rpc_work_(rpc_service_),
gen_(std::chrono::high_resolution_clock::now().time_since_epoch().count()),
object_manager_server_("ObjectManager", config_.object_manager_port,
config_.rpc_service_threads_number),
object_manager_service_(rpc_service_, *this),
client_call_manager_(main_service, config_.rpc_service_threads_number),
restore_spilled_object_(restore_spilled_object) {
restore_spilled_object_(restore_spilled_object),
pull_retry_timer_(*main_service_,
boost::posix_time::milliseconds(config.timer_freq_ms)) {
RAY_CHECK(config_.rpc_service_threads_number > 0);
main_service_ = &main_service;
const auto &object_is_local = [this](const ObjectID &object_id) {
return local_objects_.count(object_id) != 0;
};
const auto &send_pull_request = [this](const ObjectID &object_id,
const NodeID &client_id) {
SendPullRequest(object_id, client_id);
};
const auto &get_time = []() { return absl::GetCurrentTimeNanos() / 1e9; };
pull_manager_.reset(new PullManager(self_node_id_, object_is_local, send_pull_request,
restore_spilled_object_, get_time,
config.pull_timeout_ms));
push_manager_.reset(new PushManager(/* max_chunks_in_flight= */ std::max(
static_cast<int64_t>(1L),
static_cast<int64_t>(config_.max_bytes_in_flight / config_.object_chunk_size))));
pull_retry_timer_.async_wait([this](const boost::system::error_code &e) {
RAY_CHECK(!e) << "The raylet's object manager has failed unexpectedly with error: "
<< e
<< ". Please file a bug report on here: "
"https://github.com/ray-project/ray/issues";
Tick();
});
if (plasma::plasma_store_runner) {
store_notification_ = std::make_shared<ObjectStoreNotificationManager>(main_service);
plasma::plasma_store_runner->SetNotificationListener(store_notification_);
@@ -178,171 +200,48 @@ ray::Status ObjectManager::SubscribeObjDeleted(
ray::Status ObjectManager::Pull(const ObjectID &object_id,
const rpc::Address &owner_address) {
RAY_LOG(DEBUG) << "Pull on " << self_node_id_ << " of object " << object_id;
// Check if object is already local.
if (local_objects_.count(object_id) != 0) {
RAY_LOG(ERROR) << object_id << " attempted to pull an object that's already local.";
return ray::Status::OK();
}
if (pull_requests_.find(object_id) != pull_requests_.end()) {
RAY_LOG(DEBUG) << object_id << " has inflight pull_requests, skipping.";
return ray::Status::OK();
if (!pull_manager_->Pull(object_id, owner_address)) {
// If we don't need to pull, the object is either already local or this is a duplicate
// request.
return Status::OK();
}
pull_requests_.emplace(object_id, PullRequest());
const auto &callback = [this](const ObjectID &object_id,
const std::unordered_set<NodeID> &client_ids,
const std::string &spilled_url) {
pull_manager_->OnLocationChange(object_id, client_ids, spilled_url);
};
// Subscribe to object notifications. A notification will be received every
// time the set of node IDs for the object changes. Notifications will also
// be received if the list of locations is empty. The set of node IDs has
// no ordering guarantee between notifications.
return object_directory_->SubscribeObjectLocations(
object_directory_pull_callback_id_, object_id, owner_address,
[this](const ObjectID &object_id, const std::unordered_set<NodeID> &node_ids,
const std::string &spilled_url) {
// Exit if the Pull request has already been fulfilled or canceled.
auto it = pull_requests_.find(object_id);
if (it == pull_requests_.end()) {
return;
}
// Reset the list of nodes that are now expected to have the object.
// NOTE(swang): Since we are overwriting the previous list of nodes,
// we may end up sending a duplicate request to the same node as
// before.
it->second.node_locations = std::vector<NodeID>(node_ids.begin(), node_ids.end());
if (!spilled_url.empty()) {
// Try to restore the spilled object.
restore_spilled_object_(object_id, spilled_url,
[this, object_id](const ray::Status &status) {
// Fall back to fetching from another object manager.
if (!status.ok()) {
TryPull(object_id);
}
});
} else if (it->second.node_locations.empty()) {
// The object locations are now empty, so we should wait for the next
// notification about a new object location. Cancel the timer until
// the next Pull attempt since there are no more nodes to try.
if (it->second.retry_timer != nullptr) {
it->second.retry_timer->cancel();
it->second.timer_set = false;
}
} else {
// New object locations were found, so begin trying to pull from a
// node. This will be called every time a new node location
// appears.
TryPull(object_id);
}
});
return object_directory_->SubscribeObjectLocations(object_directory_pull_callback_id_,
object_id, owner_address, callback);
}
void ObjectManager::TryPull(const ObjectID &object_id) {
auto it = pull_requests_.find(object_id);
if (it == pull_requests_.end()) {
return;
}
auto &node_vector = it->second.node_locations;
// The timer should never fire if there are no expected node locations.
if (node_vector.empty()) {
return;
}
RAY_CHECK(local_objects_.count(object_id) == 0);
// Make sure that there is at least one node which is not the local node.
// TODO(rkn): It may actually be possible for this check to fail.
if (node_vector.size() == 1 && node_vector[0] == self_node_id_) {
RAY_LOG(WARNING) << "The object manager with ID " << self_node_id_
<< " is trying to pull object " << object_id
<< " but the object table suggests that this object manager "
<< "already has the object. The object may have been evicted. It is "
<< "most likely due to memory pressure, object pull has been "
<< "requested before object location is updated.";
it->second.timer_set = false;
return;
}
// Choose a random node to pull the object from.
// Generate a random index.
std::uniform_int_distribution<int> distribution(0, node_vector.size() - 1);
int node_index = distribution(gen_);
NodeID node_id = node_vector[node_index];
// If the object manager somehow ended up choosing itself, choose a different
// object manager.
if (node_id == self_node_id_) {
std::swap(node_vector[node_index], node_vector[node_vector.size() - 1]);
node_vector.pop_back();
RAY_LOG(WARNING)
<< "The object manager with ID " << self_node_id_ << " is trying to pull object "
<< object_id << " but the object table suggests that this object manager "
<< "already has the object. It is most likely due to memory pressure, object "
<< "pull has been requested before object location is updated.";
node_id = node_vector[node_index % node_vector.size()];
RAY_CHECK(node_id != self_node_id_);
}
RAY_LOG(DEBUG) << "Sending pull request from " << self_node_id_ << " to " << node_id
<< " of object " << object_id;
auto rpc_client = GetRpcClient(node_id);
void ObjectManager::SendPullRequest(const ObjectID &object_id, const NodeID &client_id) {
auto rpc_client = GetRpcClient(client_id);
if (rpc_client) {
// Try pulling from the node.
rpc_service_.post([this, object_id, node_id, rpc_client]() {
SendPullRequest(object_id, node_id, rpc_client);
// Try pulling from the client.
rpc_service_.post([this, object_id, client_id, rpc_client]() {
rpc::PullRequest pull_request;
pull_request.set_object_id(object_id.Binary());
pull_request.set_node_id(self_node_id_.Binary());
rpc_client->Pull(pull_request, [object_id, client_id](const Status &status,
const rpc::PullReply &reply) {
if (!status.ok()) {
RAY_LOG(WARNING) << "Send pull " << object_id << " request to client "
<< client_id << " failed due to" << status.message();
}
});
});
} else {
RAY_LOG(ERROR) << "Couldn't send pull request from " << self_node_id_ << " to "
<< node_id << " of object " << object_id
<< client_id << " of object " << object_id
<< " , setup rpc connection failed.";
}
// If there are more nodes to try, try them in succession, with a timeout
// in between each try.
if (!it->second.node_locations.empty()) {
if (it->second.retry_timer == nullptr) {
// Set the timer if we haven't already.
it->second.retry_timer = std::unique_ptr<boost::asio::deadline_timer>(
new boost::asio::deadline_timer(*main_service_));
}
// Wait for a timeout. If we receive the object or a caller Cancels the
// Pull within the timeout, then nothing will happen. Otherwise, the timer
// will fire and the next node in the list will be tried.
boost::posix_time::milliseconds retry_timeout(config_.pull_timeout_ms);
it->second.retry_timer->expires_from_now(retry_timeout);
it->second.retry_timer->async_wait(
[this, object_id](const boost::system::error_code &error) {
if (!error) {
// Try the Pull from the next node.
TryPull(object_id);
} else {
// Check that the error was due to the timer being canceled.
RAY_CHECK(error == boost::asio::error::operation_aborted);
}
});
// Record that we set the timer until the next attempt.
it->second.timer_set = true;
} else {
// The timer is not reset since there are no more nodes to try. Go back
// to waiting for more notifications. Once we receive a new object location
// from the object directory, then the Pull will be retried.
it->second.timer_set = false;
}
};
void ObjectManager::SendPullRequest(
const ObjectID &object_id, const NodeID &node_id,
std::shared_ptr<rpc::ObjectManagerClient> rpc_client) {
rpc::PullRequest pull_request;
pull_request.set_object_id(object_id.Binary());
pull_request.set_node_id(self_node_id_.Binary());
rpc_client->Pull(pull_request, [object_id, node_id](const Status &status,
const rpc::PullReply &reply) {
if (!status.ok()) {
RAY_LOG(WARNING) << "Send pull " << object_id << " request to node " << node_id
<< " failed due to" << status.message();
}
});
}
void ObjectManager::HandlePushTaskTimeout(const ObjectID &object_id,
@@ -528,14 +427,13 @@ void ObjectManager::SendObjectChunk(const UniqueID &push_id, const ObjectID &obj
}
void ObjectManager::CancelPull(const ObjectID &object_id) {
auto it = pull_requests_.find(object_id);
if (it == pull_requests_.end()) {
if (!pull_manager_->CancelPull(object_id)) {
// We weren't tracking a pull request for this object, so there is nothing to cancel.
return;
}
RAY_CHECK_OK(object_directory_->UnsubscribeObjectLocations(
object_directory_pull_callback_id_, object_id));
pull_requests_.erase(it);
}
ray::Status ObjectManager::Wait(
@@ -898,7 +796,7 @@ std::string ObjectManager::DebugString() const {
result << "\n- num local objects: " << local_objects_.size();
result << "\n- num active wait requests: " << active_wait_requests_.size();
result << "\n- num unfulfilled push requests: " << unfulfilled_push_requests_.size();
result << "\n- num pull requests: " << pull_requests_.size();
result << "\n- num pull requests: " << pull_manager_->NumActiveRequests();
result << "\n- num buffered profile events: " << profile_events_.size();
result << "\n- num chunks received total: " << num_chunks_received_total_;
result << "\n- num chunks received failed: " << num_chunks_received_failed_;
@@ -913,7 +811,9 @@ void ObjectManager::RecordMetrics() const {
stats::ObjectStoreAvailableMemory().Record(config_.object_store_memory - used_memory_);
stats::ObjectStoreUsedMemory().Record(used_memory_);
stats::ObjectStoreLocalObjects().Record(local_objects_.size());
stats::ObjectManagerPullRequests().Record(pull_requests_.size());
stats::ObjectManagerPullRequests().Record(pull_manager_->NumActiveRequests());
}
void ObjectManager::Tick() { pull_manager_->Tick(); }
} // namespace ray
src/ray/object_manager/object_manager.h
+19 -33
View File
@@ -32,12 +32,14 @@
#include "ray/common/id.h"
#include "ray/common/ray_config.h"
#include "ray/common/status.h"
#include "ray/object_manager/common.h"
#include "ray/object_manager/format/object_manager_generated.h"
#include "ray/object_manager/notification/object_store_notification_manager_ipc.h"
#include "ray/object_manager/object_buffer_pool.h"
#include "ray/object_manager/object_directory.h"
#include "ray/object_manager/ownership_based_object_directory.h"
#include "ray/object_manager/plasma/store_runner.h"
#include "ray/object_manager/pull_manager.h"
#include "ray/object_manager/push_manager.h"
#include "ray/rpc/object_manager/object_manager_client.h"
#include "ray/rpc/object_manager/object_manager_server.h"
@@ -49,6 +51,8 @@ struct ObjectManagerConfig {
/// from other object managers. If this is 0, the object manager will choose
/// its own port.
int object_manager_port;
/// The object manager's global timer frequency.
unsigned int timer_freq_ms;
/// The time in milliseconds to wait before retrying a pull
/// that fails due to node id lookup.
unsigned int pull_timeout_ms;
@@ -78,7 +82,6 @@ struct LocalObjectInfo {
/// Information from the object store about the object.
object_manager::protocol::ObjectInfoT object_info;
};
class ObjectStoreRunner {
public:
ObjectStoreRunner(const ObjectManagerConfig &config,
@@ -171,9 +174,8 @@ class ObjectManager : public ObjectManagerInterface,
/// Send pull request
///
/// \param object_id Object id
/// \param node_id Remote server node id
void SendPullRequest(const ObjectID &object_id, const NodeID &node_id,
std::shared_ptr<rpc::ObjectManagerClient> rpc_client);
/// \param client_id Remote server client id
void SendPullRequest(const ObjectID &object_id, const NodeID &client_id);
/// Get the rpc client according to the node ID
///
@@ -235,17 +237,6 @@ class ObjectManager : public ObjectManagerInterface,
/// \return Status of whether the pull request successfully initiated.
ray::Status Pull(const ObjectID &object_id, const rpc::Address &owner_address) override;
/// Try to Pull an object from one of its expected node locations. If there
/// are more node locations to try after this attempt, then this method
/// will try each of the other nodes in succession, with a timeout between
/// each attempt. If the object is received or if the Pull is Canceled before
/// the timeout, then no more Pull requests for this object will be sent
/// to other node managers until TryPull is called again.
///
/// \param object_id The object's object id.
/// \return Void.
void TryPull(const ObjectID &object_id);
/// Cancels all requests (Push/Pull) associated with the given ObjectID. This
/// method is idempotent.
///
@@ -293,16 +284,11 @@ class ObjectManager : public ObjectManagerInterface,
/// Record metrics.
void RecordMetrics() const;
void Tick();
private:
friend class TestObjectManager;
struct PullRequest {
PullRequest() : retry_timer(nullptr), timer_set(false), node_locations() {}
std::unique_ptr<boost::asio::deadline_timer> retry_timer;
bool timer_set;
std::vector<NodeID> node_locations;
};
struct WaitState {
WaitState(boost::asio::io_service &service, int64_t timeout_ms,
const WaitCallback &callback)
@@ -406,6 +392,10 @@ class ObjectManager : public ObjectManagerInterface,
/// Handle Push task timeout.
void HandlePushTaskTimeout(const ObjectID &object_id, const NodeID &node_id);
/// Weak reference to main service. We ensure this object is destroyed before
/// main_service_ is stopped.
boost::asio::io_service *main_service_;
NodeID self_node_id_;
const ObjectManagerConfig config_;
std::shared_ptr<ObjectDirectoryInterface> object_directory_;
@@ -416,10 +406,6 @@ class ObjectManager : public ObjectManagerInterface,
std::shared_ptr<ObjectStoreNotificationManager> store_notification_;
ObjectBufferPool buffer_pool_;
/// Weak reference to main service. We ensure this object is destroyed before
/// main_service_ is stopped.
boost::asio::io_service *main_service_;
/// Multi-thread asio service, deal with all outgoing and incoming RPC request.
boost::asio::io_service rpc_service_;
@@ -448,10 +434,6 @@ class ObjectManager : public ObjectManagerInterface,
ObjectID, std::unordered_map<NodeID, std::unique_ptr<boost::asio::deadline_timer>>>
unfulfilled_push_requests_;
/// The objects that this object manager is currently trying to fetch from
/// remote object managers.
std::unordered_map<ObjectID, PullRequest> pull_requests_;
/// Profiling events that are to be batched together and added to the profile
/// table in the GCS.
std::vector<rpc::ProfileTableData::ProfileEvent> profile_events_;
@@ -460,9 +442,6 @@ class ObjectManager : public ObjectManagerInterface,
/// and rpc thread.
std::mutex profile_mutex_;
/// Internally maintained random number generator.
std::mt19937_64 gen_;
/// The gPRC server.
rpc::GrpcServer object_manager_server_;
@@ -478,9 +457,16 @@ class ObjectManager : public ObjectManagerInterface,
const RestoreSpilledObjectCallback restore_spilled_object_;
/// Pull manager retry timer .
/* std::unique_ptr<boost::asio::deadline_timer> pull_retry_timer_; */
boost::asio::deadline_timer pull_retry_timer_;
/// Object push manager.
std::unique_ptr<PushManager> push_manager_;
/// Object pull manager.
std::unique_ptr<PullManager> pull_manager_;
/// Running sum of the amount of memory used in the object store.
int64_t used_memory_ = 0;
src/ray/object_manager/pull_manager.cc
+139
View File
@@ -0,0 +1,139 @@
#include "ray/object_manager/pull_manager.h"
namespace ray {
PullManager::PullManager(
NodeID &self_node_id, const std::function<bool(const ObjectID &)> object_is_local,
const std::function<void(const ObjectID &, const NodeID &)> send_pull_request,
const RestoreSpilledObjectCallback restore_spilled_object,
const std::function<double()> get_time, int pull_timeout_ms)
: self_node_id_(self_node_id),
object_is_local_(object_is_local),
send_pull_request_(send_pull_request),
restore_spilled_object_(restore_spilled_object),
get_time_(get_time),
pull_timeout_ms_(pull_timeout_ms),
gen_(std::chrono::high_resolution_clock::now().time_since_epoch().count()) {}
bool PullManager::Pull(const ObjectID &object_id, const rpc::Address &owner_address) {
RAY_LOG(DEBUG) << "Pull "
<< " of object " << object_id;
// Check if object is already local.
if (object_is_local_(object_id)) {
RAY_LOG(DEBUG) << object_id << " attempted to pull an object that's already local.";
return false;
}
if (pull_requests_.find(object_id) != pull_requests_.end()) {
RAY_LOG(DEBUG) << object_id << " has inflight pull_requests, skipping.";
return false;
}
pull_requests_.emplace(object_id, PullRequest(get_time_() + pull_timeout_ms_ / 1000));
return true;
}
void PullManager::OnLocationChange(const ObjectID &object_id,
const std::unordered_set<NodeID> &client_ids,
const std::string &spilled_url) {
// Exit if the Pull request has already been fulfilled or canceled.
auto it = pull_requests_.find(object_id);
if (it == pull_requests_.end()) {
return;
}
// Reset the list of clients that are now expected to have the object.
// NOTE(swang): Since we are overwriting the previous list of clients,
// we may end up sending a duplicate request to the same client as
// before.
it->second.client_locations = std::vector<NodeID>(client_ids.begin(), client_ids.end());
if (!spilled_url.empty()) {
// Try to restore the spilled object.
restore_spilled_object_(object_id, spilled_url,
[this, object_id](const ray::Status &status) {
// Fall back to fetching from another object manager.
if (!status.ok()) {
TryPull(object_id);
}
});
} else {
// New object locations were found, so begin trying to pull from a
// client. This will be called every time a new client location
// appears.
TryPull(object_id);
}
}
void PullManager::TryPull(const ObjectID &object_id) {
auto it = pull_requests_.find(object_id);
if (it == pull_requests_.end()) {
return;
}
auto &node_vector = it->second.client_locations;
// The timer should never fire if there are no expected client locations.
if (node_vector.empty()) {
return;
}
RAY_CHECK(!object_is_local_(object_id));
// Make sure that there is at least one client which is not the local client.
// TODO(rkn): It may actually be possible for this check to fail.
if (node_vector.size() == 1 && node_vector[0] == self_node_id_) {
RAY_LOG(WARNING) << "The object manager with ID " << self_node_id_
<< " is trying to pull object " << object_id
<< " but the object table suggests that this object manager "
<< "already has the object. The object may have been evicted. It is "
<< "most likely due to memory pressure, object pull has been "
<< "requested before object location is updated.";
return;
}
// Choose a random client to pull the object from.
// Generate a random index.
std::uniform_int_distribution<int> distribution(0, node_vector.size() - 1);
int node_index = distribution(gen_);
NodeID node_id = node_vector[node_index];
// If the object manager somehow ended up choosing itself, choose a different
// object manager.
if (node_id == self_node_id_) {
std::swap(node_vector[node_index], node_vector[node_vector.size() - 1]);
node_vector.pop_back();
RAY_LOG(WARNING)
<< "The object manager with ID " << self_node_id_ << " is trying to pull object "
<< object_id << " but the object table suggests that this object manager "
<< "already has the object. It is most likely due to memory pressure, object "
<< "pull has been requested before object location is updated.";
node_id = node_vector[node_index % node_vector.size()];
RAY_CHECK(node_id != self_node_id_);
}
RAY_LOG(DEBUG) << "Sending pull request from " << self_node_id_ << " to " << node_id
<< " of object " << object_id;
send_pull_request_(object_id, node_id);
}
bool PullManager::CancelPull(const ObjectID &object_id) {
auto it = pull_requests_.find(object_id);
if (it == pull_requests_.end()) {
return false;
}
pull_requests_.erase(it);
return true;
}
void PullManager::Tick() {
for (auto &pair : pull_requests_) {
const auto &object_id = pair.first;
auto &request = pair.second;
const auto time = get_time_();
if (time >= request.next_pull_time) {
TryPull(object_id);
request.next_pull_time = time + pull_timeout_ms_ / 1000;
}
}
}
int PullManager::NumActiveRequests() const { return pull_requests_.size(); }
} // namespace ray
src/ray/object_manager/pull_manager.h
+115
View File
@@ -0,0 +1,115 @@
#pragma once
#include <boost/asio.hpp>
#include <boost/asio/error.hpp>
#include <boost/bind.hpp>
#include <map>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/time/clock.h"
#include "ray/common/id.h"
#include "ray/common/ray_config.h"
#include "ray/common/status.h"
#include "ray/object_manager/common.h"
#include "ray/object_manager/format/object_manager_generated.h"
#include "ray/object_manager/notification/object_store_notification_manager_ipc.h"
#include "ray/object_manager/object_buffer_pool.h"
#include "ray/object_manager/object_directory.h"
#include "ray/object_manager/ownership_based_object_directory.h"
#include "ray/object_manager/plasma/store_runner.h"
#include "ray/rpc/object_manager/object_manager_client.h"
#include "ray/rpc/object_manager/object_manager_server.h"
namespace ray {
class PullManager {
public:
/// PullManager is responsible for managing the policy around when to send pull requests
/// and to whom. Notably, it is _not_ responsible for controlling the object directory
/// or any pubsub communications.
///
/// \param self_node_id the current node
/// \param object_is_local A callback which should return true if a given object is
/// already on the local node. \param send_pull_request A callback which should send a
/// pull request to the specified node.
/// \param restore_spilled_object A callback which should
/// retrieve an spilled object from the external store.
PullManager(
NodeID &self_node_id, const std::function<bool(const ObjectID &)> object_is_local,
const std::function<void(const ObjectID &, const NodeID &)> send_pull_request,
const RestoreSpilledObjectCallback restore_spilled_object,
const std::function<double()> get_time, int pull_timeout_ms);
/// Begin a new pull request if necessary.
///
/// \param object_id The object id to pull.
/// \param owner_address The owner of the object.
///
/// \return True if a new pull request was necessary. If true, the caller should
/// subscribe to new locations of the object, and call OnLocationChange when necessary.
bool Pull(const ObjectID &object_id, const rpc::Address &owner_address);
/// Called when the available locations for a given object change.
///
/// \param object_id The ID of the object which is now available in a new location.
/// \param client_ids The new set of nodes that the object is available on. Not
/// necessarily a super or subset of the previously available nodes. \param spilled_url
/// The location of the object if it was spilled. If non-empty, the object may no longer
/// be on any node.
void OnLocationChange(const ObjectID &object_id,
const std::unordered_set<NodeID> &client_ids,
const std::string &spilled_url);
/// Cancel an existing pull request if necessary.
///
/// \param object_id The object id that no longer needs to be pulled.
///
/// \return True if a pull was cancelled. If there was no pending pull request for the
/// object this method may return false.
bool CancelPull(const ObjectID &object_id);
/// Called when the retry timer fires. If this fires, the pull manager may try to pull
/// existing objects from other nodes if necessary.
void Tick();
/// The number of ongoing object pulls.
int NumActiveRequests() const;
private:
/// A helper structure for tracking information about each ongoing object pull.
struct PullRequest {
PullRequest(double first_retry_time)
: client_locations(), next_pull_time(first_retry_time) {}
std::vector<NodeID> client_locations;
double next_pull_time;
};
/// See the constructor's arguments.
NodeID self_node_id_;
const std::function<bool(const ObjectID &)> object_is_local_;
const std::function<void(const ObjectID &, const NodeID &)> send_pull_request_;
const RestoreSpilledObjectCallback restore_spilled_object_;
const std::function<double()> get_time_;
int pull_timeout_ms_;
/// The objects that this object manager is currently trying to fetch from
/// remote object managers.
std::unordered_map<ObjectID, PullRequest> pull_requests_;
/// Internally maintained random number generator.
std::mt19937_64 gen_;
/// Try to Pull an object from one of its expected client locations. If there
/// are more client locations to try after this attempt, then this method
/// will try each of the other clients in succession, with a timeout between
/// each attempt. If the object is received or if the Pull is Canceled before
/// the timeout, then no more Pull requests for this object will be sent
/// to other node managers until TryPull is called again.
///
/// \param object_id The object's object id.
/// \return Void.
void TryPull(const ObjectID &object_id);
};
} // namespace ray
src/ray/object_manager/test/pull_manager_test.cc
+173
View File
@@ -0,0 +1,173 @@
#include "ray/object_manager/pull_manager.h"
#include "gtest/gtest.h"
#include "ray/common/test_util.h"
namespace ray {
class PullManagerTest : public ::testing::Test {
public:
PullManagerTest()
: self_node_id_(NodeID::FromRandom()),
object_is_local_(false),
num_send_pull_request_calls_(0),
num_restore_spilled_object_calls_(0),
fake_time_(0),
pull_manager_(self_node_id_,
[this](const ObjectID &object_id) { return object_is_local_; },
[this](const ObjectID &object_id, const NodeID &node_id) {
num_send_pull_request_calls_++;
},
[this](const ObjectID &, const std::string &,
std::function<void(const ray::Status &)>) {
num_restore_spilled_object_calls_++;
},
[this]() { return fake_time_; }, 10000) {}
NodeID self_node_id_;
bool object_is_local_;
int num_send_pull_request_calls_;
int num_restore_spilled_object_calls_;
double fake_time_;
PullManager pull_manager_;
};
TEST_F(PullManagerTest, TestStaleSubscription) {
ObjectID obj1 = ObjectID::FromRandom();
rpc::Address addr1;
ASSERT_EQ(pull_manager_.NumActiveRequests(), 0);
pull_manager_.Pull(obj1, addr1);
ASSERT_EQ(pull_manager_.NumActiveRequests(), 1);
std::unordered_set<NodeID> client_ids;
pull_manager_.OnLocationChange(obj1, client_ids, "");
// There are no client ids to pull from.
ASSERT_EQ(num_send_pull_request_calls_, 0);
ASSERT_EQ(num_restore_spilled_object_calls_, 0);
pull_manager_.CancelPull(obj1);
ASSERT_EQ(num_send_pull_request_calls_, 0);
ASSERT_EQ(num_restore_spilled_object_calls_, 0);
ASSERT_EQ(pull_manager_.NumActiveRequests(), 0);
client_ids.insert(NodeID::FromRandom());
pull_manager_.OnLocationChange(obj1, client_ids, "");
// Now we're getting a notification about an object that was already cancelled.
ASSERT_EQ(num_send_pull_request_calls_, 0);
ASSERT_EQ(num_restore_spilled_object_calls_, 0);
ASSERT_EQ(pull_manager_.NumActiveRequests(), 0);
}
TEST_F(PullManagerTest, TestRestoreSpilledObject) {
ObjectID obj1 = ObjectID::FromRandom();
rpc::Address addr1;
ASSERT_EQ(pull_manager_.NumActiveRequests(), 0);
pull_manager_.Pull(obj1, addr1);
ASSERT_EQ(pull_manager_.NumActiveRequests(), 1);
std::unordered_set<NodeID> client_ids;
pull_manager_.OnLocationChange(obj1, client_ids, "remote_url/foo/bar");
// client_ids is empty here, so there's nowhere to pull from.
ASSERT_EQ(num_send_pull_request_calls_, 0);
ASSERT_EQ(num_restore_spilled_object_calls_, 1);
client_ids.insert(NodeID::FromRandom());
pull_manager_.OnLocationChange(obj1, client_ids, "remote_url/foo/bar");
// The behavior is supposed to be to always restore the spilled object if possible (even
// if it exists elsewhere in the cluster).
ASSERT_EQ(num_send_pull_request_calls_, 0);
ASSERT_EQ(num_restore_spilled_object_calls_, 2);
pull_manager_.CancelPull(obj1);
ASSERT_EQ(pull_manager_.NumActiveRequests(), 0);
}
TEST_F(PullManagerTest, TestManyUpdates) {
ObjectID obj1 = ObjectID::FromRandom();
rpc::Address addr1;
ASSERT_EQ(pull_manager_.NumActiveRequests(), 0);
pull_manager_.Pull(obj1, addr1);
ASSERT_EQ(pull_manager_.NumActiveRequests(), 1);
std::unordered_set<NodeID> client_ids;
client_ids.insert(NodeID::FromRandom());
for (int i = 0; i < 100; i++) {
pull_manager_.OnLocationChange(obj1, client_ids, "");
}
ASSERT_EQ(num_send_pull_request_calls_, 100);
ASSERT_EQ(num_restore_spilled_object_calls_, 0);
pull_manager_.CancelPull(obj1);
ASSERT_EQ(pull_manager_.NumActiveRequests(), 0);
}
TEST_F(PullManagerTest, TestRetryTimer) {
ObjectID obj1 = ObjectID::FromRandom();
rpc::Address addr1;
ASSERT_EQ(pull_manager_.NumActiveRequests(), 0);
pull_manager_.Pull(obj1, addr1);
ASSERT_EQ(pull_manager_.NumActiveRequests(), 1);
std::unordered_set<NodeID> client_ids;
client_ids.insert(NodeID::FromRandom());
// We need to call OnLocationChange at least once, to population the list of nodes with
// the object.
pull_manager_.OnLocationChange(obj1, client_ids, "");
ASSERT_EQ(num_send_pull_request_calls_, 1);
ASSERT_EQ(num_restore_spilled_object_calls_, 0);
for (; fake_time_ <= 127 * 10; fake_time_ += 0.1) {
pull_manager_.Tick();
}
// Rapid set of location changes.
for (int i = 0; i < 127; i++) {
fake_time_ += 0.1;
pull_manager_.OnLocationChange(obj1, client_ids, "");
}
// We should make a pull request every tick (even if it's a duplicate to a node we're
// already pulling from).
// OnLocationChange also doesn't count towards the retry timer.
// To the casual observer, this may seem off-by-one, but this is due to floating point
// error (0.1 + 0.1 ... 10k times > 10 == True)
ASSERT_EQ(num_send_pull_request_calls_, 127 * 2);
ASSERT_EQ(num_restore_spilled_object_calls_, 0);
pull_manager_.CancelPull(obj1);
ASSERT_EQ(pull_manager_.NumActiveRequests(), 0);
}
TEST_F(PullManagerTest, TestBasic) {
ObjectID obj1 = ObjectID::FromRandom();
rpc::Address addr1;
ASSERT_EQ(pull_manager_.NumActiveRequests(), 0);
pull_manager_.Pull(obj1, addr1);
ASSERT_EQ(pull_manager_.NumActiveRequests(), 1);
std::unordered_set<NodeID> client_ids;
client_ids.insert(NodeID::FromRandom());
pull_manager_.OnLocationChange(obj1, client_ids, "");
ASSERT_EQ(num_send_pull_request_calls_, 1);
ASSERT_EQ(num_restore_spilled_object_calls_, 0);
pull_manager_.CancelPull(obj1);
ASSERT_EQ(pull_manager_.NumActiveRequests(), 0);
}
} // namespace ray
int main(int argc, char **argv) {
::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}
src/ray/raylet/main.cc
+3
View File
@@ -234,6 +234,9 @@ int main(int argc, char *argv[]) {
ray::ObjectManagerConfig object_manager_config;
object_manager_config.object_manager_port = object_manager_port;
object_manager_config.store_socket_name = store_socket_name;
object_manager_config.timer_freq_ms =
RayConfig::instance().object_manager_timer_freq_ms();
object_manager_config.pull_timeout_ms =
RayConfig::instance().object_manager_pull_timeout_ms();
object_manager_config.push_timeout_ms =