From fcd30444a8bf75e597d0b253c0b4c49ff12b357f Mon Sep 17 00:00:00 2001 From: Melih Elibol Date: Sat, 14 Apr 2018 17:08:19 -0700 Subject: [PATCH] Single Big Object Parallel Transfer. (#1827) * cache all object info from object added store notification. * Adds parallel transfer for big objects. * documentation and clean up. * compare objects... * merge buffer_state with chunk vec. Make separate buffer state for get and create. * use references for Get. Allow partial failure of Create. * single plasma client. * changes based on review. * update documentation and add parameters for object manager in main.cc. * review feedback. * use vector consturctor. * linting * remove profile visualizations. * test fixes. * linting. * kill specific pids and use less memory. * linting. * simplify tests. * Asynchronous IO for ObjectManager messages and object transfer. * Revert "Asynchronous IO for ObjectManager messages and object transfer." This reverts commit 4af43b159babc04daf80d1543e27c2cb46b7b19d. * update test configuration to reflect changes in #1891 * review feedback. * linting. --- src/ray/CMakeLists.txt | 2 +- .../object_manager/format/object_manager.fbs | 8 +- src/ray/object_manager/object_buffer_pool.cc | 187 +++++++++++++++++ src/ray/object_manager/object_buffer_pool.h | 196 +++++++++++++++++ src/ray/object_manager/object_directory.h | 2 +- src/ray/object_manager/object_manager.cc | 198 +++++++++--------- src/ray/object_manager/object_manager.h | 24 ++- .../object_store_client_pool.cc | 36 ---- .../object_manager/object_store_client_pool.h | 64 ------ .../test/object_manager_stress_test.cc | 45 ++-- .../test/object_manager_test.cc | 31 ++- src/ray/object_manager/transfer_queue.cc | 36 +--- src/ray/object_manager/transfer_queue.h | 64 ++---- src/ray/raylet/main.cc | 8 + src/ray/raylet/raylet.cc | 9 +- src/ray/test/run_object_manager_tests.sh | 5 +- 16 files changed, 608 insertions(+), 307 deletions(-) create mode 100644 src/ray/object_manager/object_buffer_pool.cc create mode 100644 src/ray/object_manager/object_buffer_pool.h delete mode 100644 src/ray/object_manager/object_store_client_pool.cc delete mode 100644 src/ray/object_manager/object_store_client_pool.h diff --git a/src/ray/CMakeLists.txt b/src/ray/CMakeLists.txt index e19e9d9bb..53edaf7c8 100644 --- a/src/ray/CMakeLists.txt +++ b/src/ray/CMakeLists.txt @@ -39,7 +39,7 @@ set(RAY_SRCS common/client_connection.cc object_manager/object_manager_client_connection.cc object_manager/connection_pool.cc - object_manager/object_store_client_pool.cc + object_manager/object_buffer_pool.cc object_manager/object_store_notification_manager.cc object_manager/object_directory.cc object_manager/transfer_queue.cc diff --git a/src/ray/object_manager/format/object_manager.fbs b/src/ray/object_manager/format/object_manager.fbs index 7d833d728..d47905287 100644 --- a/src/ray/object_manager/format/object_manager.fbs +++ b/src/ray/object_manager/format/object_manager.fbs @@ -10,8 +10,12 @@ enum MessageType:int { table PushRequestMessage { // The object ID being transferred. object_id: string; - // The size of the object being transferred. - object_size: ulong; + // The index of the chunk being transferred. + chunk_index: ulong; + // The total size of the object + metadata. + data_size: ulong; + // The metadata size. + metadata_size: ulong; } table PullRequestMessage { diff --git a/src/ray/object_manager/object_buffer_pool.cc b/src/ray/object_manager/object_buffer_pool.cc new file mode 100644 index 000000000..beab84736 --- /dev/null +++ b/src/ray/object_manager/object_buffer_pool.cc @@ -0,0 +1,187 @@ +#include "ray/object_manager/object_buffer_pool.h" + +namespace ray { + +ObjectBufferPool::ObjectBufferPool(const std::string &store_socket_name, + uint64_t chunk_size, int release_delay) + : chunk_size_(chunk_size) { + store_socket_name_ = store_socket_name; + ARROW_CHECK_OK(store_client_.Connect(store_socket_name_.c_str(), "", release_delay)); +} + +ObjectBufferPool::~ObjectBufferPool() { + // Abort everything in progress. + auto get_buf_state_copy = get_buffer_state_; + for (const auto &pair : get_buf_state_copy) { + AbortGet(pair.first); + } + auto create_buf_state_copy = create_buffer_state_; + for (const auto &pair : create_buf_state_copy) { + AbortCreate(pair.first); + } + RAY_CHECK(get_buffer_state_.empty()); + RAY_CHECK(create_buffer_state_.empty()); + ARROW_CHECK_OK(store_client_.Disconnect()); +} + +uint64_t ObjectBufferPool::GetNumChunks(uint64_t data_size) { + return (data_size + chunk_size_ - 1) / chunk_size_; +} + +uint64_t ObjectBufferPool::GetBufferLength(uint64_t chunk_index, uint64_t data_size) { + return (chunk_index + 1) * chunk_size_ > data_size ? data_size % chunk_size_ + : chunk_size_; +} + +std::pair ObjectBufferPool::GetChunk( + const ObjectID &object_id, uint64_t data_size, uint64_t metadata_size, + uint64_t chunk_index) { + std::lock_guard lock(pool_mutex_); + RAY_LOG(DEBUG) << "GetChunk " << object_id << " " << data_size << " " << metadata_size; + if (get_buffer_state_.count(object_id) == 0) { + plasma::ObjectBuffer object_buffer; + plasma::ObjectID plasma_id = ObjectID(object_id).to_plasma_id(); + ARROW_CHECK_OK(store_client_.Get(&plasma_id, 1, 0, &object_buffer)); + if (object_buffer.data == nullptr) { + RAY_LOG(ERROR) << "Failed to get object"; + return std::pair( + errored_chunk_, + ray::Status::IOError("Unable to obtain object chunk, object not local.")); + } + RAY_CHECK(object_buffer.metadata->data() == + object_buffer.data->data() + object_buffer.data->size()); + RAY_CHECK(data_size == static_cast(object_buffer.data->size() + + object_buffer.metadata_size)); + auto *data = const_cast(object_buffer.data->data()); + uint64_t num_chunks = GetNumChunks(data_size); + get_buffer_state_.emplace( + std::piecewise_construct, std::forward_as_tuple(object_id), + std::forward_as_tuple(BuildChunks(object_id, data, data_size))); + RAY_CHECK(get_buffer_state_[object_id].chunk_info.size() == num_chunks); + } + get_buffer_state_[object_id].references++; + return std::pair( + get_buffer_state_[object_id].chunk_info[chunk_index], ray::Status::OK()); +} + +void ObjectBufferPool::ReleaseGetChunk(const ObjectID &object_id, uint64_t chunk_index) { + std::lock_guard lock(pool_mutex_); + GetBufferState &buffer_state = get_buffer_state_[object_id]; + buffer_state.references--; + RAY_LOG(DEBUG) << "ReleaseBuffer " << object_id << " " << buffer_state.references; + if (buffer_state.references == 0) { + ARROW_CHECK_OK(store_client_.Release(ObjectID(object_id).to_plasma_id())); + get_buffer_state_.erase(object_id); + } +} + +void ObjectBufferPool::AbortGet(const ObjectID &object_id) { + std::lock_guard lock(pool_mutex_); + ARROW_CHECK_OK(store_client_.Release(ObjectID(object_id).to_plasma_id())); + get_buffer_state_.erase(object_id); +} + +std::pair ObjectBufferPool::CreateChunk( + const ObjectID &object_id, uint64_t data_size, uint64_t metadata_size, + uint64_t chunk_index) { + std::lock_guard lock(pool_mutex_); + RAY_LOG(DEBUG) << "CreateChunk " << object_id << " " << data_size << " " + << metadata_size; + if (create_buffer_state_.count(object_id) == 0) { + const plasma::ObjectID plasma_id = ObjectID(object_id).to_plasma_id(); + int64_t object_size = data_size - metadata_size; + // Try to create shared buffer. + std::shared_ptr data; + arrow::Status s = + store_client_.Create(plasma_id, object_size, NULL, metadata_size, &data); + std::vector buffer; + if (!s.ok()) { + // Create failed. The object may already exist locally. If something else went + // wrong, another chunk will succeed in creating the buffer, and this + // chunk will eventually make it here via pull requests. + return std::pair( + errored_chunk_, ray::Status::IOError(s.message())); + } + // Read object into store. + uint8_t *mutable_data = data->mutable_data(); + uint64_t num_chunks = GetNumChunks(data_size); + create_buffer_state_.emplace( + std::piecewise_construct, std::forward_as_tuple(object_id), + std::forward_as_tuple(BuildChunks(object_id, mutable_data, data_size))); + RAY_CHECK(create_buffer_state_[object_id].chunk_info.size() == num_chunks); + } + if (create_buffer_state_[object_id].chunk_state[chunk_index] != + CreateChunkState::AVAILABLE) { + // There can be only one reference to this chunk at any given time. + return std::pair( + errored_chunk_, + ray::Status::IOError("Chunk already referenced by another thread.")); + } + create_buffer_state_[object_id].chunk_state[chunk_index] = CreateChunkState::REFERENCED; + return std::pair( + create_buffer_state_[object_id].chunk_info[chunk_index], ray::Status::OK()); +} + +void ObjectBufferPool::AbortCreateChunk(const ObjectID &object_id, + const uint64_t chunk_index) { + std::lock_guard lock(pool_mutex_); + RAY_CHECK(create_buffer_state_[object_id].chunk_state[chunk_index] == + CreateChunkState::REFERENCED); + create_buffer_state_[object_id].chunk_state[chunk_index] = CreateChunkState::AVAILABLE; + if (create_buffer_state_[object_id].num_seals_remaining == + create_buffer_state_[object_id].chunk_state.size()) { + // If chunk_state is AVAILABLE at every chunk_index and + // num_seals_remaining == num_chunks, this is back to the initial state + // right before the first CreateChunk. + bool abort = true; + for (auto chunk_state : create_buffer_state_[object_id].chunk_state) { + abort &= chunk_state == CreateChunkState::AVAILABLE; + } + if (abort) { + AbortCreate(object_id); + } + } +} + +void ObjectBufferPool::SealChunk(const ObjectID &object_id, const uint64_t chunk_index) { + std::lock_guard lock(pool_mutex_); + RAY_CHECK(create_buffer_state_[object_id].chunk_state[chunk_index] == + CreateChunkState::REFERENCED); + create_buffer_state_[object_id].chunk_state[chunk_index] = CreateChunkState::SEALED; + create_buffer_state_[object_id].num_seals_remaining--; + RAY_LOG(DEBUG) << "SealChunk" << object_id << " " + << create_buffer_state_[object_id].num_seals_remaining; + if (create_buffer_state_[object_id].num_seals_remaining == 0) { + const plasma::ObjectID plasma_id = ObjectID(object_id).to_plasma_id(); + ARROW_CHECK_OK(store_client_.Seal(plasma_id)); + ARROW_CHECK_OK(store_client_.Release(plasma_id)); + create_buffer_state_.erase(object_id); + } +} + +void ObjectBufferPool::AbortCreate(const ObjectID &object_id) { + const plasma::ObjectID plasma_id = ObjectID(object_id).to_plasma_id(); + ARROW_CHECK_OK(store_client_.Release(plasma_id)); + ARROW_CHECK_OK(store_client_.Abort(plasma_id)); + create_buffer_state_.erase(object_id); +} + +std::vector ObjectBufferPool::BuildChunks( + const ObjectID &object_id, uint8_t *data, uint64_t data_size) { + uint64_t space_remaining = data_size; + std::vector chunks; + int64_t position = 0; + while (space_remaining) { + position = data_size - space_remaining; + if (space_remaining < chunk_size_) { + chunks.emplace_back(chunks.size(), data + position, space_remaining); + space_remaining = 0; + } else { + chunks.emplace_back(chunks.size(), data + position, chunk_size_); + space_remaining -= chunk_size_; + } + } + return chunks; +} + +} // namespace ray diff --git a/src/ray/object_manager/object_buffer_pool.h b/src/ray/object_manager/object_buffer_pool.h new file mode 100644 index 000000000..3edc1be30 --- /dev/null +++ b/src/ray/object_manager/object_buffer_pool.h @@ -0,0 +1,196 @@ +#ifndef RAY_OBJECT_MANAGER_OBJECT_BUFFER_POOL_H +#define RAY_OBJECT_MANAGER_OBJECT_BUFFER_POOL_H + +#include +#include +#include +#include + +#include +#include +#include + +#include "plasma/client.h" +#include "plasma/events.h" +#include "plasma/plasma.h" + +#include "ray/id.h" +#include "ray/status.h" + +namespace ray { + +/// \class ObjectBufferPool Exposes chunks of object buffers for use by the ObjectManager. +class ObjectBufferPool { + public: + /// Information needed to read or write an object chunk. + /// This is the structure returned whenever an object chunk is + /// accessed via Get and Create. + struct ChunkInfo { + ChunkInfo(uint64_t chunk_index, uint8_t *data, uint64_t buffer_length) + : chunk_index(chunk_index), data(data), buffer_length(buffer_length){}; + /// A pointer to the start position of this object chunk. + uint64_t chunk_index; + /// A pointer to the start position of this object chunk. + uint8_t *data; + /// The size of this object chunk. + uint64_t buffer_length; + }; + + /// Constructor. + /// + /// \param store_socket_name The socket name of the store to which plasma clients + /// connect. + /// \param chunk_size The chunk size into which objects are to be split. + /// \param release_delay The number of release calls before objects are released + /// from the store client (FIFO). + ObjectBufferPool(const std::string &store_socket_name, const uint64_t chunk_size, + const int release_delay); + + ~ObjectBufferPool(); + + /// This object cannot be copied due to pool_mutex. + RAY_DISALLOW_COPY_AND_ASSIGN(ObjectBufferPool); + + /// Computes the number of chunks needed to transfer an object and its metadata. + /// + /// \param data_size The size of the object + metadata. + /// \return The number of chunks into which the object will be split. + uint64_t GetNumChunks(uint64_t data_size); + + /// Computes the buffer length of a chunk of an object. + /// + /// \param chunk_index The chunk index for which to obtain the buffer length. + /// \param data_size The size of the object + metadata. + /// \return The buffer length of the chunk at chunk_index. + uint64_t GetBufferLength(uint64_t chunk_index, uint64_t data_size); + + /// Returns a chunk of an object at the given chunk_index. The object chunk serves + /// as the data that is to be written to a connection as part of sending an object to + /// a remote node. + /// + /// \param object_id The ObjectID. + /// \param data_size The sum of the object size and metadata size. + /// \param metadata_size The size of the metadata. + /// \param chunk_index The index of the chunk. + /// \return A pair consisting of a ChunkInfo and status of invoking this method. + /// An IOError status is returned if the Get call on the plasma store fails. + std::pair GetChunk( + const ObjectID &object_id, uint64_t data_size, uint64_t metadata_size, + uint64_t chunk_index); + + /// When a chunk is done being used as part of a get, this method releases the chunk. + /// If all chunks of an object are released, the object buffer will be released. + /// + /// \param object_id The object_id of the buffer to release. + /// \param chunk_index The index of the chunk. + void ReleaseGetChunk(const ObjectID &object_id, uint64_t chunk_index); + + /// Returns a chunk of an empty object at the given chunk_index. The object chunk + /// serves as the buffer that is to be written to by a connection receiving an object + /// from a remote node. Only one thread is permitted to create the object chunk at + /// chunk_index. Multiple threads attempting to create the same object chunk will + /// result in one succeeding. The ObjectManager is responsible for handling + /// create failures. This method will fail if it's invoked on a chunk_index on which + /// SealChunk has already been invoked. + /// + /// \param object_id The ObjectID. + /// \param data_size The sum of the object size and metadata size. + /// \param metadata_size The size of the metadata. + /// \param chunk_index The index of the chunk. + /// \return A pair consisting of ChunkInfo and status of invoking this method. + /// An IOError status is returned if object creation on the store client fails, + /// or if create is invoked consecutively on the same chunk + /// (with no intermediate AbortCreateChunk). + std::pair CreateChunk( + const ObjectID &object_id, uint64_t data_size, uint64_t metadata_size, + uint64_t chunk_index); + + /// Abort the create operation associated with a chunk at chunk_index. + /// This method will fail if it's invoked on a chunk_index on which + /// CreateChunk was not first invoked, or a chunk_index on which + /// SealChunk has already been invoked. + /// + /// \param object_id The ObjectID. + /// \param chunk_index The index of the chunk. + void AbortCreateChunk(const ObjectID &object_id, uint64_t chunk_index); + + /// Seal the object associated with a create operation. This is invoked whenever + /// a chunk is successfully written to. + /// This method will fail if it's invoked on a chunk_index on which + /// CreateChunk was not first invoked, or a chunk_index on which + /// SealChunk or AbortCreateChunk has already been invoked. + /// + /// \param object_id The ObjectID. + /// \param chunk_index The index of the chunk. + void SealChunk(const ObjectID &object_id, uint64_t chunk_index); + + private: + /// Abort the create operation associated with an object. This destroys the buffer + /// state, including create operations in progress for all chunks of the object. + void AbortCreate(const ObjectID &object_id); + + /// Abort the get operation associated with an object. + void AbortGet(const ObjectID &object_id); + + /// Splits an object into ceil(data_size/chunk_size) chunks, which will + /// either be read or written to in parallel. + std::vector BuildChunks(const ObjectID &object_id, uint8_t *data, + uint64_t data_size); + + /// Holds the state of a get buffer. + struct GetBufferState { + GetBufferState() {} + GetBufferState(std::vector chunk_info) : chunk_info(chunk_info) {} + /// A vector maintaining information about the chunks which comprise + /// an object. + std::vector chunk_info; + /// The number of references that currently rely on this buffer. + /// Once this reaches 0, the buffer is released and this object is erased + /// from get_buffer_state_. + uint64_t references = 0; + }; + + /// The state of a chunk associated with a create operation. + enum class CreateChunkState : uint { AVAILABLE = 0, REFERENCED, SEALED }; + + /// Holds the state of a create buffer. + struct CreateBufferState { + CreateBufferState() {} + CreateBufferState(std::vector chunk_info) + : chunk_info(chunk_info), + chunk_state(chunk_info.size(), CreateChunkState::AVAILABLE), + num_seals_remaining(chunk_info.size()) {} + /// A vector maintaining information about the chunks which comprise + /// an object. + std::vector chunk_info; + /// The state of each chunk, which is used to enforce strict state + /// transitions of each chunk. + std::vector chunk_state; + /// The number of chunks left to seal before the buffer is sealed. + uint64_t num_seals_remaining; + }; + + /// Returned when GetChunk or CreateChunk fails. + const ChunkInfo errored_chunk_ = {0, nullptr, 0}; + + /// Mutex on public methods for thread-safe operations on + /// get_buffer_state_, create_buffer_state_, and store_client_. + std::mutex pool_mutex_; + /// Determines the maximum chunk size to be transferred by a single thread. + const uint64_t chunk_size_; + /// The state of a buffer that's currently being used. + std::unordered_map + get_buffer_state_; + /// The state of a buffer that's currently being used. + std::unordered_map + create_buffer_state_; + + /// Plasma client pool. + plasma::PlasmaClient store_client_; + /// Socket name of plasma store. + std::string store_socket_name_; +}; + +} // namespace ray + +#endif // RAY_OBJECT_MANAGER_OBJECT_BUFFER_POOL_H diff --git a/src/ray/object_manager/object_directory.h b/src/ray/object_manager/object_directory.h index 686f67a39..df326125c 100644 --- a/src/ray/object_manager/object_directory.h +++ b/src/ray/object_manager/object_directory.h @@ -13,6 +13,7 @@ namespace ray { +/// Connection information for remote object managers. struct RemoteConnectionInfo { RemoteConnectionInfo() = default; RemoteConnectionInfo(const ClientID &id, const std::string &ip_address, @@ -23,7 +24,6 @@ struct RemoteConnectionInfo { uint16_t port; }; -/// Connection information for remote object managers. class ObjectDirectoryInterface { public: ObjectDirectoryInterface() = default; diff --git a/src/ray/object_manager/object_manager.cc b/src/ray/object_manager/object_manager.cc index f76aa4dc6..2b2845962 100644 --- a/src/ray/object_manager/object_manager.cc +++ b/src/ray/object_manager/object_manager.cc @@ -14,7 +14,10 @@ ObjectManager::ObjectManager(asio::io_service &main_service, : client_id_(gcs_client->client_table().GetLocalClientId()), object_directory_(new ObjectDirectory(gcs_client)), store_notification_(main_service, config.store_socket_name), - store_pool_(config.store_socket_name), + // release_delay of 2 * config.max_sends is to ensure the pool does not release + // an object prematurely whenever we reach the maximum number of sends. + buffer_pool_(config.store_socket_name, config.object_chunk_size, + /*release_delay=*/2 * config.max_sends), object_manager_service_(std::move(object_manager_service)), work_(*object_manager_service_), connection_pool_(), @@ -39,7 +42,10 @@ ObjectManager::ObjectManager(asio::io_service &main_service, std::unique_ptr od) : object_directory_(std::move(od)), store_notification_(main_service, config.store_socket_name), - store_pool_(config.store_socket_name), + // release_delay of 2 * config.max_sends is to ensure the pool does not release + // an object prematurely whenever we reach the maximum number of sends. + buffer_pool_(config.store_socket_name, config.object_chunk_size, + /*release_delay=*/2 * config.max_sends), object_manager_service_(std::move(object_manager_service)), work_(*object_manager_service_), connection_pool_(), @@ -59,12 +65,7 @@ ObjectManager::ObjectManager(asio::io_service &main_service, StartIOService(); } -ObjectManager::~ObjectManager() { - object_manager_service_->stop(); - for (int i = 0; i < num_threads_; ++i) { - io_threads_[i].join(); - } -} +ObjectManager::~ObjectManager() { StopIOService(); } void ObjectManager::StartIOService() { for (int i = 0; i < num_threads_; ++i) { @@ -74,6 +75,13 @@ void ObjectManager::StartIOService() { void ObjectManager::IOServiceLoop() { object_manager_service_->run(); } +void ObjectManager::StopIOService() { + object_manager_service_->stop(); + for (int i = 0; i < num_threads_; ++i) { + io_threads_[i].join(); + } +} + void ObjectManager::NotifyDirectoryObjectAdd(const ObjectInfoT &object_info) { ObjectID object_id = ObjectID::from_binary(object_info.object_id); local_objects_[object_id] = object_info; @@ -197,13 +205,28 @@ ray::Status ObjectManager::PullSendRequest(const ObjectID &object_id, } ray::Status ObjectManager::Push(const ObjectID &object_id, const ClientID &client_id) { - // TODO(hme): Cache this data in ObjectDirectory. - // Okay for now since the GCS client caches this data. + if (local_objects_.count(object_id) == 0) { + // TODO(hme): Do not retry indefinitely... + main_service_->post( + [this, object_id, client_id]() { RAY_CHECK_OK(Push(object_id, client_id)); }); + return ray::Status::OK(); + } + main_service_->dispatch([this, object_id, client_id]() { + // TODO(hme): Cache this data in ObjectDirectory. + // Okay for now since the GCS client caches this data. Status status = object_directory_->GetInformation( client_id, [this, object_id, client_id](const RemoteConnectionInfo &info) { - transfer_queue_.QueueSend(client_id, object_id, info); + ObjectInfoT object_info = local_objects_[object_id]; + uint64_t data_size = + static_cast(object_info.data_size + object_info.metadata_size); + uint64_t metadata_size = static_cast(object_info.metadata_size); + uint64_t num_chunks = buffer_pool_.GetNumChunks(data_size); + for (uint64_t chunk_index = 0; chunk_index < num_chunks; ++chunk_index) { + transfer_queue_.QueueSend(client_id, object_id, data_size, metadata_size, + chunk_index, info); + } RAY_CHECK_OK(DequeueTransfers()); }, [](const Status &status) { @@ -226,8 +249,9 @@ ray::Status ObjectManager::DequeueTransfers() { object_manager_service_->dispatch([this, req]() { RAY_LOG(DEBUG) << "DequeueSend " << client_id_ << " " << req.object_id << " " << num_transfers_send_ << "/" << config_.max_sends; - RAY_CHECK_OK( - ExecuteSendObject(req.object_id, req.client_id, req.connection_info)); + RAY_CHECK_OK(ExecuteSendObject(req.client_id, req.object_id, req.data_size, + req.metadata_size, req.chunk_index, + req.connection_info)); }); } else { std::atomic_fetch_sub(&num_transfers_send_, 1); @@ -248,7 +272,8 @@ ray::Status ObjectManager::DequeueTransfers() { object_manager_service_->dispatch([this, req]() { RAY_LOG(DEBUG) << "DequeueReceive " << client_id_ << " " << req.object_id << " " << num_transfers_receive_ << "/" << config_.max_receives; - RAY_CHECK_OK(ExecuteReceiveObject(req.client_id, req.object_id, req.object_size, + RAY_CHECK_OK(ExecuteReceiveObject(req.client_id, req.object_id, req.data_size, + req.metadata_size, req.chunk_index, req.conn)); }); } else { @@ -273,100 +298,75 @@ ray::Status ObjectManager::TransferCompleted(TransferQueue::TransferType type) { }; ray::Status ObjectManager::ExecuteSendObject( - const ObjectID &object_id, const ClientID &client_id, + const ClientID &client_id, const ObjectID &object_id, uint64_t data_size, + uint64_t metadata_size, uint64_t chunk_index, const RemoteConnectionInfo &connection_info) { + RAY_LOG(DEBUG) << "ExecuteSendObject " << client_id << " " << object_id << " " + << chunk_index; ray::Status status; std::shared_ptr conn; status = connection_pool_.GetSender(ConnectionPool::ConnectionType::TRANSFER, client_id, &conn); - if (!status.ok()) { - // TODO(hme): Keep track of retries, - // and only retry on object not local - // for now. - RAY_CHECK_OK(Push(object_id, conn->GetClientID())); - return Status::OK(); - } if (conn == nullptr) { conn = CreateSenderConnection(ConnectionPool::ConnectionType::TRANSFER, connection_info); connection_pool_.RegisterSender(ConnectionPool::ConnectionType::TRANSFER, client_id, conn); } - status = SendObjectHeaders(object_id, conn); - if (!status.ok()) { - RAY_CHECK_OK(Push(object_id, conn->GetClientID())); - return Status::OK(); - } + status = SendObjectHeaders(object_id, data_size, metadata_size, chunk_index, conn); return Status::OK(); } -ray::Status ObjectManager::SendObjectHeaders(const ObjectID &object_id_const, +ray::Status ObjectManager::SendObjectHeaders(const ObjectID &object_id, + uint64_t data_size, uint64_t metadata_size, + uint64_t chunk_index, std::shared_ptr conn) { - ObjectID object_id = ObjectID(object_id_const); - // Allocate and append the request to the transfer queue. - plasma::ObjectBuffer object_buffer; - plasma::ObjectID plasma_id = object_id.to_plasma_id(); - std::shared_ptr store_client = store_pool_.GetObjectStore(); - ARROW_CHECK_OK(store_client->Get(&plasma_id, 1, 0, &object_buffer)); - if (object_buffer.data == nullptr) { - RAY_LOG(ERROR) << "Failed to get object"; - // If the object wasn't locally available, exit immediately. If the object - // later appears locally, the requesting plasma manager should request the - // transfer again. - RAY_CHECK_OK( - connection_pool_.ReleaseSender(ConnectionPool::ConnectionType::TRANSFER, conn)); - store_pool_.ReleaseObjectStore(store_client); + std::pair chunk_status = + buffer_pool_.GetChunk(object_id, data_size, metadata_size, chunk_index); + ObjectBufferPool::ChunkInfo chunk_info = chunk_status.first; + + if (!chunk_status.second.ok()) { + // This is the first thread to invoke GetChunk => Get failed on the + // plasma client. + // No reference is acquired for this chunk, so no need to release the chunk. + // TODO(hme): Retry send here? If so, store RemoteConnectionInfo in SenderConnection. RAY_CHECK_OK(TransferCompleted(TransferQueue::TransferType::SEND)); - return ray::Status::IOError( - "Unable to transfer object to requesting plasma manager, object not local."); + return chunk_status.second; } - RAY_CHECK(object_buffer.metadata->data() == - object_buffer.data->data() + object_buffer.data->size()); - - TransferQueue::SendContext context; - context.client_id = conn->GetClientID(); - context.object_id = object_id; - context.object_size = static_cast(object_buffer.data->size()); - context.data = const_cast(object_buffer.data->data()); - UniqueID context_id = transfer_queue_.AddContext(context); - // Create buffer. flatbuffers::FlatBufferBuilder fbb; // TODO(hme): use to_flatbuf auto message = object_manager_protocol::CreatePushRequestMessage( - fbb, fbb.CreateString(object_id.binary()), context.object_size); + fbb, fbb.CreateString(object_id.binary()), chunk_index, data_size, metadata_size); fbb.Finish(message); ray::Status status = conn->WriteMessage(object_manager_protocol::MessageType_PushRequest, fbb.GetSize(), fbb.GetBufferPointer()); RAY_CHECK_OK(status); - // TODO(hme): Make this async. - return SendObjectData(conn, context_id, store_client); + return SendObjectData(object_id, chunk_info, conn); } -ray::Status ObjectManager::SendObjectData( - std::shared_ptr conn, const UniqueID &context_id, - std::shared_ptr store_client) { - TransferQueue::SendContext context = transfer_queue_.GetContext(context_id); +ray::Status ObjectManager::SendObjectData(const ObjectID &object_id, + const ObjectBufferPool::ChunkInfo &chunk_info, + std::shared_ptr conn) { + // TransferQueue::SendContext context = transfer_queue_.GetContext(context_id); boost::system::error_code ec; std::vector buffer; - buffer.push_back(asio::buffer(context.data, context.object_size)); + buffer.push_back(asio::buffer(chunk_info.data, chunk_info.buffer_length)); conn->WriteBuffer(buffer, ec); ray::Status status = ray::Status::OK(); if (ec.value() != 0) { - // push failed. - // TODO(hme): Trash sender. + // Push failed. Deal with partial objects on the receiving end. + // TODO(hme): Try to invoke disconnect on sender connection, then remove it. status = ray::Status::IOError(ec.message()); } // Do this regardless of whether it failed or succeeded. - ARROW_CHECK_OK(store_client->Release(context.object_id.to_plasma_id())); - store_pool_.ReleaseObjectStore(store_client); + buffer_pool_.ReleaseGetChunk(object_id, chunk_info.chunk_index); RAY_CHECK_OK( connection_pool_.ReleaseSender(ConnectionPool::ConnectionType::TRANSFER, conn)); - RAY_CHECK_OK(transfer_queue_.RemoveContext(context_id)); - RAY_LOG(DEBUG) << "SendCompleted " << client_id_ << " " << context.object_id << " " + RAY_LOG(DEBUG) << "SendCompleted " << client_id_ << " " << object_id << " " << num_transfers_send_ << "/" << config_.max_sends; RAY_CHECK_OK(TransferCompleted(TransferQueue::TransferType::SEND)); return status; @@ -469,45 +469,53 @@ void ObjectManager::ReceivePushRequest(std::shared_ptr conn auto object_header = flatbuffers::GetRoot(message); ObjectID object_id = ObjectID::from_binary(object_header->object_id()->str()); - int64_t object_size = (int64_t)object_header->object_size(); - transfer_queue_.QueueReceive(conn->GetClientID(), object_id, object_size, conn); + uint64_t chunk_index = object_header->chunk_index(); + uint64_t data_size = object_header->data_size(); + uint64_t metadata_size = object_header->metadata_size(); + transfer_queue_.QueueReceive(conn->GetClientID(), object_id, data_size, metadata_size, + chunk_index, conn); + RAY_LOG(DEBUG) << "ReceivePushRequest " << conn->GetClientID() << " " << object_id + << " " << chunk_index; RAY_CHECK_OK(DequeueTransfers()); } ray::Status ObjectManager::ExecuteReceiveObject( - const ClientID &client_id, const ObjectID &object_id, uint64_t object_size, + const ClientID &client_id, const ObjectID &object_id, uint64_t data_size, + uint64_t metadata_size, uint64_t chunk_index, std::shared_ptr conn) { - boost::system::error_code ec; - int64_t metadata_size = 0; - const plasma::ObjectID plasma_id = ObjectID(object_id).to_plasma_id(); - // Try to create shared buffer. - std::shared_ptr data; - std::shared_ptr store_client = store_pool_.GetObjectStore(); - arrow::Status s = - store_client->Create(plasma_id, object_size, NULL, metadata_size, &data); - std::vector buffer; - if (s.ok()) { - // Read object into store. - uint8_t *mutable_data = data->mutable_data(); - buffer.push_back(asio::buffer(mutable_data, object_size)); + RAY_LOG(DEBUG) << "ExecuteReceiveObject " << client_id << " " << object_id << " " + << chunk_index; + + std::pair chunk_status = + buffer_pool_.CreateChunk(object_id, data_size, metadata_size, chunk_index); + ObjectBufferPool::ChunkInfo chunk_info = chunk_status.first; + if (chunk_status.second.ok()) { + // Avoid handling this chunk if it's already being handled by another process. + std::vector buffer; + buffer.push_back(asio::buffer(chunk_info.data, chunk_info.buffer_length)); + boost::system::error_code ec; conn->ReadBuffer(buffer, ec); - if (!ec.value()) { - ARROW_CHECK_OK(store_client->Seal(plasma_id)); - ARROW_CHECK_OK(store_client->Release(plasma_id)); + if (ec.value() == 0) { + buffer_pool_.SealChunk(object_id, chunk_index); } else { - ARROW_CHECK_OK(store_client->Release(plasma_id)); - ARROW_CHECK_OK(store_client->Abort(plasma_id)); - RAY_LOG(ERROR) << "Receive Failed"; + buffer_pool_.AbortCreateChunk(object_id, chunk_index); + // TODO(hme): This chunk failed, so create a pull request for this chunk. } } else { - RAY_LOG(ERROR) << "Buffer Create Failed: " << s.message(); + RAY_LOG(ERROR) << "Buffer Create Failed: " << chunk_status.second.message(); // Read object into empty buffer. - std::vector mutable_data; - mutable_data.resize(object_size + metadata_size); - buffer.push_back(asio::buffer(mutable_data, object_size + metadata_size)); + uint64_t buffer_length = buffer_pool_.GetBufferLength(chunk_index, data_size); + std::vector mutable_vec; + mutable_vec.resize(buffer_length); + std::vector buffer; + buffer.push_back(asio::buffer(mutable_vec, buffer_length)); + boost::system::error_code ec; conn->ReadBuffer(buffer, ec); + if (ec.value() != 0) { + RAY_LOG(ERROR) << ec.message(); + } + // TODO(hme): If the object isn't local, create a pull request for this chunk. } - store_pool_.ReleaseObjectStore(store_client); conn->ProcessMessages(); RAY_LOG(DEBUG) << "ReceiveCompleted " << client_id_ << " " << object_id << " " << num_transfers_receive_ << "/" << config_.max_receives; diff --git a/src/ray/object_manager/object_manager.h b/src/ray/object_manager/object_manager.h index 1201669df..ed3b7c27a 100644 --- a/src/ray/object_manager/object_manager.h +++ b/src/ray/object_manager/object_manager.h @@ -22,9 +22,9 @@ #include "ray/object_manager/connection_pool.h" #include "ray/object_manager/format/object_manager_generated.h" +#include "ray/object_manager/object_buffer_pool.h" #include "ray/object_manager/object_directory.h" #include "ray/object_manager/object_manager_client_connection.h" -#include "ray/object_manager/object_store_client_pool.h" #include "ray/object_manager/object_store_notification_manager.h" #include "ray/object_manager/transfer_queue.h" @@ -38,6 +38,8 @@ struct ObjectManagerConfig { int max_sends = 2; /// Maximum number of receives allowed. int max_receives = 2; + /// Object chunk size, in bytes + uint64_t object_chunk_size = std::pow(10, 8); // TODO(hme): Implement num retries (to avoid infinite retries). std::string store_socket_name; }; @@ -153,7 +155,7 @@ class ObjectManager { ObjectManagerConfig config_; std::unique_ptr object_directory_; ObjectStoreNotificationManager store_notification_; - ObjectStoreClientPool store_pool_; + ObjectBufferPool buffer_pool_; /// An io service for creating connections to other object managers. /// This runs on a thread pool. @@ -252,19 +254,22 @@ class ObjectManager { /// Begin executing a send. /// Executes on object_manager_service_ thread pool. - ray::Status ExecuteSendObject(const ObjectID &object_id, const ClientID &client_id, + ray::Status ExecuteSendObject(const ClientID &client_id, const ObjectID &object_id, + uint64_t data_size, uint64_t metadata_size, + uint64_t chunk_index, const RemoteConnectionInfo &connection_info); /// This method synchronously sends the object id and object size /// to the remote object manager. /// Executes on object_manager_service_ thread pool. - ray::Status SendObjectHeaders(const ObjectID &object_id, - std::shared_ptr client); + ray::Status SendObjectHeaders(const ObjectID &object_id, uint64_t data_size, + uint64_t metadata_size, uint64_t chunk_index, + std::shared_ptr conn); /// This method initiates the actual object transfer. /// Executes on object_manager_service_ thread pool. - ray::Status SendObjectData(std::shared_ptr conn, - const UniqueID &context_id, - std::shared_ptr store_client); + ray::Status SendObjectData(const ObjectID &object_id, + const ObjectBufferPool::ChunkInfo &chunk_info, + std::shared_ptr conn); /// Invoked when a remote object manager pushes an object to this object manager. /// This will queue the receive. @@ -272,7 +277,8 @@ class ObjectManager { const uint8_t *message); /// Execute a receive that was in the queue. ray::Status ExecuteReceiveObject(const ClientID &client_id, const ObjectID &object_id, - uint64_t object_size, + uint64_t data_size, uint64_t metadata_size, + uint64_t chunk_index, std::shared_ptr conn); /// Handles receiving a pull request message. diff --git a/src/ray/object_manager/object_store_client_pool.cc b/src/ray/object_manager/object_store_client_pool.cc deleted file mode 100644 index 7534bf6e9..000000000 --- a/src/ray/object_manager/object_store_client_pool.cc +++ /dev/null @@ -1,36 +0,0 @@ -#include "object_store_client_pool.h" - -namespace ray { - -ObjectStoreClientPool::ObjectStoreClientPool(const std::string &store_socket_name) - : store_socket_name_(store_socket_name) {} - -ObjectStoreClientPool::~ObjectStoreClientPool() { - for (const auto &client : clients) { - ARROW_CHECK_OK(client->Disconnect()); - } -} - -std::shared_ptr ObjectStoreClientPool::GetObjectStore() { - std::lock_guard lock(pool_mutex); - if (available_clients.empty()) { - Add(); - } - std::shared_ptr client = available_clients.back(); - available_clients.pop_back(); - return client; -} - -void ObjectStoreClientPool::ReleaseObjectStore( - std::shared_ptr client) { - std::lock_guard lock(pool_mutex); - available_clients.push_back(client); -} - -void ObjectStoreClientPool::Add() { - clients.emplace_back(new plasma::PlasmaClient()); - ARROW_CHECK_OK(clients.back()->Connect(store_socket_name_.c_str(), "", - PLASMA_DEFAULT_RELEASE_DELAY)); - available_clients.push_back(clients.back()); -} -} // namespace ray diff --git a/src/ray/object_manager/object_store_client_pool.h b/src/ray/object_manager/object_store_client_pool.h deleted file mode 100644 index b6a611af5..000000000 --- a/src/ray/object_manager/object_store_client_pool.h +++ /dev/null @@ -1,64 +0,0 @@ -#ifndef RAY_OBJECT_MANAGER_OBJECT_STORE_CLIENT_POOL_H -#define RAY_OBJECT_MANAGER_OBJECT_STORE_CLIENT_POOL_H - -#include -#include -#include -#include - -#include -#include -#include - -#include "plasma/client.h" -#include "plasma/events.h" -#include "plasma/plasma.h" - -#include "object_directory.h" -#include "ray/id.h" -#include "ray/status.h" - -namespace ray { - -/// \class ObjectStoreClientPool -/// -/// Provides connections to the object store. Enables concurrent communication with -/// the object store. -class ObjectStoreClientPool { - public: - /// Constructor. - /// - /// \param store_socket_name The object store socket name. - ObjectStoreClientPool(const std::string &store_socket_name); - - ~ObjectStoreClientPool(); - - /// This object cannot be copied due to pool_mutex. - RAY_DISALLOW_COPY_AND_ASSIGN(ObjectStoreClientPool); - - /// Provides a connection to the object store from the object store pool. - /// This removes the object store client from the pool of available clients. - /// - /// \return A connection to the object store. - std::shared_ptr GetObjectStore(); - - /// Releases a client object and puts it back into the object store pool - /// for reuse. - /// Once a client is released, it is assumed that it is not being used. - /// \param client The client to return. - /// \param client - void ReleaseObjectStore(std::shared_ptr client); - - private: - /// Adds a client to the client pool and mark it as available. - void Add(); - - std::mutex pool_mutex; - std::vector> available_clients; - std::vector> clients; - std::string store_socket_name_; -}; - -} // namespace ray - -#endif // RAY_OBJECT_MANAGER_OBJECT_STORE_CLIENT_POOL_H diff --git a/src/ray/object_manager/test/object_manager_stress_test.cc b/src/ray/object_manager/test/object_manager_stress_test.cc index 9c013550b..5b3e168db 100644 --- a/src/ray/object_manager/test/object_manager_stress_test.cc +++ b/src/ray/object_manager/test/object_manager_stress_test.cc @@ -96,14 +96,25 @@ class TestObjectManagerBase : public ::testing::Test { std::string StartStore(const std::string &id) { std::string store_id = "/tmp/store"; store_id = store_id + id; + std::string store_pid = store_id + ".pid"; std::string plasma_command = store_executable + " -m 1000000000 -s " + store_id + - " 1> /dev/null 2> /dev/null &"; + " 1> /dev/null 2> /dev/null &" + " echo $! > " + + store_pid; + RAY_LOG(DEBUG) << plasma_command; int ec = system(plasma_command.c_str()); RAY_CHECK(ec == 0); + sleep(1); return store_id; } + void StopStore(std::string store_id) { + std::string store_pid = store_id + ".pid"; + std::string kill_1 = "kill -9 `cat " + store_pid + "`"; + int s = system(kill_1.c_str()); + ASSERT_TRUE(!s); + } + void SetUp() { flushall_redis(); @@ -111,30 +122,32 @@ class TestObjectManagerBase : public ::testing::Test { object_manager_service_2.reset(new boost::asio::io_service()); // start store - std::string store_sock_1 = StartStore("1"); - std::string store_sock_2 = StartStore("2"); + store_id_1 = StartStore(UniqueID::from_random().hex()); + store_id_2 = StartStore(UniqueID::from_random().hex()); // start first server gcs_client_1 = std::shared_ptr(new gcs::AsyncGcsClient()); ObjectManagerConfig om_config_1; - om_config_1.store_socket_name = store_sock_1; + om_config_1.store_socket_name = store_id_1; om_config_1.max_sends = 2; om_config_1.max_receives = 2; + om_config_1.object_chunk_size = static_cast(std::pow(10, 3)); server1.reset(new MockServer(main_service, std::move(object_manager_service_1), om_config_1, gcs_client_1)); // start second server gcs_client_2 = std::shared_ptr(new gcs::AsyncGcsClient()); ObjectManagerConfig om_config_2; - om_config_2.store_socket_name = store_sock_2; + om_config_2.store_socket_name = store_id_2; om_config_2.max_sends = 2; om_config_2.max_receives = 2; + om_config_2.object_chunk_size = static_cast(std::pow(10, 3)); server2.reset(new MockServer(main_service, std::move(object_manager_service_2), om_config_2, gcs_client_2)); // connect to stores. - ARROW_CHECK_OK(client1.Connect(store_sock_1, "", PLASMA_DEFAULT_RELEASE_DELAY)); - ARROW_CHECK_OK(client2.Connect(store_sock_2, "", PLASMA_DEFAULT_RELEASE_DELAY)); + ARROW_CHECK_OK(client1.Connect(store_id_1, "", PLASMA_DEFAULT_RELEASE_DELAY)); + ARROW_CHECK_OK(client2.Connect(store_id_2, "", PLASMA_DEFAULT_RELEASE_DELAY)); } void TearDown() { @@ -145,8 +158,8 @@ class TestObjectManagerBase : public ::testing::Test { this->server1.reset(); this->server2.reset(); - int s = system("killall plasma_store &"); - ASSERT_TRUE(!s); + StopStore(store_id_1); + StopStore(store_id_2); } ObjectID WriteDataToClient(plasma::PlasmaClient &client, int64_t data_size) { @@ -179,6 +192,9 @@ class TestObjectManagerBase : public ::testing::Test { plasma::PlasmaClient client2; std::vector v1; std::vector v2; + + std::string store_id_1; + std::string store_id_2; }; class StressTestObjectManager : public TestObjectManagerBase { @@ -257,7 +273,7 @@ class StressTestObjectManager : public TestObjectManagerBase { async_loop_index += 1; if ((uint)async_loop_index < async_loop_patterns.size()) { TransferPattern pattern = async_loop_patterns[async_loop_index]; - TransferTestExecute(100, 100, pattern); + TransferTestExecute(100, 3 * std::pow(10, 3) - 1, pattern); } else { main_service.stop(); } @@ -279,11 +295,14 @@ class StressTestObjectManager : public TestObjectManagerBase { void CompareObjects(ObjectID &object_id_1, ObjectID &object_id_2) { plasma::ObjectBuffer object_buffer_1 = GetObject(client1, object_id_1); - plasma::ObjectBuffer object_buffer_2 = GetObject(client1, object_id_1); + plasma::ObjectBuffer object_buffer_2 = GetObject(client2, object_id_2); uint8_t *data_1 = const_cast(object_buffer_1.data->data()); uint8_t *data_2 = const_cast(object_buffer_2.data->data()); - ASSERT_EQ(object_buffer_1.data->size(), object_buffer_2.data->size()); - for (int i = -1; ++i < object_buffer_1.data->size();) { + ASSERT_EQ(object_buffer_1.data->size(), object_buffer_2.data_size); + ASSERT_EQ(object_buffer_1.metadata_size, object_buffer_2.metadata_size); + int64_t total_size = object_buffer_1.data->size() + object_buffer_1.metadata_size; + RAY_LOG(DEBUG) << "total_size " << total_size; + for (int i = -1; ++i < total_size;) { ASSERT_TRUE(data_1[i] == data_2[i]); } } diff --git a/src/ray/object_manager/test/object_manager_test.cc b/src/ray/object_manager/test/object_manager_test.cc index 0fa059662..d35969588 100644 --- a/src/ray/object_manager/test/object_manager_test.cc +++ b/src/ray/object_manager/test/object_manager_test.cc @@ -87,14 +87,24 @@ class TestObjectManager : public ::testing::Test { std::string StartStore(const std::string &id) { std::string store_id = "/tmp/store"; store_id = store_id + id; + std::string store_pid = store_id + ".pid"; std::string plasma_command = store_executable + " -m 1000000000 -s " + store_id + - " 1> /dev/null 2> /dev/null &"; + " 1> /dev/null 2> /dev/null &" + " echo $! > " + + store_pid; + RAY_LOG(DEBUG) << plasma_command; int ec = system(plasma_command.c_str()); RAY_CHECK(ec == 0); + sleep(1); return store_id; } + void StopStore(std::string store_id) { + std::string store_pid = store_id + ".pid"; + std::string kill_1 = "kill -9 `cat " + store_pid + "`"; + ASSERT_TRUE(!system(kill_1.c_str())); + } + void SetUp() { flushall_redis(); @@ -102,26 +112,26 @@ class TestObjectManager : public ::testing::Test { object_manager_service_2.reset(new boost::asio::io_service()); // start store - std::string store_sock_1 = StartStore("1"); - std::string store_sock_2 = StartStore("2"); + store_id_1 = StartStore(UniqueID::from_random().hex()); + store_id_2 = StartStore(UniqueID::from_random().hex()); // start first server gcs_client_1 = std::shared_ptr(new gcs::AsyncGcsClient()); ObjectManagerConfig om_config_1; - om_config_1.store_socket_name = store_sock_1; + om_config_1.store_socket_name = store_id_1; server1.reset(new MockServer(main_service, std::move(object_manager_service_1), om_config_1, gcs_client_1)); // start second server gcs_client_2 = std::shared_ptr(new gcs::AsyncGcsClient()); ObjectManagerConfig om_config_2; - om_config_2.store_socket_name = store_sock_2; + om_config_2.store_socket_name = store_id_2; server2.reset(new MockServer(main_service, std::move(object_manager_service_2), om_config_2, gcs_client_2)); // connect to stores. - ARROW_CHECK_OK(client1.Connect(store_sock_1, "", PLASMA_DEFAULT_RELEASE_DELAY)); - ARROW_CHECK_OK(client2.Connect(store_sock_2, "", PLASMA_DEFAULT_RELEASE_DELAY)); + ARROW_CHECK_OK(client1.Connect(store_id_1, "", PLASMA_DEFAULT_RELEASE_DELAY)); + ARROW_CHECK_OK(client2.Connect(store_id_2, "", PLASMA_DEFAULT_RELEASE_DELAY)); } void TearDown() { @@ -132,8 +142,8 @@ class TestObjectManager : public ::testing::Test { this->server1.reset(); this->server2.reset(); - int s = system("killall plasma_store &"); - ASSERT_TRUE(!s); + StopStore(store_id_1); + StopStore(store_id_2); } ObjectID WriteDataToClient(plasma::PlasmaClient &client, int64_t data_size) { @@ -166,6 +176,9 @@ class TestObjectManager : public ::testing::Test { plasma::PlasmaClient client2; std::vector v1; std::vector v2; + + std::string store_id_1; + std::string store_id_2; }; class TestObjectManagerCommands : public TestObjectManager { diff --git a/src/ray/object_manager/transfer_queue.cc b/src/ray/object_manager/transfer_queue.cc index 1d0dc0e33..9e25c1fed 100644 --- a/src/ray/object_manager/transfer_queue.cc +++ b/src/ray/object_manager/transfer_queue.cc @@ -3,9 +3,10 @@ namespace ray { void TransferQueue::QueueSend(const ClientID &client_id, const ObjectID &object_id, - const RemoteConnectionInfo &info) { - WriteLock guard(send_mutex); - SendRequest req = {client_id, object_id, info}; + uint64_t data_size, uint64_t metadata_size, + uint64_t chunk_index, const RemoteConnectionInfo &info) { + std::unique_lock guard(send_mutex_); + SendRequest req = {client_id, object_id, data_size, metadata_size, chunk_index, info}; // TODO(hme): Use a set to speed this up. if (std::find(send_queue_.begin(), send_queue_.end(), req) != send_queue_.end()) { // already queued. @@ -15,10 +16,12 @@ void TransferQueue::QueueSend(const ClientID &client_id, const ObjectID &object_ } void TransferQueue::QueueReceive(const ClientID &client_id, const ObjectID &object_id, - uint64_t object_size, + uint64_t data_size, uint64_t metadata_size, + uint64_t chunk_index, std::shared_ptr conn) { - WriteLock guard(receive_mutex); - ReceiveRequest req = {client_id, object_id, object_size, conn}; + std::unique_lock guard(receive_mutex_); + ReceiveRequest req = {client_id, object_id, data_size, + metadata_size, chunk_index, conn}; if (std::find(receive_queue_.begin(), receive_queue_.end(), req) != receive_queue_.end()) { // already queued. @@ -28,7 +31,7 @@ void TransferQueue::QueueReceive(const ClientID &client_id, const ObjectID &obje } bool TransferQueue::DequeueSendIfPresent(TransferQueue::SendRequest *send_ptr) { - WriteLock guard(send_mutex); + std::unique_lock guard(send_mutex_); if (send_queue_.empty()) { return false; } @@ -38,7 +41,7 @@ bool TransferQueue::DequeueSendIfPresent(TransferQueue::SendRequest *send_ptr) { } bool TransferQueue::DequeueReceiveIfPresent(TransferQueue::ReceiveRequest *receive_ptr) { - WriteLock guard(receive_mutex); + std::unique_lock guard(receive_mutex_); if (receive_queue_.empty()) { return false; } @@ -47,21 +50,4 @@ bool TransferQueue::DequeueReceiveIfPresent(TransferQueue::ReceiveRequest *recei return true; } -UniqueID TransferQueue::AddContext(SendContext &context) { - WriteLock guard(context_mutex); - UniqueID id = UniqueID::from_random(); - send_context_set_.emplace(id, context); - return id; -} - -TransferQueue::SendContext &TransferQueue::GetContext(const UniqueID &id) { - ReadLock guard(context_mutex); - return send_context_set_[id]; -} - -ray::Status TransferQueue::RemoveContext(const UniqueID &id) { - WriteLock guard(context_mutex); - send_context_set_.erase(id); - return Status::OK(); -} } // namespace ray diff --git a/src/ray/object_manager/transfer_queue.h b/src/ray/object_manager/transfer_queue.h index 8570d1ec8..6528a91c1 100644 --- a/src/ray/object_manager/transfer_queue.h +++ b/src/ray/object_manager/transfer_queue.h @@ -26,21 +26,17 @@ class TransferQueue { public: enum TransferType { SEND = 1, RECEIVE }; - /// Context maintained during an object send. - struct SendContext { - ClientID client_id; - ObjectID object_id; - uint64_t object_size; - uint8_t *data; - }; - /// The structure used in the send queue. struct SendRequest { ClientID client_id; ObjectID object_id; + uint64_t data_size; + uint64_t metadata_size; + uint64_t chunk_index; RemoteConnectionInfo connection_info; bool operator==(const SendRequest &rhs) const { - return client_id == rhs.client_id && object_id == rhs.object_id; + return client_id == rhs.client_id && object_id == rhs.object_id && + chunk_index == rhs.chunk_index; } }; @@ -48,10 +44,14 @@ class TransferQueue { struct ReceiveRequest { ClientID client_id; ObjectID object_id; - uint64_t object_size; + uint64_t data_size; + uint64_t metadata_size; + uint64_t chunk_index; std::shared_ptr conn; bool operator==(const ReceiveRequest &rhs) const { - return client_id == rhs.client_id && object_id == rhs.object_id; + return client_id == rhs.client_id && object_id == rhs.object_id && + chunk_index == rhs.chunk_index; + ; } }; @@ -61,11 +61,13 @@ class TransferQueue { /// /// \param client_id The ClientID to which the object needs to be sent. /// \param object_id The ObjectID of the object to be sent. - void QueueSend(const ClientID &client_id, const ObjectID &object_id, + void QueueSend(const ClientID &client_id, const ObjectID &object_id, uint64_t data_size, + uint64_t metadata_size, uint64_t chunk_index, const RemoteConnectionInfo &info); /// If send_queue_ is not empty, removes a SendRequest from send_queue_ and assigns /// it to send_ptr. The queue is FIFO. + /// /// \param send_ptr A pointer to an empty SendRequest. /// \return A bool indicating whether the queue was empty at the time this method /// was invoked. @@ -76,51 +78,27 @@ class TransferQueue { /// \param client_id The ClientID from which the object is being received. /// \param object_id The ObjectID of the object to be received. void QueueReceive(const ClientID &client_id, const ObjectID &object_id, - uint64_t object_size, std::shared_ptr conn); + uint64_t data_size, uint64_t metadata_size, uint64_t chunk_index, + std::shared_ptr conn); /// If receive_queue_ is not empty, removes a ReceiveRequest from receive_queue_ and - /// assigns - /// it to receive_ptr. The queue is FIFO. + /// assigns it to receive_ptr. The queue is FIFO. + /// /// \param receive_ptr A pointer to an empty ReceiveRequest. /// \return A bool indicating whether the queue was empty at the time this method /// was invoked. bool DequeueReceiveIfPresent(TransferQueue::ReceiveRequest *receive_ptr); - /// Maintain ownership over SendContext for sends in transit. - /// - /// \param context The context to maintain. - /// \return A unique identifier identifying the context that was added. - UniqueID AddContext(SendContext &context); - - /// Gets the SendContext associated with the given id. - /// - /// \param id The unique identifier of the context. - /// \return The context. - SendContext &GetContext(const UniqueID &id); - - /// Removes the context associated with the given id. - /// - /// \param id The unique identifier of the context. - /// \return The status of invoking this method. - ray::Status RemoveContext(const UniqueID &id); - /// This object cannot be copied for thread-safety. RAY_DISALLOW_COPY_AND_ASSIGN(TransferQueue); private: - // TODO(hme): make this a shared mutex. - typedef std::mutex Lock; - typedef std::unique_lock WriteLock; - // TODO(hme): make this a shared lock. - typedef std::unique_lock ReadLock; - Lock send_mutex; - Lock receive_mutex; - Lock context_mutex; - + std::mutex send_mutex_; + std::mutex receive_mutex_; std::deque send_queue_; std::deque receive_queue_; - std::unordered_map send_context_set_; }; + } // namespace ray #endif // RAY_OBJECT_MANAGER_TRANSFER_QUEUE_H diff --git a/src/ray/raylet/main.cc b/src/ray/raylet/main.cc index ebdc9a3f8..65dbbe725 100644 --- a/src/ray/raylet/main.cc +++ b/src/ray/raylet/main.cc @@ -47,6 +47,14 @@ int main(int argc, char *argv[]) { // Configuration for the object manager. ray::ObjectManagerConfig object_manager_config; object_manager_config.store_socket_name = store_socket_name; + // Time out in milliseconds to wait before retrying a failed pull. + object_manager_config.pull_timeout_ms = 100; + // Maximum number of sends allowed. + object_manager_config.max_sends = 2; + // Maximum number of receives allowed. + object_manager_config.max_receives = 2; + // Object chunk size, in bytes. + object_manager_config.object_chunk_size = static_cast(std::pow(10, 8)); // initialize mock gcs & object directory auto gcs_client = std::make_shared(); diff --git a/src/ray/raylet/raylet.cc b/src/ray/raylet/raylet.cc index 98ba2747e..a6e001cfe 100644 --- a/src/ray/raylet/raylet.cc +++ b/src/ray/raylet/raylet.cc @@ -37,16 +37,13 @@ Raylet::Raylet(boost::asio::io_service &main_service, DoAcceptNodeManager(); RAY_CHECK_OK(RegisterGcs(node_ip_address, socket_name_, - object_manager_config.store_socket_name, - redis_address, redis_port, main_service, - node_manager_config)); + object_manager_config.store_socket_name, redis_address, + redis_port, main_service, node_manager_config)); RAY_CHECK_OK(RegisterPeriodicTimer(main_service)); } -Raylet::~Raylet() { - RAY_CHECK_OK(gcs_client_->client_table().Disconnect()); -} +Raylet::~Raylet() { RAY_CHECK_OK(gcs_client_->client_table().Disconnect()); } ray::Status Raylet::RegisterPeriodicTimer(boost::asio::io_service &io_service) { boost::posix_time::milliseconds timer_period_ms(100); diff --git a/src/ray/test/run_object_manager_tests.sh b/src/ray/test/run_object_manager_tests.sh index f55b51823..af9a979ed 100644 --- a/src/ray/test/run_object_manager_tests.sh +++ b/src/ray/test/run_object_manager_tests.sh @@ -30,12 +30,11 @@ echo "$REDIS_DIR/redis-server --loglevel warning --loadmodule $REDIS_MODULE --po echo "$REDIS_DIR/redis-cli -p 6379 shutdown" # Allow cleanup commands to fail. -killall plasma_store || true -$REDIS_DIR/redis-cli -p 6379 shutdown || true +# killall plasma_store || true +# $REDIS_DIR/redis-cli -p 6379 shutdown || true sleep 1s $REDIS_DIR/redis-server --loglevel warning --loadmodule $REDIS_MODULE --port 6379 & sleep 1s - # Run tests. $CORE_DIR/src/ray/object_manager/object_manager_stress_test $STORE_EXEC sleep 1s