Merge pull request #47 from amplab/internode

Internode

lib/orchpy/orchpy/worker.py

@@ -14,11 +14,11 @@ class Worker(object):
 
   def put_object(self, objref, value):
     """Put `value` in the local object store with objref `objref`. This assumes that the value for `objref` has not yet been placed in the local object store."""
-    if type(value) == np.ndarray:
-      orchpy.lib.put_arrow(self.handle, objref, value)
-    else:
-      object_capsule, contained_objrefs = serialization.serialize(self.handle, value) # contained_objrefs is a list of the objrefs contained in object_capsule
-      orchpy.lib.put_object(self.handle, objref, object_capsule, contained_objrefs)
+    # if type(value) == np.ndarray:
+    #   orchpy.lib.put_arrow(self.handle, objref, value)
+    # else:
+    object_capsule, contained_objrefs = serialization.serialize(self.handle, value) # contained_objrefs is a list of the objrefs contained in object_capsule
+    orchpy.lib.put_object(self.handle, objref, object_capsule, contained_objrefs)
 
   def get_object(self, objref):
     """
@@ -27,11 +27,11 @@ class Worker(object):
 
     WARNING: get_object can only be called on a canonical objref.
     """
-    if orchpy.lib.is_arrow(self.handle, objref):
-      return orchpy.lib.get_arrow(self.handle, objref)
-    else:
-      object_capsule = orchpy.lib.get_object(self.handle, objref)
-      return serialization.deserialize(self.handle, object_capsule)
+    # if orchpy.lib.is_arrow(self.handle, objref):
+    #   return orchpy.lib.get_arrow(self.handle, objref)
+    # else:
+    object_capsule = orchpy.lib.get_object(self.handle, objref)
+    return serialization.deserialize(self.handle, object_capsule)
 
   def alias_objrefs(self, alias_objref, target_objref):
     """Make `alias_objref` refer to the same object that `target_objref` refers to."""

protos/orchestra.proto

