Added handles for python to access buffers in objstore (#282)

src/raylib.cc

@@ -495,6 +495,50 @@ static PyObject* put_arrow(PyObject* self, PyObject* args) {
   Py_RETURN_NONE;
 }
 
+static PyObject* allocate_buffer(PyObject* self, PyObject* args) {
+  Worker* worker;
+  ObjRef objref;
+  SegmentId segmentid;
+  long size;
+  if (!PyArg_ParseTuple(args, "O&O&l", &PyObjectToWorker, &worker, &PyObjectToObjRef, &objref, &size)) {
+    return NULL;
+  }
+  void* address = reinterpret_cast<void*>(const_cast<char*>(worker->allocate_buffer(objref, size, segmentid)));
+  std::vector<npy_intp> dim({size});
+  PyObject* t = PyTuple_New(2);
+  PyTuple_SetItem(t, 0, PyArray_SimpleNewFromData(1, dim.data(), NPY_BYTE, address));
+  PyTuple_SetItem(t, 1, PyInt_FromLong(segmentid));
+  return t;
+}
+
+static PyObject* finish_buffer(PyObject* self, PyObject* args) {
+  Worker* worker;
+  ObjRef objref;
+  long segmentid;
+  long metadata_offset;
+  if (!PyArg_ParseTuple(args, "O&O&ll", &PyObjectToWorker, &worker, &PyObjectToObjRef, &objref, &segmentid, &metadata_offset)) {
+    return NULL;
+  }
+  return worker->finish_buffer(objref, segmentid, metadata_offset);
+}
+
+static PyObject* get_buffer(PyObject* self, PyObject* args) {
+  Worker* worker;
+  ObjRef objref;
+  int64_t size;
+  SegmentId segmentid;
+  if (!PyArg_ParseTuple(args, "O&O&", &PyObjectToWorker, &worker, &PyObjectToObjRef, &objref)) {
+    return NULL;
+  }
+  void* address = reinterpret_cast<void*>(const_cast<char*>(worker->get_buffer(objref, size, segmentid)));
+  std::vector<npy_intp> dim({size});
+  PyObject* t = PyTuple_New(2);
+  PyTuple_SetItem(t, 0, PyArray_SimpleNewFromData(1, dim.data(), NPY_BYTE, address));
+  PyTuple_SetItem(t, 1, PyInt_FromLong(segmentid));
+  return t;
+}
+
+
 static PyObject* get_arrow(PyObject* self, PyObject* args) {
   Worker* worker;
   ObjRef objref;
@@ -926,6 +970,9 @@ static PyMethodDef RayLibMethods[] = {
  { "serialize_object", serialize_object, METH_VARARGS, "serialize an object to protocol buffers" },
  { "deserialize_object", deserialize_object, METH_VARARGS, "deserialize an object from protocol buffers" },
  { "put_arrow", put_arrow, METH_VARARGS, "put an arrow array on the local object store"},
+ { "allocate_buffer", allocate_buffer, METH_VARARGS, "Allocates and returns buffer for objref."},
+ { "finish_buffer", finish_buffer, METH_VARARGS, "Makes the buffer immutable and closes memory segment of objref."},
+ { "get_buffer", get_buffer, METH_VARARGS, "Gets buffer for objref"},
  { "get_arrow", get_arrow, METH_VARARGS, "get an arrow array from the local object store"},
  { "is_arrow", is_arrow, METH_VARARGS, "is the object in the local object store an arrow object?"},
  { "unmap_object", unmap_object, METH_VARARGS, "unmap the object from the client's shared memory pool"},

src/worker.cc

@@ -207,6 +207,47 @@ PyObject* Worker::put_arrow(ObjRef objref, PyObject* value) {
   Py_RETURN_NONE;
 }
 
+const char* Worker::allocate_buffer(ObjRef objref, int64_t size, SegmentId& segmentid) {
+  RAY_CHECK(connected_, "Attempted to perform put_arrow but failed.");
+  ObjRequest request;
+  request.workerid = workerid_;
+  request.type = ObjRequestType::ALLOC;
+  request.objref = objref;
+  request.size = size;
+  RAY_CHECK(request_obj_queue_.send(&request), "error sending over IPC");
+  ObjHandle result;
+  RAY_CHECK(receive_obj_queue_.receive(&result), "error receiving over IPC");
+  const char* address = reinterpret_cast<const char*>(segmentpool_->get_address(result));
+  segmentid = result.segmentid();
+  return address;
+}
+
+PyObject* Worker::finish_buffer(ObjRef objref, SegmentId segmentid, int64_t metadata_offset) {
+  segmentpool_->unmap_segment(segmentid);
+  ObjRequest request;
+  request.workerid = workerid_;
+  request.objref = objref;
+  request.type = ObjRequestType::WORKER_DONE;
+  request.metadata_offset = metadata_offset;
+  RAY_CHECK(request_obj_queue_.send(&request), "error sending over IPC");
+  Py_RETURN_NONE;
+}
+
+const char* Worker::get_buffer(ObjRef objref, int64_t &size, SegmentId& segmentid) {
+  RAY_CHECK(connected_, "Attempted to perform get_arrow but failed.");
+  ObjRequest request;
+  request.workerid = workerid_;
+  request.type = ObjRequestType::GET;
+  request.objref = objref;
+  RAY_CHECK(request_obj_queue_.send(&request), "error sending over IPC");
+  ObjHandle result;
+  RAY_CHECK(receive_obj_queue_.receive(&result), "error receiving over IPC");
+  const char* address = reinterpret_cast<const char*>(segmentpool_->get_address(result));
+  size = result.size();
+  segmentid = result.segmentid();
+  return address;
+}
+
 // returns python list containing the value represented by objref and the
 // segmentid in which the object is stored
 PyObject* Worker::get_arrow(ObjRef objref, SegmentId& segmentid) {

src/worker.h

@@ -61,6 +61,12 @@ class Worker {
   slice get_object(ObjRef objref);
   // stores an arrow object to the local object store
   PyObject* put_arrow(ObjRef objref, PyObject* array);
+  // Allocates buffer for objref with size of size
+  const char* allocate_buffer(ObjRef objref, int64_t size, SegmentId& segmentid);
+  // Finishes buffer with segmentid and an offset of metadata_ofset
+  PyObject* finish_buffer(ObjRef objref, SegmentId segmentid, int64_t metadata_offset);
+  // Gets the buffer for objref
+  const char* get_buffer(ObjRef objref, int64_t& size, SegmentId& segmentid);
   // gets an arrow object from the local object store
   PyObject* get_arrow(ObjRef objref, SegmentId& segmentid);
   // determine if the object stored in objref is an arrow object // TODO(pcm): more general mechanism for this?

test/runtest.py

@@ -136,6 +136,14 @@ class ObjStoreTest(unittest.TestCase):
     self.assertEqual(data[0], result[0])
     self.assertTrue(np.alltrue(data[1] == result[1]))
 
+    # Getting a buffer after modifying it before it finishes should return updated buffer
+    objref =  ray.lib.get_objref(w1.handle)
+    buf = ray.libraylib.allocate_buffer(w1.handle, objref, 100)
+    buf[0][0] = 1
+    ray.libraylib.finish_buffer(w1.handle, objref, buf[1], 0)
+    completedbuffer = ray.libraylib.get_buffer(w1.handle, objref)
+    self.assertEqual(completedbuffer[0][0], 1)
+
     ray.services.cleanup()
 
 class WorkerTest(unittest.TestCase):