From b68512ec7db523b4857d015c153ae7dfa8d7af70 Mon Sep 17 00:00:00 2001 From: Wapaul1 Date: Wed, 20 Jul 2016 17:36:15 -0700 Subject: [PATCH] Added handles for python to access buffers in objstore (#282) --- src/raylib.cc | 47 +++++++++++++++++++++++++++++++++++++++++++++++ src/worker.cc | 41 +++++++++++++++++++++++++++++++++++++++++ src/worker.h | 6 ++++++ test/runtest.py | 8 ++++++++ 4 files changed, 102 insertions(+) diff --git a/src/raylib.cc b/src/raylib.cc index b84478d35..bf0a3f7dd 100644 --- a/src/raylib.cc +++ b/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(const_cast(worker->allocate_buffer(objref, size, segmentid))); + std::vector 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(const_cast(worker->get_buffer(objref, size, segmentid))); + std::vector 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"}, diff --git a/src/worker.cc b/src/worker.cc index a2ecb601c..e2b8bed88 100644 --- a/src/worker.cc +++ b/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(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(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) { diff --git a/src/worker.h b/src/worker.h index 7a26da407..14d8141a4 100644 --- a/src/worker.h +++ b/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? diff --git a/test/runtest.py b/test/runtest.py index ba7087e02..7c91c940b 100644 --- a/test/runtest.py +++ b/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):