mirror of
https://github.com/wassname/ray.git
synced 2026-06-30 08:10:28 +08:00
Wapaul1
345 lines
16 KiB
C++
345 lines
16 KiB
C++
#include "objstore.h"
|
|
|
|
#include <chrono>
|
|
#include "utils.h"
|
|
|
|
const size_t ObjStoreService::CHUNK_SIZE = 8 * 1024;
|
|
|
|
// this method needs to be protected by a objstore_lock_
|
|
// TODO(rkn): Make sure that we do not in fact need the objstore_lock_. We want multiple deliveries to be able to happen simultaneously.
|
|
void ObjStoreService::pull_data_from(ObjRef objref, ObjStore::Stub& stub) {
|
|
RAY_LOG(RAY_DEBUG, "Objstore " << objstoreid_ << " is beginning to pull objref " << objref);
|
|
ObjChunk chunk;
|
|
ClientContext context;
|
|
StreamObjToRequest stream_request;
|
|
stream_request.set_objref(objref);
|
|
std::unique_ptr<ClientReader<ObjChunk> > reader(stub.StreamObjTo(&context, stream_request));
|
|
|
|
size_t total_size = 0;
|
|
ObjHandle handle;
|
|
if (reader->Read(&chunk)) {
|
|
total_size = chunk.total_size();
|
|
handle = alloc(objref, total_size);
|
|
}
|
|
size_t num_bytes = 0;
|
|
segmentpool_lock_.lock();
|
|
uint8_t* data = segmentpool_->get_address(handle);
|
|
segmentpool_lock_.unlock();
|
|
do {
|
|
RAY_CHECK_LE(num_bytes + chunk.data().size(), total_size, "The reader attempted to stream too many bytes.");
|
|
std::memcpy(data, chunk.data().c_str(), chunk.data().size());
|
|
data += chunk.data().size();
|
|
num_bytes += chunk.data().size();
|
|
} while (reader->Read(&chunk));
|
|
Status status = reader->Finish(); // Right now we don't use the status.
|
|
|
|
// finalize object
|
|
RAY_CHECK_EQ(num_bytes, total_size, "Streamed objref " << objref << ", but num_bytes != total_size");
|
|
object_ready(objref, chunk.metadata_offset());
|
|
RAY_LOG(RAY_DEBUG, "finished streaming data, objref was " << objref << " and size was " << num_bytes);
|
|
}
|
|
|
|
ObjStoreService::ObjStoreService(const std::string& objstore_address, std::shared_ptr<Channel> scheduler_channel)
|
|
: scheduler_stub_(Scheduler::NewStub(scheduler_channel)), objstore_address_(objstore_address) {
|
|
recv_queue_.connect(std::string("queue:") + objstore_address + std::string(":obj"), true);
|
|
ClientContext context;
|
|
RegisterObjStoreRequest request;
|
|
request.set_objstore_address(objstore_address);
|
|
RegisterObjStoreReply reply;
|
|
scheduler_stub_->RegisterObjStore(&context, request, &reply);
|
|
objstoreid_ = reply.objstoreid();
|
|
segmentpool_ = std::make_shared<MemorySegmentPool>(objstoreid_, true);
|
|
}
|
|
|
|
// this method needs to be protected by a objstores_lock_
|
|
ObjStore::Stub& ObjStoreService::get_objstore_stub(const std::string& objstore_address) {
|
|
auto iter = objstores_.find(objstore_address);
|
|
if (iter != objstores_.end())
|
|
return *(iter->second);
|
|
auto channel = grpc::CreateChannel(objstore_address, grpc::InsecureChannelCredentials());
|
|
objstores_.emplace(objstore_address, ObjStore::NewStub(channel));
|
|
return *objstores_[objstore_address];
|
|
}
|
|
|
|
Status ObjStoreService::StartDelivery(ServerContext* context, const StartDeliveryRequest* request, AckReply* reply) {
|
|
// TODO(rkn): We're pushing the delivery task onto a new thread so that this method can return immediately. This matters
|
|
// because the scheduler holds a lock while DeliverObj is being called. The correct solution is to make DeliverObj
|
|
// an asynchronous call (and similarly with the rest of the object store service methods).
|
|
std::string address = request->objstore_address();
|
|
ObjRef objref = request->objref();
|
|
{
|
|
std::lock_guard<std::mutex> memory_lock(memory_lock_);
|
|
if (objref >= memory_.size()) {
|
|
memory_.resize(objref + 1, std::make_pair(ObjHandle(), MemoryStatusType::NOT_PRESENT));
|
|
}
|
|
if (memory_[objref].second == MemoryStatusType::NOT_PRESENT) {
|
|
}
|
|
else {
|
|
RAY_CHECK_NEQ(memory_[objref].second, MemoryStatusType::DEALLOCATED, "Objstore " << objstoreid_ << " is attempting to get objref " << objref << ", but memory_[objref] == DEALLOCATED.");
|
|
RAY_LOG(RAY_DEBUG, "Objstore " << objstoreid_ << " already has objref " << objref << " or it is already being shipped, so no need to pull it again.");
|
|
return Status::OK;
|
|
}
|
|
memory_[objref].second = MemoryStatusType::PRE_ALLOCED;
|
|
}
|
|
delivery_threads_.push_back(std::make_shared<std::thread>([this, address, objref]() {
|
|
std::lock_guard<std::mutex> objstores_lock(objstores_lock_);
|
|
ObjStore::Stub& stub = get_objstore_stub(address);
|
|
pull_data_from(objref, stub);
|
|
}));
|
|
return Status::OK;
|
|
}
|
|
|
|
Status ObjStoreService::ObjStoreInfo(ServerContext* context, const ObjStoreInfoRequest* request, ObjStoreInfoReply* reply) {
|
|
std::lock_guard<std::mutex> memory_lock(memory_lock_);
|
|
for (size_t i = 0; i < memory_.size(); ++i) {
|
|
if (memory_[i].second == MemoryStatusType::READY) { // is the object available?
|
|
reply->add_objref(i);
|
|
}
|
|
}
|
|
/*
|
|
for (int i = 0; i < request->objref_size(); ++i) {
|
|
ObjRef objref = request->objref(i);
|
|
Obj* obj = new Obj();
|
|
std::string data(memory_[objref].ptr.data, memory_[objref].ptr.len); // copies, but for debugging should be ok
|
|
obj->ParseFromString(data);
|
|
reply->mutable_obj()->AddAllocated(obj);
|
|
}
|
|
*/
|
|
return Status::OK;
|
|
}
|
|
|
|
Status ObjStoreService::StreamObjTo(ServerContext* context, const StreamObjToRequest* request, ServerWriter<ObjChunk>* writer) {
|
|
RAY_LOG(RAY_DEBUG, "begin to stream data from object store " << objstoreid_);
|
|
ObjChunk chunk;
|
|
ObjRef objref = request->objref();
|
|
memory_lock_.lock();
|
|
RAY_CHECK_LT(objref, memory_.size(), "Objstore " << objstoreid_ << " is attempting to use objref " << objref << " in StreamObjTo, but this objref is not present in the object store.");
|
|
RAY_CHECK_EQ(memory_[objref].second, MemoryStatusType::READY, "Objstore " << objstoreid_ << " is attempting to stream objref " << objref << ", but memory_[objref].second != MemoryStatusType::READY.");
|
|
ObjHandle handle = memory_[objref].first;
|
|
memory_lock_.unlock(); // TODO(rkn): Make sure we don't still need to hold on to this lock.
|
|
segmentpool_lock_.lock();
|
|
const uint8_t* head = segmentpool_->get_address(handle);
|
|
segmentpool_lock_.unlock();
|
|
size_t size = handle.size();
|
|
for (size_t i = 0; i < size; i += CHUNK_SIZE) {
|
|
chunk.set_metadata_offset(handle.metadata_offset());
|
|
chunk.set_total_size(size);
|
|
chunk.set_data(head + i, std::min(CHUNK_SIZE, size - i));
|
|
RAY_CHECK(writer->Write(chunk), "stream connection prematurely closed")
|
|
}
|
|
return Status::OK;
|
|
}
|
|
|
|
Status ObjStoreService::NotifyAlias(ServerContext* context, const NotifyAliasRequest* request, AckReply* reply) {
|
|
// NotifyAlias assumes that the objstore already holds canonical_objref
|
|
ObjRef alias_objref = request->alias_objref();
|
|
ObjRef canonical_objref = request->canonical_objref();
|
|
RAY_LOG(RAY_DEBUG, "Aliasing objref " << alias_objref << " with objref " << canonical_objref);
|
|
{
|
|
std::lock_guard<std::mutex> memory_lock(memory_lock_);
|
|
RAY_CHECK_LT(canonical_objref, memory_.size(), "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not in the objstore.")
|
|
RAY_CHECK_NEQ(memory_[canonical_objref].second, MemoryStatusType::NOT_READY, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not ready yet in the objstore.")
|
|
RAY_CHECK_NEQ(memory_[canonical_objref].second, MemoryStatusType::NOT_PRESENT, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not present in the objstore.")
|
|
RAY_CHECK_NEQ(memory_[canonical_objref].second, MemoryStatusType::DEALLOCATED, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " has already been deallocated.")
|
|
if (alias_objref >= memory_.size()) {
|
|
memory_.resize(alias_objref + 1, std::make_pair(ObjHandle(), MemoryStatusType::NOT_PRESENT));
|
|
}
|
|
memory_[alias_objref].first = memory_[canonical_objref].first;
|
|
memory_[alias_objref].second = MemoryStatusType::READY;
|
|
}
|
|
ObjRequest done_request;
|
|
done_request.type = ObjRequestType::ALIAS_DONE;
|
|
done_request.objref = alias_objref;
|
|
recv_queue_.send(&done_request);
|
|
return Status::OK;
|
|
}
|
|
|
|
Status ObjStoreService::DeallocateObject(ServerContext* context, const DeallocateObjectRequest* request, AckReply* reply) {
|
|
ObjRef canonical_objref = request->canonical_objref();
|
|
RAY_LOG(RAY_REFCOUNT, "Deallocating canonical_objref " << canonical_objref);
|
|
std::lock_guard<std::mutex> memory_lock(memory_lock_);
|
|
RAY_CHECK_EQ(memory_[canonical_objref].second, MemoryStatusType::READY, "Attempting to deallocate canonical_objref " << canonical_objref << ", but memory_[canonical_objref].second = " << memory_[canonical_objref].second);
|
|
RAY_CHECK_LT(canonical_objref, memory_.size(), "Attempting to deallocate canonical_objref " << canonical_objref << ", but it is not in the objstore.");
|
|
segmentpool_lock_.lock();
|
|
segmentpool_->deallocate(memory_[canonical_objref].first);
|
|
segmentpool_lock_.unlock();
|
|
memory_[canonical_objref].second = MemoryStatusType::DEALLOCATED;
|
|
return Status::OK;
|
|
}
|
|
|
|
// This table describes how the memory status changes in response to requests.
|
|
//
|
|
// MemoryStatus | ObjRequest | New MemoryStatus | action performed
|
|
// -------------+-------------+------------------+----------------------------
|
|
// NOT_PRESENT | ALLOC | NOT_READY | allocate object
|
|
// NOT_READY | WORKER_DONE | READY | send ObjReady to scheduler
|
|
// NOT_READY | GET | NOT_READY | add to pull queue
|
|
// READY | GET | READY | return handle
|
|
// READY | DEALLOC | DEALLOCATED | deallocate
|
|
// -------------+-------------+------------------+----------------------------
|
|
void ObjStoreService::process_objstore_request(const ObjRequest request) {
|
|
switch (request.type) {
|
|
case ObjRequestType::ALIAS_DONE: {
|
|
process_pulls_for_objref(request.objref);
|
|
}
|
|
break;
|
|
default: {
|
|
RAY_CHECK(false, "Attempting to process request of type " << request.type << ". This code should be unreachable.");
|
|
}
|
|
}
|
|
}
|
|
|
|
void ObjStoreService::process_worker_request(const ObjRequest request) {
|
|
if (request.workerid >= send_queues_.size()) {
|
|
send_queues_.resize(request.workerid + 1);
|
|
}
|
|
if (!send_queues_[request.workerid].connected()) {
|
|
std::string queue_name = std::string("queue:") + objstore_address_ + std::string(":worker:") + std::to_string(request.workerid) + std::string(":obj");
|
|
send_queues_[request.workerid].connect(queue_name, false);
|
|
}
|
|
{
|
|
std::lock_guard<std::mutex> memory_lock(memory_lock_);
|
|
if (request.objref >= memory_.size()) {
|
|
memory_.resize(request.objref + 1, std::make_pair(ObjHandle(), MemoryStatusType::NOT_PRESENT));
|
|
}
|
|
}
|
|
switch (request.type) {
|
|
case ObjRequestType::ALLOC: {
|
|
ObjHandle handle = alloc(request.objref, request.size); // This method acquires memory_lock_
|
|
send_queues_[request.workerid].send(&handle);
|
|
}
|
|
break;
|
|
case ObjRequestType::GET: {
|
|
std::lock_guard<std::mutex> memory_lock(memory_lock_);
|
|
std::pair<ObjHandle, MemoryStatusType>& item = memory_[request.objref];
|
|
if (item.second == MemoryStatusType::READY) {
|
|
RAY_LOG(RAY_DEBUG, "Responding to GET request: returning objref " << request.objref);
|
|
send_queues_[request.workerid].send(&item.first);
|
|
} else if (item.second == MemoryStatusType::NOT_READY || item.second == MemoryStatusType::NOT_PRESENT || item.second == MemoryStatusType::PRE_ALLOCED) {
|
|
std::lock_guard<std::mutex> lock(pull_queue_lock_);
|
|
pull_queue_.push_back(std::make_pair(request.workerid, request.objref));
|
|
} else {
|
|
RAY_CHECK(false, "A worker requested objref " << request.objref << ", but memory_[objref].second = " << memory_[request.objref].second);
|
|
}
|
|
}
|
|
break;
|
|
case ObjRequestType::WORKER_DONE: {
|
|
object_ready(request.objref, request.metadata_offset); // This method acquires memory_lock_
|
|
}
|
|
break;
|
|
default: {
|
|
RAY_CHECK(false, "Attempting to process request of type " << request.type << ". This code should be unreachable.");
|
|
}
|
|
}
|
|
}
|
|
|
|
void ObjStoreService::process_requests() {
|
|
// TODO(rkn): Should memory_lock_ be used in this method?
|
|
ObjRequest request;
|
|
while (true) {
|
|
recv_queue_.receive(&request);
|
|
switch (request.type) {
|
|
case ObjRequestType::ALLOC: {
|
|
RAY_LOG(RAY_VERBOSE, "Request (worker " << request.workerid << " to objstore " << objstoreid_ << "): Allocate object with objref " << request.objref << " and size " << request.size);
|
|
process_worker_request(request);
|
|
}
|
|
break;
|
|
case ObjRequestType::GET: {
|
|
RAY_LOG(RAY_VERBOSE, "Request (worker " << request.workerid << " to objstore " << objstoreid_ << "): Get object with objref " << request.objref);
|
|
process_worker_request(request);
|
|
}
|
|
break;
|
|
case ObjRequestType::WORKER_DONE: {
|
|
RAY_LOG(RAY_VERBOSE, "Request (worker " << request.workerid << " to objstore " << objstoreid_ << "): Finalize object with objref " << request.objref);
|
|
process_worker_request(request);
|
|
}
|
|
break;
|
|
case ObjRequestType::ALIAS_DONE: {
|
|
process_objstore_request(request);
|
|
}
|
|
break;
|
|
default: {
|
|
RAY_CHECK(false, "Attempting to process request of type " << request.type << ". This code should be unreachable.");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
void ObjStoreService::process_pulls_for_objref(ObjRef objref) {
|
|
std::pair<ObjHandle, MemoryStatusType>& item = memory_[objref];
|
|
std::lock_guard<std::mutex> pull_queue_lock(pull_queue_lock_);
|
|
for (size_t i = 0; i < pull_queue_.size(); ++i) {
|
|
if (pull_queue_[i].second == objref) {
|
|
ObjHandle& elem = memory_[objref].first;
|
|
send_queues_[pull_queue_[i].first].send(&item.first);
|
|
// Remove the pull task from the queue
|
|
std::swap(pull_queue_[i], pull_queue_[pull_queue_.size() - 1]);
|
|
pull_queue_.pop_back();
|
|
i -= 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
ObjHandle ObjStoreService::alloc(ObjRef objref, size_t size) {
|
|
segmentpool_lock_.lock();
|
|
ObjHandle handle = segmentpool_->allocate(size);
|
|
segmentpool_lock_.unlock();
|
|
std::lock_guard<std::mutex> memory_lock(memory_lock_);
|
|
RAY_LOG(RAY_VERBOSE, "Allocating space for objref " << objref << " on object store " << objstoreid_);
|
|
RAY_CHECK(memory_[objref].second == MemoryStatusType::NOT_PRESENT || memory_[objref].second == MemoryStatusType::PRE_ALLOCED, "Attempting to allocate space for objref " << objref << ", but memory_[objref].second = " << memory_[objref].second);
|
|
memory_[objref].first = handle;
|
|
memory_[objref].second = MemoryStatusType::NOT_READY;
|
|
return handle;
|
|
}
|
|
|
|
void ObjStoreService::object_ready(ObjRef objref, size_t metadata_offset) {
|
|
{
|
|
std::lock_guard<std::mutex> memory_lock(memory_lock_);
|
|
std::pair<ObjHandle, MemoryStatusType>& item = memory_[objref];
|
|
RAY_CHECK_EQ(item.second, MemoryStatusType::NOT_READY, "A worker notified the object store that objref " << objref << " has been written to the object store, but memory_[objref].second != NOT_READY.");
|
|
item.first.set_metadata_offset(metadata_offset);
|
|
item.second = MemoryStatusType::READY;
|
|
}
|
|
process_pulls_for_objref(objref);
|
|
// Tell the scheduler that the object arrived
|
|
// TODO(pcm): put this in a separate thread so we don't have to pay the latency here
|
|
ClientContext objready_context;
|
|
ObjReadyRequest objready_request;
|
|
objready_request.set_objref(objref);
|
|
objready_request.set_objstoreid(objstoreid_);
|
|
AckReply objready_reply;
|
|
scheduler_stub_->ObjReady(&objready_context, objready_request, &objready_reply);
|
|
}
|
|
|
|
void ObjStoreService::start_objstore_service() {
|
|
communicator_thread_ = std::thread([this]() {
|
|
RAY_LOG(RAY_INFO, "started object store communicator server");
|
|
process_requests();
|
|
});
|
|
}
|
|
|
|
void start_objstore(const char* scheduler_addr, const char* objstore_addr) {
|
|
auto scheduler_channel = grpc::CreateChannel(scheduler_addr, grpc::InsecureChannelCredentials());
|
|
RAY_LOG(RAY_INFO, "object store " << objstore_addr << " connected to scheduler " << scheduler_addr);
|
|
std::string objstore_address(objstore_addr);
|
|
ObjStoreService service(objstore_address, scheduler_channel);
|
|
service.start_objstore_service();
|
|
std::string::iterator split_point = split_ip_address(objstore_address);
|
|
std::string port;
|
|
port.assign(split_point, objstore_address.end());
|
|
ServerBuilder builder;
|
|
builder.AddListeningPort(std::string("0.0.0.0:") + port, grpc::InsecureServerCredentials());
|
|
builder.RegisterService(&service);
|
|
std::unique_ptr<Server> server(builder.BuildAndStart());
|
|
|
|
server->Wait();
|
|
}
|
|
|
|
int main(int argc, char** argv) {
|
|
RAY_CHECK_EQ(argc, 3, "object store: expected two arguments (scheduler ip address and object store ip address)");
|
|
|
|
start_objstore(argv[1], argv[2]);
|
|
|
|
return 0;
|
|
}
|