[Streaming] Streaming data transfer and python integration (#6185)

This commit is contained in:
Chaokun Yang
2019-12-10 20:33:24 +08:00
committed by Hao Chen
parent c1d4ab8bb4
commit 6272907a57
93 changed files with 8434 additions and 1480 deletions
+153
View File
@@ -0,0 +1,153 @@
# cython: profile=False
# distutils: language = c++
# cython: embedsignature = True
# cython: language_level = 3
# flake8: noqa
from libc.stdint cimport *
from libcpp cimport bool as c_bool
from libcpp.memory cimport shared_ptr
from libcpp.vector cimport vector as c_vector
from libcpp.list cimport list as c_list
from cpython cimport PyObject
cimport cpython
cdef inline object PyObject_to_object(PyObject* o):
# Cast to "object" increments reference count
cdef object result = <object> o
cpython.Py_DECREF(result)
return result
from ray.includes.common cimport (
CLanguage,
CRayObject,
CRayStatus,
CRayFunction
)
from ray.includes.unique_ids cimport (
CActorID,
CJobID,
CTaskID,
CObjectID,
)
from ray.includes.libcoreworker cimport CCoreWorker
cdef extern from "status.h" namespace "ray::streaming" nogil:
cdef cppclass CStreamingStatus "ray::streaming::StreamingStatus":
pass
cdef CStreamingStatus StatusOK "ray::streaming::StreamingStatus::OK"
cdef CStreamingStatus StatusReconstructTimeOut "ray::streaming::StreamingStatus::ReconstructTimeOut"
cdef CStreamingStatus StatusQueueIdNotFound "ray::streaming::StreamingStatus::QueueIdNotFound"
cdef CStreamingStatus StatusResubscribeFailed "ray::streaming::StreamingStatus::ResubscribeFailed"
cdef CStreamingStatus StatusEmptyRingBuffer "ray::streaming::StreamingStatus::EmptyRingBuffer"
cdef CStreamingStatus StatusFullChannel "ray::streaming::StreamingStatus::FullChannel"
cdef CStreamingStatus StatusNoSuchItem "ray::streaming::StreamingStatus::NoSuchItem"
cdef CStreamingStatus StatusInitQueueFailed "ray::streaming::StreamingStatus::InitQueueFailed"
cdef CStreamingStatus StatusGetBundleTimeOut "ray::streaming::StreamingStatus::GetBundleTimeOut"
cdef CStreamingStatus StatusSkipSendEmptyMessage "ray::streaming::StreamingStatus::SkipSendEmptyMessage"
cdef CStreamingStatus StatusInterrupted "ray::streaming::StreamingStatus::Interrupted"
cdef CStreamingStatus StatusWaitQueueTimeOut "ray::streaming::StreamingStatus::WaitQueueTimeOut"
cdef CStreamingStatus StatusOutOfMemory "ray::streaming::StreamingStatus::OutOfMemory"
cdef CStreamingStatus StatusInvalid "ray::streaming::StreamingStatus::Invalid"
cdef CStreamingStatus StatusUnknownError "ray::streaming::StreamingStatus::UnknownError"
cdef CStreamingStatus StatusTailStatus "ray::streaming::StreamingStatus::TailStatus"
cdef cppclass CStreamingCommon "ray::streaming::StreamingCommon":
void SetConfig(const uint8_t *, uint32_t size)
cdef extern from "runtime_context.h" namespace "ray::streaming" nogil:
cdef cppclass CRuntimeContext "ray::streaming::RuntimeContext":
CRuntimeContext()
void SetConfig(const uint8_t *data, uint32_t size)
inline void MarkMockTest()
inline c_bool IsMockTest()
cdef extern from "message/message.h" namespace "ray::streaming" nogil:
cdef cppclass CStreamingMessageType "ray::streaming::StreamingMessageType":
pass
cdef CStreamingMessageType MessageTypeBarrier "ray::streaming::StreamingMessageType::Barrier"
cdef CStreamingMessageType MessageTypeMessage "ray::streaming::StreamingMessageType::Message"
cdef cppclass CStreamingMessage "ray::streaming::StreamingMessage":
inline uint8_t *RawData() const
inline uint32_t GetDataSize() const
inline CStreamingMessageType GetMessageType() const
inline uint64_t GetMessageSeqId() const
cdef extern from "message/message_bundle.h" namespace "ray::streaming" nogil:
cdef cppclass CStreamingMessageBundleType "ray::streaming::StreamingMessageBundleType":
pass
cdef CStreamingMessageBundleType BundleTypeEmpty "ray::streaming::StreamingMessageBundleType::Empty"
cdef CStreamingMessageBundleType BundleTypeBarrier "ray::streaming::StreamingMessageBundleType::Barrier"
cdef CStreamingMessageBundleType BundleTypeBundle "ray::streaming::StreamingMessageBundleType::Bundle"
cdef cppclass CStreamingMessageBundleMeta "ray::streaming::StreamingMessageBundleMeta":
CStreamingMessageBundleMeta()
inline uint64_t GetMessageBundleTs() const
inline uint64_t GetLastMessageId() const
inline uint32_t GetMessageListSize() const
inline CStreamingMessageBundleType GetBundleType() const
inline c_bool IsBarrier()
inline c_bool IsBundle()
ctypedef shared_ptr[CStreamingMessageBundleMeta] CStreamingMessageBundleMetaPtr
uint32_t kMessageBundleHeaderSize "ray::streaming::kMessageBundleHeaderSize"
cdef cppclass CStreamingMessageBundle "ray::streaming::StreamingMessageBundle"(CStreamingMessageBundleMeta):
@staticmethod
void GetMessageListFromRawData(const uint8_t *data, uint32_t size, uint32_t msg_nums,
c_list[shared_ptr[CStreamingMessage]] &msg_list);
cdef extern from "queue/queue_client.h" namespace "ray::streaming" nogil:
cdef cppclass CReaderClient "ray::streaming::ReaderClient":
CReaderClient(CCoreWorker *core_worker,
CRayFunction &async_func,
CRayFunction &sync_func)
void OnReaderMessage(shared_ptr[CLocalMemoryBuffer] buffer);
shared_ptr[CLocalMemoryBuffer] OnReaderMessageSync(shared_ptr[CLocalMemoryBuffer] buffer);
cdef cppclass CWriterClient "ray::streaming::WriterClient":
CWriterClient(CCoreWorker *core_worker,
CRayFunction &async_func,
CRayFunction &sync_func)
void OnWriterMessage(shared_ptr[CLocalMemoryBuffer] buffer);
shared_ptr[CLocalMemoryBuffer] OnWriterMessageSync(shared_ptr[CLocalMemoryBuffer] buffer);
cdef extern from "data_reader.h" namespace "ray::streaming" nogil:
cdef cppclass CDataBundle "ray::streaming::DataBundle":
uint8_t *data
uint32_t data_size
CObjectID c_from "from"
uint64_t seq_id
CStreamingMessageBundleMetaPtr meta
cdef cppclass CDataReader "ray::streaming::DataReader"(CStreamingCommon):
CDataReader(shared_ptr[CRuntimeContext] &runtime_context)
void Init(const c_vector[CObjectID] &input_ids,
const c_vector[CActorID] &actor_ids,
const c_vector[uint64_t] &seq_ids,
const c_vector[uint64_t] &msg_ids,
int64_t timer_interval);
CStreamingStatus GetBundle(const uint32_t timeout_ms,
shared_ptr[CDataBundle] &message)
void Stop()
cdef extern from "data_writer.h" namespace "ray::streaming" nogil:
cdef cppclass CDataWriter "ray::streaming::DataWriter"(CStreamingCommon):
CDataWriter(shared_ptr[CRuntimeContext] &runtime_context)
CStreamingStatus Init(const c_vector[CObjectID] &channel_ids,
const c_vector[CActorID] &actor_ids,
const c_vector[uint64_t] &message_ids,
const c_vector[uint64_t] &queue_size_vec);
long WriteMessageToBufferRing(
const CObjectID &q_id, uint8_t *data, uint32_t data_size)
void Run()
void Stop()
cdef extern from "ray/common/buffer.h" nogil:
cdef cppclass CLocalMemoryBuffer "ray::LocalMemoryBuffer":
uint8_t *Data() const
size_t Size() const
+323
View File
@@ -0,0 +1,323 @@
# flake8: noqa
from libc.stdint cimport *
from libcpp cimport bool as c_bool
from libcpp.memory cimport shared_ptr, make_shared, dynamic_pointer_cast
from libcpp.string cimport string as c_string
from libcpp.vector cimport vector as c_vector
from libcpp.list cimport list as c_list
from ray.includes.common cimport (
CRayFunction,
LANGUAGE_PYTHON,
CBuffer
)
from ray.includes.unique_ids cimport (
CActorID,
CObjectID
)
from ray._raylet cimport (
Buffer,
CoreWorker,
ActorID,
ObjectID,
string_vector_from_list
)
from ray.includes.libcoreworker cimport CCoreWorker
cimport ray.streaming.includes.libstreaming as libstreaming
from ray.streaming.includes.libstreaming cimport (
CStreamingStatus,
CStreamingMessage,
CStreamingMessageBundle,
CRuntimeContext,
CDataBundle,
CDataWriter,
CDataReader,
CReaderClient,
CWriterClient,
CLocalMemoryBuffer,
)
import logging
from ray.function_manager import FunctionDescriptor
channel_logger = logging.getLogger(__name__)
cdef class ReaderClient:
cdef:
CReaderClient *client
def __cinit__(self,
CoreWorker worker,
async_func: FunctionDescriptor,
sync_func: FunctionDescriptor):
cdef:
CCoreWorker *core_worker = worker.core_worker.get()
CRayFunction async_native_func
CRayFunction sync_native_func
async_native_func = CRayFunction(
LANGUAGE_PYTHON, string_vector_from_list(async_func.get_function_descriptor_list()))
sync_native_func = CRayFunction(
LANGUAGE_PYTHON, string_vector_from_list(sync_func.get_function_descriptor_list()))
self.client = new CReaderClient(core_worker, async_native_func, sync_native_func)
def __dealloc__(self):
del self.client
self.client = NULL
def on_reader_message(self, const unsigned char[:] value):
cdef:
size_t size = value.nbytes
shared_ptr[CLocalMemoryBuffer] local_buf = \
make_shared[CLocalMemoryBuffer](<uint8_t *>(&value[0]), size, True)
with nogil:
self.client.OnReaderMessage(local_buf)
def on_reader_message_sync(self, const unsigned char[:] value):
cdef:
size_t size = value.nbytes
shared_ptr[CLocalMemoryBuffer] local_buf = \
make_shared[CLocalMemoryBuffer](<uint8_t *>(&value[0]), size, True)
shared_ptr[CLocalMemoryBuffer] result_buffer
with nogil:
result_buffer = self.client.OnReaderMessageSync(local_buf)
return Buffer.make(dynamic_pointer_cast[CBuffer, CLocalMemoryBuffer](result_buffer))
cdef class WriterClient:
cdef:
CWriterClient * client
def __cinit__(self,
CoreWorker worker,
async_func: FunctionDescriptor,
sync_func: FunctionDescriptor):
cdef:
CCoreWorker *core_worker = worker.core_worker.get()
CRayFunction async_native_func
CRayFunction sync_native_func
async_native_func = CRayFunction(
LANGUAGE_PYTHON, string_vector_from_list(async_func.get_function_descriptor_list()))
sync_native_func = CRayFunction(
LANGUAGE_PYTHON, string_vector_from_list(sync_func.get_function_descriptor_list()))
self.client = new CWriterClient(core_worker, async_native_func, sync_native_func)
def __dealloc__(self):
del self.client
self.client = NULL
def on_writer_message(self, const unsigned char[:] value):
cdef:
size_t size = value.nbytes
shared_ptr[CLocalMemoryBuffer] local_buf = \
make_shared[CLocalMemoryBuffer](<uint8_t *>(&value[0]), size, True)
with nogil:
self.client.OnWriterMessage(local_buf)
def on_writer_message_sync(self, const unsigned char[:] value):
cdef:
size_t size = value.nbytes
shared_ptr[CLocalMemoryBuffer] local_buf = \
make_shared[CLocalMemoryBuffer](<uint8_t *>(&value[0]), size, True)
shared_ptr[CLocalMemoryBuffer] result_buffer
with nogil:
result_buffer = self.client.OnWriterMessageSync(local_buf)
return Buffer.make(dynamic_pointer_cast[CBuffer, CLocalMemoryBuffer](result_buffer))
cdef class DataWriter:
cdef:
CDataWriter *writer
def __init__(self):
raise Exception("use create() to create DataWriter")
@staticmethod
def create(list py_output_channels,
list output_actor_ids: list[ActorID],
uint64_t queue_size,
list py_msg_ids,
bytes config_bytes,
c_bool is_mock):
cdef:
c_vector[CObjectID] channel_ids = bytes_list_to_qid_vec(py_output_channels)
c_vector[CActorID] actor_ids
c_vector[uint64_t] msg_ids
CDataWriter *c_writer
cdef const unsigned char[:] config_data
for actor_id in output_actor_ids:
actor_ids.push_back((<ActorID>actor_id).data)
for py_msg_id in py_msg_ids:
msg_ids.push_back(<uint64_t>py_msg_id)
cdef shared_ptr[CRuntimeContext] ctx = make_shared[CRuntimeContext]()
if is_mock:
ctx.get().MarkMockTest()
if config_bytes:
config_data = config_bytes
channel_logger.info("load config, config bytes size: %s", config_data.nbytes)
ctx.get().SetConfig(<uint8_t *>(&config_data[0]), config_data.nbytes)
c_writer = new CDataWriter(ctx)
cdef:
c_vector[CObjectID] remain_id_vec
c_vector[uint64_t] queue_size_vec
for i in range(channel_ids.size()):
queue_size_vec.push_back(queue_size)
cdef CStreamingStatus status = c_writer.Init(channel_ids, actor_ids, msg_ids, queue_size_vec)
if remain_id_vec.size() != 0:
channel_logger.warning("failed queue amounts => %s", remain_id_vec.size())
if <uint32_t>status != <uint32_t> libstreaming.StatusOK:
msg = "initialize writer failed, status={}".format(<uint32_t>status)
channel_logger.error(msg)
del c_writer
import ray.streaming.runtime.transfer as transfer
raise transfer.ChannelInitException(msg, qid_vector_to_list(remain_id_vec))
c_writer.Run()
channel_logger.info("create native writer succeed")
cdef DataWriter writer = DataWriter.__new__(DataWriter)
writer.writer = c_writer
return writer
def __dealloc__(self):
if self.writer != NULL:
del self.writer
channel_logger.info("deleted DataWriter")
self.writer = NULL
def write(self, ObjectID qid, const unsigned char[:] value):
"""support zero-copy bytes, bytearray, array of unsigned char"""
cdef:
CObjectID native_id = qid.data
uint64_t msg_id
uint8_t *data = <uint8_t *>(&value[0])
uint32_t size = value.nbytes
with nogil:
msg_id = self.writer.WriteMessageToBufferRing(native_id, data, size)
return msg_id
def stop(self):
self.writer.Stop()
channel_logger.info("stopped DataWriter")
cdef class DataReader:
cdef:
CDataReader *reader
readonly bytes meta
readonly bytes data
def __init__(self):
raise Exception("use create() to create DataReader")
@staticmethod
def create(list py_input_queues,
list input_actor_ids: list[ActorID],
list py_seq_ids,
list py_msg_ids,
int64_t timer_interval,
c_bool is_recreate,
bytes config_bytes,
c_bool is_mock):
cdef:
c_vector[CObjectID] queue_id_vec = bytes_list_to_qid_vec(py_input_queues)
c_vector[CActorID] actor_ids
c_vector[uint64_t] seq_ids
c_vector[uint64_t] msg_ids
CDataReader *c_reader
cdef const unsigned char[:] config_data
for actor_id in input_actor_ids:
actor_ids.push_back((<ActorID>actor_id).data)
for py_seq_id in py_seq_ids:
seq_ids.push_back(<uint64_t>py_seq_id)
for py_msg_id in py_msg_ids:
msg_ids.push_back(<uint64_t>py_msg_id)
cdef shared_ptr[CRuntimeContext] ctx = make_shared[CRuntimeContext]()
if config_bytes:
config_data = config_bytes
channel_logger.info("load config, config bytes size: %s", config_data.nbytes)
ctx.get().SetConfig(<uint8_t *>(&(config_data[0])), config_data.nbytes)
if is_mock:
ctx.get().MarkMockTest()
c_reader = new CDataReader(ctx)
c_reader.Init(queue_id_vec, actor_ids, seq_ids, msg_ids, timer_interval)
channel_logger.info("create native reader succeed")
cdef DataReader reader = DataReader.__new__(DataReader)
reader.reader = c_reader
return reader
def __dealloc__(self):
if self.reader != NULL:
del self.reader
channel_logger.info("deleted DataReader")
self.reader = NULL
def read(self, uint32_t timeout_millis):
cdef:
shared_ptr[CDataBundle] bundle
CStreamingStatus status
with nogil:
status = self.reader.GetBundle(timeout_millis, bundle)
cdef uint32_t bundle_type = <uint32_t>(bundle.get().meta.get().GetBundleType())
if <uint32_t> status != <uint32_t> libstreaming.StatusOK:
if <uint32_t> status == <uint32_t> libstreaming.StatusInterrupted:
# avoid cyclic import
import ray.streaming.runtime.transfer as transfer
raise transfer.ChannelInterruptException("reader interrupted")
elif <uint32_t> status == <uint32_t> libstreaming.StatusInitQueueFailed:
raise Exception("init channel failed")
elif <uint32_t> status == <uint32_t> libstreaming.StatusWaitQueueTimeOut:
raise Exception("wait channel object timeout")
cdef:
uint32_t msg_nums
CObjectID queue_id
c_list[shared_ptr[CStreamingMessage]] msg_list
list msgs = []
uint64_t timestamp
uint64_t msg_id
if bundle_type == <uint32_t> libstreaming.BundleTypeBundle:
msg_nums = bundle.get().meta.get().GetMessageListSize()
CStreamingMessageBundle.GetMessageListFromRawData(
bundle.get().data + libstreaming.kMessageBundleHeaderSize,
bundle.get().data_size - libstreaming.kMessageBundleHeaderSize,
msg_nums,
msg_list)
timestamp = bundle.get().meta.get().GetMessageBundleTs()
for msg in msg_list:
msg_bytes = msg.get().RawData()[:msg.get().GetDataSize()]
qid_bytes = queue_id.Binary()
msg_id = msg.get().GetMessageSeqId()
msgs.append((msg_bytes, msg_id, timestamp, qid_bytes))
return msgs
elif bundle_type == <uint32_t> libstreaming.BundleTypeEmpty:
return []
else:
raise Exception("Unsupported bundle type {}".format(bundle_type))
def stop(self):
self.reader.Stop()
channel_logger.info("stopped DataReader")
cdef c_vector[CObjectID] bytes_list_to_qid_vec(list py_queue_ids) except *:
assert len(py_queue_ids) > 0
cdef:
c_vector[CObjectID] queue_id_vec
c_string q_id_data
for q_id in py_queue_ids:
q_id_data = q_id
assert q_id_data.size() == CObjectID.Size()
obj_id = CObjectID.FromBinary(q_id_data)
queue_id_vec.push_back(obj_id)
return queue_id_vec
cdef c_vector[c_string] qid_vector_to_list(c_vector[CObjectID] queue_id_vec):
queues = []
for obj_id in queue_id_vec:
queues.append(obj_id.Binary())
return queues