def write(self, content: Union[bytes, MemoryPointer]):
if not self._initialized:
self._write_init()
self._initialized = True
cupy_pointer = MemoryPointer(self._cupy_memory, self._offset)
if isinstance(content, bytes):
content_length = len(content)
source_mem = np.frombuffer(content, dtype='uint8').ctypes.data_as(
ctypes.c_void_p)
else:
content_length = content.mem.size
source_mem = content
cupy_pointer.copy_from(source_mem, content_length)
self._offset += content_length
def read(self, size=-1) -> Buffer:
if not self._initialized:
self._read_init()
self._initialized = True
size = self._size if size < 0 else size
cupy_pointer = MemoryPointer(self._cupy_memory, self._offset)
self._offset += size
return Buffer(cupy_pointer.ptr, size=size)
def ascupy(ceo_cu_array):
from cupy.cuda import UnownedMemory, MemoryPointer
from cupy import ndarray
ptr = UnownedMemory(ceo_cu_array.dev_ptr, ceo_cu_array.nbytes,
ceo_cu_array)
return ndarray(shape=ceo_cu_array.shape,
dtype=ceo_cu_array.type,
memptr=MemoryPointer(ptr, 0))
def read(self, size=-1):
from cudf.core import Buffer
from cupy.cuda import MemoryPointer
if not self._initialized:
self._read_init()
self._initialized = True
size = self._size if size < 0 else size
cupy_pointer = MemoryPointer(self._cupy_memory, self._offset)
self._offset += size
return Buffer(cupy_pointer.ptr, size=size)
def write(self, content):
from cudf.core.buffer import Buffer
from cupy.cuda import MemoryPointer
from cupy.cuda.memory import UnownedMemory
if not self._has_write_headers:
self._headers = headers = pickle.loads(content)
buffer_types = headers["buffer_types"]
for buffer_type, size in buffer_types:
if buffer_type == "cuda":
self._buffers.append(Buffer.empty(size))
else:
self._buffers.append(BytesIO())
self._has_write_headers = True
return
cur_buf = self._buffers[self._cur_buffer_index]
cur_buf_size = self._headers["buffer_types"][self._cur_buffer_index][1]
if isinstance(cur_buf, Buffer):
cur_cupy_memory = UnownedMemory(cur_buf.ptr, len(cur_buf), cur_buf)
cupy_pointer = MemoryPointer(cur_cupy_memory, self._offset)
if isinstance(content, bytes):
content_length = len(content)
source_mem = np.frombuffer(
content, dtype="uint8").ctypes.data_as(ctypes.c_void_p)
else:
source_mem = MemoryPointer(
UnownedMemory(content.ptr, len(content), content), 0)
content_length = source_mem.mem.size
cupy_pointer.copy_from(source_mem, content_length)
else:
content_length = len(content)
cur_buf.write(content)
if content_length + self._offset >= cur_buf_size:
if isinstance(cur_buf, BytesIO):
self._buffers[self._cur_buffer_index] = cur_buf.getvalue()
self._cur_buffer_index += 1
self._offset = 0
else:
self._offset += content_length
def read(self, size: int):
# we read cuda_header first and then read cuda buffers one by one,
# the return value's size is not exactly the specified size.
from cudf.core.buffer import Buffer
from cupy.cuda import MemoryPointer
from cupy.cuda.memory import UnownedMemory
if not self._has_read_headers:
self._has_read_headers = True
return pickle.dumps(self._headers)
if len(self._buffers) == 0:
return ""
cur_buf = self._buffers[0]
# current buf read to end
if cur_buf is None:
# empty cuda buffer
content = Buffer.empty(0)
self._offset = 0
self._buffers.pop(0)
return content
elif size >= cur_buf.size - self._offset:
if isinstance(cur_buf, UnownedMemory):
cupy_pointer = MemoryPointer(cur_buf, self._offset)
content = Buffer(cupy_pointer.ptr,
size=cur_buf.size - self._offset)
else:
content = cur_buf[self._offset:self._offset + size]
self._offset = 0
self._buffers.pop(0)
return content
else:
if isinstance(cur_buf, UnownedMemory):
cupy_pointer = MemoryPointer(cur_buf, self._offset)
self._offset += size
return Buffer(cupy_pointer.ptr, size=size)
else:
self._offset += size
return cur_buf[self._offset, self._offset + size]