@@ -149,10 +149,10 @@ message SchedulerInfoReply {
 // Object stores
 
 service ObjStore {
-  // Request to deliver an object to another object store (called by the scheduler)
-  rpc DeliverObj(DeliverObjRequest) returns (AckReply);
+  // Tell the object store to begin pulling an object from another object store (called by the scheduler)
+  rpc StartDelivery(StartDeliveryRequest) returns (AckReply);
   // Accept incoming data from another object store, as a stream of object chunks
-  rpc StreamObj(stream ObjChunk) returns (AckReply);
+  rpc StreamObjTo(StreamObjToRequest) returns (stream ObjChunk);
   // Notify the object store about objref aliasing. This is called by the scheduler
   rpc NotifyAlias(NotifyAliasRequest) returns (AckReply);
   // Tell the object store to deallocate an object held by the object store. This is called by the scheduler.
@@ -161,8 +161,8 @@ service ObjStore {
   rpc ObjStoreInfo(ObjStoreInfoRequest) returns (ObjStoreInfoReply);
 }
 
-message DeliverObjRequest {
-  string objstore_address = 1; // Object store to deliver the object to
+message StartDeliveryRequest {
+  string objstore_address = 1; // Object store to pull the object from
   uint64 objref = 2; // Reference of object that gets delivered
 }
 
@@ -174,9 +174,13 @@ message RegisterObjReply {
   uint64 handle = 1; // Handle to memory segment where object is stored
 }
 
-message ObjChunk {
+message StreamObjToRequest {
   uint64 objref = 1; // Object reference of the object being streamed
-  uint64 totalsize = 2; // Total size of the object
+}
+
+message ObjChunk {
+  uint64 total_size = 1; // Total size of the object
+  uint64 metadata_offset = 2; // Offset of the arrow metadata
   bytes data = 3; // Data for this chunk of the object
 }
 
@@ -193,12 +197,6 @@ message GetObjRequest {
   uint64 objref = 1; // Object reference of the object being requested by the worker
 }
 
-message GetObjReply {
-  string bucket = 1; // Name of the shared memory segment where the object is stored
-  uint64 handle = 2; // Shared memory pointer to the object
-  uint64 size = 3; // Total size of the object in bytes
-}
-
 // These messages are for getting information about the object store state
 
 message ObjStoreInfoRequest {

src/ipc.cc

@@ -31,9 +31,9 @@ MemorySegmentPool::MemorySegmentPool(ObjStoreId objstoreid, bool create) : objst
 // creates a memory segment if it is not already there; if the pool is in create mode,
 // space is allocated, if it is in open mode, the shared memory is mapped into the process
 void MemorySegmentPool::open_segment(SegmentId segmentid, size_t size) {
-  ORCH_LOG(ORCH_DEBUG, "opening segmentid " << segmentid);
+  ORCH_LOG(ORCH_DEBUG, "Opening segmentid " << segmentid << " on object store " << objstoreid_ << " with create_mode_ = " << create_mode_);
   if (segmentid != segments_.size() && create_mode_) {
-    ORCH_LOG(ORCH_FATAL, "Attempting to open segmentid " << segmentid << " on the object store, but segments_.size() = " << segments_.size());
+    ORCH_LOG(ORCH_FATAL, "Object store " << objstoreid_ << " is attempting to open segmentid " << segmentid << " on the object store, but segments_.size() = " << segments_.size());
   }
   if (segmentid >= segments_.size()) { // resize and initialize segments_
     int current_size = segments_.size();
@@ -90,7 +90,12 @@ void MemorySegmentPool::deallocate(ObjHandle pointer) {
 // returns address of the object refered to by the handle, needs to be called on
 // the process that will use the address
 uint8_t* MemorySegmentPool::get_address(ObjHandle pointer) {
-  open_segment(pointer.segmentid());
+  if (create_mode_ && segments_[pointer.segmentid()].second != SegmentStatusType::OPENED) {
+    ORCH_LOG(ORCH_FATAL, "Object store " << objstoreid_ << " is attempting to call get_address on segmentid " << pointer.segmentid() << ", which has not been opened yet.");
+  }
+  if (!create_mode_) {
+    open_segment(pointer.segmentid());
+  }
   managed_shared_memory* segment = segments_[pointer.segmentid()].first.get();
   return static_cast<uint8_t*>(segment->get_address_from_handle(pointer.ipcpointer()));
 }

src/objstore.cc

@@ -1,25 +1,44 @@
 #include "objstore.h"
 #include <chrono>
 
-const size_t ObjStoreClient::CHUNK_SIZE = 8 * 1024;
+const size_t ObjStoreService::CHUNK_SIZE = 8 * 1024;
 
 // this method needs to be protected by a objstore_lock_
-Status ObjStoreClient::upload_data_to(slice data, ObjRef objref, ObjStore::Stub& stub) {
+// 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) {
+  ORCH_LOG(ORCH_DEBUG, "Objstore " << objstoreid_ << " is beginning to pull objref " << objref);
   ObjChunk chunk;
   ClientContext context;
-  AckReply reply;
-  std::unique_ptr<ClientWriter<ObjChunk> > writer(stub.StreamObj(&context, &reply));
-  const uint8_t* head = data.data;
-  for (size_t i = 0; i < data.len; i += CHUNK_SIZE) {
-    chunk.set_objref(objref);
-    chunk.set_totalsize(data.len);
-    chunk.set_data(head + i, std::min(CHUNK_SIZE, data.len - i));
-    if (!writer->Write(chunk)) {
-      ORCH_LOG(ORCH_FATAL, "stream connection prematurely closed")
-    }
+  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);
   }
-  writer->WritesDone();
-  return writer->Finish();
+  size_t num_bytes = 0;
+  segmentpool_lock_.lock();
+  uint8_t* data = segmentpool_->get_address(handle);
+  segmentpool_lock_.unlock();
+  do {
+    if (num_bytes + chunk.data().size() > total_size) {
+      ORCH_LOG(ORCH_FATAL, "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
+  if (num_bytes != total_size) {
+    ORCH_LOG(ORCH_FATAL, "Streamed objref " << objref << ", but num_bytes != total_size");
+  }
+  object_ready(objref, chunk.metadata_offset());
+  ORCH_LOG(ORCH_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)
@@ -44,15 +63,35 @@ ObjStore::Stub& ObjStoreService::get_objstore_stub(const std::string& objstore_a
   return *objstores_[objstore_address];
 }
 
-/*
-Status ObjStoreService::DeliverObj(ServerContext* context, const DeliverObjRequest* request, AckReply* reply) {
-  std::lock_guard<std::mutex> objstores_lock(objstores_lock_);
-  ObjStore::Stub& stub = get_objstore_stub(request->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();
-  Status status = ObjStoreClient::upload_data_to(memory_[objref].ptr, objref, stub);
-  return status;
+  {
+    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 if (memory_[objref].second == MemoryStatusType::DEALLOCATED) {
+      ORCH_LOG(ORCH_FATAL, "Objstore " << objstoreid_ << " is attempting to get objref " << objref << ", but memory_[objref] == DEALLOCATED.");
+    }
+    else {
+      ORCH_LOG(ORCH_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_);
@@ -73,70 +112,59 @@ Status ObjStoreService::ObjStoreInfo(ServerContext* context, const ObjStoreInfoR
   return Status::OK;
 }
 
-/*
-Status ObjStoreService::StreamObj(ServerContext* context, ServerReader<ObjChunk>* reader, AckReply* reply) {
-  ORCH_LOG(ORCH_VERBOSE, "begin to stream data to object store " << objstoreid_);
-  memory_lock_.lock();
+Status ObjStoreService::StreamObjTo(ServerContext* context, const StreamObjToRequest* request, ServerWriter<ObjChunk>* writer) {
+  ORCH_LOG(ORCH_DEBUG, "begin to stream data from object store " << objstoreid_);
   ObjChunk chunk;
-  ObjRef objref = 0;
-  size_t totalsize = 0;
-  if (reader->Read(&chunk)) {
-    objref = chunk.objref();
-    totalsize = chunk.totalsize();
-    allocate_memory(objref, totalsize);
+  ObjRef objref = request->objref();
+  memory_lock_.lock();
+  if (objref >= memory_.size()) {
+    ORCH_LOG(ORCH_FATAL, "Objstore " << objstoreid_ << " is attempting to use objref " << objref << " in StreamObjTo, but this objref is not present in the object store.");
   }
-  size_t num_bytes = 0;
-  char* data = memory_[objref].ptr.data;
-
-  do {
-    if (num_bytes + chunk.data().size() > totalsize) {
-      memory_lock_.unlock();
-      return Status::CANCELLED;
+  if (memory_[objref].second != MemoryStatusType::READY) {
+    ORCH_LOG(ORCH_FATAL, "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));
+    if (!writer->Write(chunk)) {
+      ORCH_LOG(ORCH_FATAL, "stream connection prematurely closed")
     }
-    std::memcpy(data, chunk.data().c_str(), chunk.data().size());
-    data += chunk.data().size();
-    num_bytes += chunk.data().size();
-  } while (reader->Read(&chunk));
-
-  ORCH_LOG(ORCH_VERBOSE, "finished streaming data, objref was " << objref << " and size was " << num_bytes);
-
-  memory_lock_.unlock();
-
-  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);
-
+  }
   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();
   ORCH_LOG(ORCH_DEBUG, "Aliasing objref " << alias_objref << " with objref " << canonical_objref);
-  std::lock_guard<std::mutex> memory_lock(memory_lock_);
-  if (canonical_objref >= memory_.size()) {
-    ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not in the objstore.")
+  {
+    std::lock_guard<std::mutex> memory_lock(memory_lock_);
+    if (canonical_objref >= memory_.size()) {
+      ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not in the objstore.")
+    }
+    if (memory_[canonical_objref].second == MemoryStatusType::NOT_READY) {
+      ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not ready yet in the objstore.")
+    }
+    if (memory_[canonical_objref].second == MemoryStatusType::NOT_PRESENT) {
+      ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not present in the objstore.")
+    }
+    if (memory_[canonical_objref].second == MemoryStatusType::DEALLOCATED) {
+      ORCH_LOG(ORCH_FATAL, "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;
   }
-  if (memory_[canonical_objref].second == MemoryStatusType::NOT_READY) {
-    ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not ready yet in the objstore.")
-  }
-  if (memory_[canonical_objref].second == MemoryStatusType::NOT_PRESENT) {
-    ORCH_LOG(ORCH_FATAL, "Attempting to alias objref " << alias_objref << " with objref " << canonical_objref << ", but objref " << canonical_objref << " is not present in the objstore.")
-  }
-  if (memory_[canonical_objref].second == MemoryStatusType::DEALLOCATED) {
-    ORCH_LOG(ORCH_FATAL, "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;
@@ -154,7 +182,9 @@ Status ObjStoreService::DeallocateObject(ServerContext* context, const Deallocat
   if (canonical_objref >= memory_.size()) {
     ORCH_LOG(ORCH_FATAL, "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;
 }
@@ -197,15 +227,8 @@ void ObjStoreService::process_worker_request(const ObjRequest request) {
   }
   switch (request.type) {
     case ObjRequestType::ALLOC: {
-        // TODO(rkn): Does segmentpool_ need a lock around it?
-        ObjHandle reply = segmentpool_->allocate(request.size);
-        send_queues_[request.workerid].send(&reply);
-        std::lock_guard<std::mutex> memory_lock(memory_lock_);
-        if (memory_[request.objref].second != MemoryStatusType::NOT_PRESENT) {
-          ORCH_LOG(ORCH_FATAL, "Attempting to allocate space for objref " << request.objref << ", but memory_[objref].second != MemoryStatusType::NOT_PRESENT, it equals " << memory_[request.objref].second);
-        }
-        memory_[request.objref].first = reply;
-        memory_[request.objref].second = MemoryStatusType::NOT_READY;
+        ObjHandle handle = alloc(request.objref, request.size); // This method acquires memory_lock_
+        send_queues_[request.workerid].send(&handle);
       }
       break;
     case ObjRequestType::GET: {
@@ -214,7 +237,7 @@ void ObjStoreService::process_worker_request(const ObjRequest request) {
         if (item.second == MemoryStatusType::READY) {
           ORCH_LOG(ORCH_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) {
+        } 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 {
@@ -223,24 +246,7 @@ void ObjStoreService::process_worker_request(const ObjRequest request) {
       }
       break;
     case ObjRequestType::WORKER_DONE: {
-        {
-          std::lock_guard<std::mutex> memory_lock(memory_lock_);
-          std::pair<ObjHandle, MemoryStatusType>& item = memory_[request.objref];
-          if (item.second != MemoryStatusType::NOT_READY) {
-            ORCH_LOG(ORCH_FATAL, "A worker notified the object store that objref " << request.objref << " has been written to the object store, but memory_[objref].second != NOT_READY.");
-          }
-          item.first.set_metadata_offset(request.metadata_offset);
-          item.second = MemoryStatusType::READY;
-        }
-        process_pulls_for_objref(request.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(request.objref);
-        objready_request.set_objstoreid(objstoreid_);
-        AckReply objready_reply;
-        scheduler_stub_->ObjReady(&objready_context, objready_request, &objready_reply);
+        object_ready(request.objref, request.metadata_offset); // This method acquires memory_lock_
       }
       break;
     default: {
@@ -293,6 +299,40 @@ void ObjStoreService::process_pulls_for_objref(ObjRef objref) {
   }
 }
 
+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_);
+  if (memory_[objref].second != MemoryStatusType::NOT_PRESENT && memory_[objref].second != MemoryStatusType::PRE_ALLOCED) {
+    ORCH_LOG(ORCH_FATAL, "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];
+    if (item.second != MemoryStatusType::NOT_READY) {
+      ORCH_LOG(ORCH_FATAL, "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]() {
     ORCH_LOG(ORCH_INFO, "started object store communicator server");

src/objstore.h

@@ -17,40 +17,50 @@ using grpc::ServerBuilder;
 using grpc::ServerReader;
 using grpc::ServerContext;
 using grpc::ClientContext;
-using grpc::ClientWriter;
+using grpc::ServerWriter;
+using grpc::ClientReader;
 using grpc::Status;
 
 using grpc::Channel;
 
-class ObjStoreClient {
-public:
-  static const size_t CHUNK_SIZE;
-  static Status upload_data_to(slice data, ObjRef objref, ObjStore::Stub& stub);
-};
-
-enum MemoryStatusType {READY = 0, NOT_READY = 1, DEALLOCATED = 2, NOT_PRESENT = 3};
+// READY:       This is used to indicate that the object has been copied from a
+//              worker and is ready to be used.
+// NOT_READY:   This is used to indicate that memory has been allocated for the
+//              object, but the object hasn't been copied from a worker yet.
+// DEALLOCATED: This is used to indicate that the object has been deallocated.
+// NOT_PRESENT: This is used to indicate that space has not been allocated for
+//              this object in this object store.
+// PRE_ALLOCED: This is used to indicate that the memory has not yet been
+//              alloced, but it will be alloced soon. This is set when we call
+//              StartDelivery.
+enum MemoryStatusType {READY = 0, NOT_READY = 1, DEALLOCATED = 2, NOT_PRESENT = 3, PRE_ALLOCED = 4};
 
 class ObjStoreService final : public ObjStore::Service {
 public:
   ObjStoreService(const std::string& objstore_address, std::shared_ptr<Channel> scheduler_channel);
 
-  // Status DeliverObj(ServerContext* context, const DeliverObjRequest* request, AckReply* reply) override;
-  // Status StreamObj(ServerContext* context, ServerReader<ObjChunk>* reader, AckReply* reply) override;
+  Status StartDelivery(ServerContext* context, const StartDeliveryRequest* request, AckReply* reply) override;
+  Status StreamObjTo(ServerContext* context, const StreamObjToRequest* request, ServerWriter<ObjChunk>* writer) override;
   Status NotifyAlias(ServerContext* context, const NotifyAliasRequest* request, AckReply* reply) override;
   Status DeallocateObject(ServerContext* context, const DeallocateObjectRequest* request, AckReply* reply) override;
   Status ObjStoreInfo(ServerContext* context, const ObjStoreInfoRequest* request, ObjStoreInfoReply* reply) override;
   void start_objstore_service();
 private:
+  void pull_data_from(ObjRef objref, ObjStore::Stub& stub);
   // check if we already connected to the other objstore, if yes, return reference to connection, otherwise connect
   ObjStore::Stub& get_objstore_stub(const std::string& objstore_address);
   void process_worker_request(const ObjRequest request);
   void process_objstore_request(const ObjRequest request);
   void process_requests();
   void process_pulls_for_objref(ObjRef objref);
+  ObjHandle alloc(ObjRef objref, size_t size);
+  void object_ready(ObjRef objref, size_t metadata_offset);
 
+  static const size_t CHUNK_SIZE;
   std::string objstore_address_;
   ObjStoreId objstoreid_; // id of this objectstore in the scheduler object store table
   std::shared_ptr<MemorySegmentPool> segmentpool_;
+  std::mutex segmentpool_lock_;
   std::vector<std::pair<ObjHandle, MemoryStatusType> > memory_; // object reference -> (memory address, memory status)
   std::mutex memory_lock_;
   std::unordered_map<std::string, std::unique_ptr<ObjStore::Stub>> objstores_;
@@ -58,9 +68,13 @@ private:
   std::unique_ptr<Scheduler::Stub> scheduler_stub_;
   std::vector<std::pair<WorkerId, ObjRef> > pull_queue_;
   std::mutex pull_queue_lock_;
-  MessageQueue<ObjRequest> recv_queue_;
-  std::vector<MessageQueue<ObjHandle> > send_queues_; // workerid -> queue
+  MessageQueue<ObjRequest> recv_queue_; // This queue is used by workers to send tasks to the object store.
+  std::vector<MessageQueue<ObjHandle> > send_queues_; // This maps workerid -> queue. The object store uses these queues to send replies to the relevant workers.
   std::thread communicator_thread_;
+
+  std::vector<std::shared_ptr<std::thread> > delivery_threads_; // TODO(rkn): document
+  // TODO(rkn): possibly add lock, and properly remove these threads from the delivery_threads_ when the deliveries are done
+
 };
 
 #endif

src/scheduler.cc

@@ -191,8 +191,14 @@ Status SchedulerService::SchedulerInfo(ServerContext* context, const SchedulerIn
   return Status::OK;
 }
 
+// TODO(rkn): This could execute multiple times with the same arguments before
+// the delivery finishes, but we only want it to happen once. Currently, the
+// redundancy is handled by the object store, which will only execute the
+// delivery once. However, we may want to handle it in the scheduler in the
+// future.
+//
+// deliver_object assumes that the aliasing for objref has already been completed. That is, has_canonical_objref(objref) == true
 void SchedulerService::deliver_object(ObjRef objref, ObjStoreId from, ObjStoreId to) {
-  // deliver_object assumes that the aliasing for objref has already been completed. That is, has_canonical_objref(objref) == true
   if (from == to) {
     ORCH_LOG(ORCH_FATAL, "attempting to deliver objref " << objref << " from objstore " << from << " to itself.");
   }
@@ -201,12 +207,12 @@ void SchedulerService::deliver_object(ObjRef objref, ObjStoreId from, ObjStoreId
   }
   ClientContext context;
   AckReply reply;
-  DeliverObjRequest request;
+  StartDeliveryRequest request;
   ObjRef canonical_objref = get_canonical_objref(objref);
   request.set_objref(canonical_objref);
   std::lock_guard<std::mutex> lock(objstores_lock_);
-  request.set_objstore_address(objstores_[to].address);
-  objstores_[from].objstore_stub->DeliverObj(&context, request, &reply);
+  request.set_objstore_address(objstores_[from].address);
+  objstores_[to].objstore_stub->StartDelivery(&context, request, &reply);
 }
 
 void SchedulerService::schedule() {
@@ -403,14 +409,16 @@ void SchedulerService::get_info(const SchedulerInfoRequest& request, SchedulerIn
 
 }
 
+// pick_objstore assumes that objtable_lock_ has been acquired
+// pick_objstore must be called with a canonical_objref
 ObjStoreId SchedulerService::pick_objstore(ObjRef canonical_objref) {
-  // pick_objstore must be called with a canonical_objref
   std::mt19937 rng;
   if (!is_canonical(canonical_objref)) {
     ORCH_LOG(ORCH_FATAL, "Attempting to call pick_objstore with a non-canonical objref, (objref " << canonical_objref << ")");
   }
   std::uniform_int_distribution<int> uni(0, objtable_[canonical_objref].size() - 1);
-  return uni(rng);
+  ObjStoreId objstoreid = objtable_[canonical_objref][uni(rng)];
+  return objstoreid;
 }
 
 bool SchedulerService::is_canonical(ObjRef objref) {
@@ -433,7 +441,7 @@ void SchedulerService::perform_pulls() {
       continue;
     }
     ObjRef canonical_objref = get_canonical_objref(objref);
-    ORCH_LOG(ORCH_DEBUG, "attempting to pull objref " << pull.second << " with canonical objref " << canonical_objref);
+    ORCH_LOG(ORCH_DEBUG, "attempting to pull objref " << pull.second << " with canonical objref " << canonical_objref << " to objstore " << get_store(workerid));
 
     objtable_lock_.lock();
     int num_stores = objtable_[canonical_objref].size();
@@ -570,16 +578,17 @@ void SchedulerService::deallocate_object(ObjRef canonical_objref) {
   // so the lock must before outside of these methods (it is acquired in
   // DecrementRefCount).
   ORCH_LOG(ORCH_REFCOUNT, "Deallocating canonical_objref " << canonical_objref << ".");
-  ClientContext context;
-  AckReply reply;
-  DeallocateObjectRequest request;
-  request.set_canonical_objref(canonical_objref);
   {
     std::lock_guard<std::mutex> objtable_lock(objtable_lock_);
     auto &objstores = objtable_[canonical_objref];
     std::lock_guard<std::mutex> objstores_lock(objstores_lock_); // TODO(rkn): Should this be inside the for loop instead?
     for (int i = 0; i < objstores.size(); ++i) {
+      ClientContext context;
+      AckReply reply;
+      DeallocateObjectRequest request;
+      request.set_canonical_objref(canonical_objref);
       ObjStoreId objstoreid = objstores[i];
+      ORCH_LOG(ORCH_REFCOUNT, "Attempting to deallocate canonical_objref " << canonical_objref << " from objstore " << objstoreid);
       objstores_[objstoreid].objstore_stub->DeallocateObject(&context, request, &reply);
     }
   }

src/serialize.cc

@@ -156,7 +156,7 @@ std::shared_ptr<arrow::Array> read_flat_array(BufferMemorySource* source, int64_
   std::shared_ptr<ipc::RowBatchReader> reader;
   Status s = ipc::RowBatchReader::Open(source, metadata_offset, &reader);
   if (!s.ok()) {
-    ORCH_LOG(ORCH_FATAL, s.ToString());
+    ORCH_LOG(ORCH_FATAL, "Error in read_flat_array: " << s.ToString());
   }
   auto field = std::make_shared<arrow::Field>("data", npy_traits<NpyType>::primitive_type);
   std::shared_ptr<arrow::Schema> schema(new arrow::Schema({field}));

test/arrays_test.py

@@ -81,7 +81,7 @@ class ArraysDistTest(unittest.TestCase):
   def testMethods(self):
     test_dir = os.path.dirname(os.path.abspath(__file__))
     test_path = os.path.join(test_dir, "testrecv.py")
-    services.start_cluster(return_drivers=False, num_workers_per_objstore=8, worker_path=test_path)
+    services.start_cluster(return_drivers=False, num_objstores=2, num_workers_per_objstore=8, worker_path=test_path)
 
     x = dist.zeros([9, 25, 51], "float")
     y = dist.assemble(x)

test/runtest.py

@@ -78,11 +78,40 @@ class ObjStoreTest(unittest.TestCase):
       self.assertEqual(result, data)
 
     # pushing an object, shipping it to another worker, and pulling it shouldn't change it
-    # for data in ["h", "h" * 10000, 0, 0.0]:
-    #   objref = worker.push(data, worker1)
-    #   response = objstore1_stub.DeliverObj(orchestra_pb2.DeliverObjRequest(objref=objref.val, objstore_address=address(IP_ADDRESS, objstore2_port)), TIMEOUT_SECONDS)
-    #   result = worker.pull(objref, worker2)
-    #   self.assertEqual(result, data)
+    for data in ["h", "h" * 10000, 0, 0.0, [1, 2, 3, "a", (1, 2)], ("a", ("b", 3))]:
+      objref = worker.push(data, w1)
+      result = worker.pull(objref, w2)
+      self.assertEqual(result, data)
+
+    # pushing an array, shipping it to another worker, and pulling it shouldn't change it
+    for data in [np.zeros([10, 20]), np.random.normal(size=[45, 25])]:
+      objref = worker.push(data, w1)
+      result = worker.pull(objref, w2)
+      self.assertTrue(np.alltrue(result == data))
+
+    """
+    # pulling multiple times shouldn't matter
+    for data in [np.zeros([10, 20]), np.random.normal(size=[45, 25]), np.zeros([10, 20], dtype=np.dtype("float64")), np.zeros([10, 20], dtype=np.dtype("float32")), np.zeros([10, 20], dtype=np.dtype("int64")), np.zeros([10, 20], dtype=np.dtype("int32"))]:
+      objref = worker.push(data, w1)
+      result = worker.pull(objref, w2)
+      result = worker.pull(objref, w2)
+      result = worker.pull(objref, w2)
+      self.assertTrue(np.alltrue(result == data))
+    """
+
+    # shipping a numpy array inside something else should be fine
+    data = ("a", np.random.normal(size=[10, 10]))
+    objref = worker.push(data, w1)
+    result = worker.pull(objref, w2)
+    self.assertTrue(data[0] == result[0])
+    self.assertTrue(np.alltrue(data[1] == result[1]))
+
+    # shipping a numpy array inside something else should be fine
+    data = ["a", np.random.normal(size=[10, 10])]
+    objref = worker.push(data, w1)
+    result = worker.pull(objref, w2)
+    self.assertTrue(data[0] == result[0])
+    self.assertTrue(np.alltrue(data[1] == result[1]))
 
     services.cleanup()