From 4273e04657fcd27990f38787bb982ba829174e7f Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Wed, 2 Aug 2023 13:45:41 +0200 Subject: [PATCH 01/22] Introduce and use BufferOwner --- python/cudf/cudf/core/abc.py | 9 +- python/cudf/cudf/core/buffer/__init__.py | 6 +- python/cudf/cudf/core/buffer/buffer.py | 241 ++++++++++++----- .../core/buffer/exposure_tracked_buffer.py | 192 ++++--------- .../cudf/cudf/core/buffer/spillable_buffer.py | 256 ++++++++---------- python/cudf/cudf/core/buffer/utils.py | 16 +- python/cudf/cudf/tests/test_copying.py | 13 +- python/cudf/cudf/tests/test_spilling.py | 41 +-- 8 files changed, 389 insertions(+), 385 deletions(-) diff --git a/python/cudf/cudf/core/abc.py b/python/cudf/cudf/core/abc.py index adf9fe39e4f..a833663ce00 100644 --- a/python/cudf/cudf/core/abc.py +++ b/python/cudf/cudf/core/abc.py @@ -89,7 +89,14 @@ def device_serialize(self): """ header, frames = self.serialize() assert all( - isinstance(f, (cudf.core.buffer.Buffer, memoryview)) + isinstance( + f, + ( + cudf.core.buffer.BufferOwner, + cudf.core.buffer.Buffer, + memoryview, + ), + ) for f in frames ) header["type-serialized"] = pickle.dumps(type(self)) diff --git a/python/cudf/cudf/core/buffer/__init__.py b/python/cudf/cudf/core/buffer/__init__.py index d8883bd97e5..b5799c076f9 100644 --- a/python/cudf/cudf/core/buffer/__init__.py +++ b/python/cudf/cudf/core/buffer/__init__.py @@ -1,6 +1,10 @@ # Copyright (c) 2022-2023, NVIDIA CORPORATION. -from cudf.core.buffer.buffer import Buffer, cuda_array_interface_wrapper +from cudf.core.buffer.buffer import ( + Buffer, + BufferOwner, + cuda_array_interface_wrapper, +) from cudf.core.buffer.exposure_tracked_buffer import ExposureTrackedBuffer from cudf.core.buffer.spillable_buffer import SpillableBuffer, SpillLock from cudf.core.buffer.utils import ( diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index 59d20a2784d..e6f285dd1a8 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -5,7 +5,17 @@ import math import pickle from types import SimpleNamespace -from typing import Any, Dict, Literal, Mapping, Optional, Sequence, Tuple +from typing import ( + Any, + Dict, + Literal, + Mapping, + Optional, + Sequence, + Tuple, + Type, + TypeVar, +) import numpy from typing_extensions import Self @@ -16,6 +26,32 @@ from cudf.core.abc import Serializable from cudf.utils.string import format_bytes +T = TypeVar("T") + + +def get_owner(data, klass: Type[T]) -> Optional[T]: + """Get the owner of `data`, if any exist + + Search through the stack of data owners in order to find an + owner of type `klass` (not subclasses). + + Parameters + ---------- + data + The data object + + Return + ------ + klass or None + The owner of `data` if `klass` or None. + """ + + if type(data) is klass: + return data + if hasattr(data, "owner"): + return get_owner(data.owner, klass) + return None + def host_memory_allocation(nbytes: int) -> memoryview: """Allocate host memory using NumPy @@ -90,22 +126,13 @@ def cuda_array_interface_wrapper( ) -class Buffer(Serializable): - """A Buffer represents device memory. - - Use the factory function `as_buffer` to create a Buffer instance. - """ +class BufferOwner(Serializable): + """A owning buffer represents device memory.""" _ptr: int _size: int _owner: object - def __init__(self): - raise ValueError( - f"do not create a {self.__class__} directly, please " - "use the factory function `cudf.core.buffer.as_buffer`" - ) - @classmethod def _from_device_memory(cls, data: Any) -> Self: """Create a Buffer from an object exposing `__cuda_array_interface__`. @@ -168,53 +195,6 @@ def _from_host_memory(cls, data: Any) -> Self: # Create from device memory return cls._from_device_memory(buf) - def _getitem(self, offset: int, size: int) -> Self: - """ - Sub-classes can overwrite this to implement __getitem__ - without having to handle non-slice inputs. - """ - return self._from_device_memory( - cuda_array_interface_wrapper( - ptr=self.get_ptr(mode="read") + offset, - size=size, - owner=self.owner, - ) - ) - - def __getitem__(self, key: slice) -> Self: - """Create a new slice of the buffer.""" - if not isinstance(key, slice): - raise TypeError( - "Argument 'key' has incorrect type " - f"(expected slice, got {key.__class__.__name__})" - ) - start, stop, step = key.indices(self.size) - if step != 1: - raise ValueError("slice must be C-contiguous") - return self._getitem(offset=start, size=stop - start) - - def copy(self, deep: bool = True) -> Self: - """ - Return a copy of Buffer. - - Parameters - ---------- - deep : bool, default True - If True, returns a deep copy of the underlying Buffer data. - If False, returns a shallow copy of the Buffer pointing to - the same underlying data. - - Returns - ------- - Buffer - """ - if deep: - return self._from_device_memory( - rmm.DeviceBuffer(ptr=self.get_ptr(mode="read"), size=self.size) - ) - else: - return self[:] - @property def size(self) -> int: """Size of the buffer in bytes.""" @@ -277,6 +257,131 @@ def memoryview( ) return memoryview(host_buf).toreadonly() + def __repr__(self) -> str: + return ( + f"<{self.__class__.__name__} size={format_bytes(self._size)} " + f"ptr={hex(self._ptr)} owner={repr(self._owner)}>" + ) + + +class Buffer(Serializable): + """A Buffer represents memory. + + Use the factory function `as_buffer` to create a Buffer instance. + """ + + def __init__( + self, + *, + owner: BufferOwner, + offset: int, + size: int, + ) -> None: + if size < 0: + raise ValueError("size cannot be negative") + if offset < 0: + raise ValueError("offset cannot be negative") + if offset + size > owner.size: + raise ValueError( + "offset+size cannot be greater than the size of owner" + ) + self._owner = owner + self._offset = offset + self._size = size + + @property + def size(self) -> int: + """Size of the buffer in bytes.""" + return self._size + + @property + def nbytes(self) -> int: + """Size of the buffer in bytes.""" + return self._size + + @property + def owner(self) -> Any: + """Object owning the memory of the buffer.""" + return self._owner + + def __getitem__(self, key: slice) -> Self: + """Create a new slice of the buffer.""" + if not isinstance(key, slice): + raise TypeError( + "Argument 'key' has incorrect type " + f"(expected slice, got {key.__class__.__name__})" + ) + start, stop, step = key.indices(self.size) + if step != 1: + raise ValueError("slice must be C-contiguous") + return self.__class__( + owner=self._owner, offset=self._offset + start, size=stop - start + ) + + def get_ptr(self, *, mode: Literal["read", "write"]) -> int: + return self._owner.get_ptr(mode=mode) + self._offset + + def memoryview( + self, *, offset: int = 0, size: Optional[int] = None + ) -> memoryview: + size = self._size if size is None else size + return self._owner.memoryview(offset=self._offset + offset, size=size) + + def copy(self, deep: bool = True) -> Self: + """Return a copy of Buffer. + + What actually happens when `deep == False` depends on the + "copy_on_write" option. When copy-on-write is enabled, a shallow copy + becomes a deep copy if the buffer has been exposed. This is because we + have no control over knowing if the data is being modified when the + buffer has been exposed to third-party. + + Parameters + ---------- + deep : bool, default True + The semantics when copy-on-write is disabled: + - If deep=True, returns a deep copy of the underlying data. + - If deep=False, returns a shallow copy of the Buffer pointing + to the same underlying data. + The semantics when copy-on-write is enabled: + - From the users perspective, always a deep copy of the + underlying data. However, the data isn't actually copied + until someone writers to the returned buffer. + + Returns + ------- + BufferSlice + A slice pointing to either a new or the existing owner + depending on the expose status of the owner and the + copy-on-write option (see above). + """ + + # When doing a shallow copy, we just return a new slice + if not deep: + return self.__class__( + owner=self._owner, offset=self._offset, size=self._size + ) + + # Otherwise, we create a new copy of the memory this slice represents + owner = self._owner._from_device_memory( + rmm.DeviceBuffer( + ptr=self._owner.get_ptr(mode="read") + self._offset, + size=self.size, + ) + ) + return self.__class__(owner=owner, offset=0, size=owner.size) + + @property + def __cuda_array_interface__(self) -> Mapping: + """Implementation of the CUDA Array Interface.""" + return { + "data": (self.get_ptr(mode="write"), False), + "shape": (self.size,), + "strides": None, + "typestr": "|u1", + "version": 0, + } + def serialize(self) -> Tuple[dict, list]: """Serialize the buffer into header and frames. @@ -291,6 +396,7 @@ def serialize(self) -> Tuple[dict, list]: """ header: Dict[str, Any] = {} header["type-serialized"] = pickle.dumps(type(self)) + header["owner-type-serialized"] = pickle.dumps(type(self._owner)) header["frame_count"] = 1 frames = [self] return header, frames @@ -317,16 +423,21 @@ def deserialize(cls, header: dict, frames: list) -> Self: if isinstance(frame, cls): return frame # The frame is already deserialized + owner_type: BufferOwner = pickle.loads(header["owner-type-serialized"]) if hasattr(frame, "__cuda_array_interface__"): - return cls._from_device_memory(frame) - return cls._from_host_memory(frame) + owner = owner_type._from_device_memory(frame) + else: + owner = owner_type._from_host_memory(frame) + return cls( + owner=owner, + offset=0, + size=owner.size, + ) def __repr__(self) -> str: - klass = self.__class__ - name = f"{klass.__module__}.{klass.__qualname__}" return ( - f"<{name} size={format_bytes(self._size)} " - f"ptr={hex(self._ptr)} owner={repr(self._owner)}>" + f"<{self.__class__.__name__} size={format_bytes(self._size)} " + f"offset={format_bytes(self._offset)} of {self._owner}>" ) diff --git a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py index f2ac6301944..e90a64fc906 100644 --- a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py +++ b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py @@ -3,52 +3,21 @@ from __future__ import annotations import weakref -from typing import ( - Any, - Container, - Literal, - Mapping, - Optional, - Type, - TypeVar, - cast, -) +from typing import Any, Literal, Mapping, Optional, Type from typing_extensions import Self import cudf -from cudf.core.buffer.buffer import Buffer, get_ptr_and_size -from cudf.utils.string import format_bytes - -T = TypeVar("T", bound="ExposureTrackedBuffer") - - -def get_owner(data, klass: Type[T]) -> Optional[T]: - """Get the owner of `data`, if any exist - - Search through the stack of data owners in order to find an - owner of type `klass` (not subclasses). - - Parameters - ---------- - data - The data object - - Return - ------ - klass or None - The owner of `data` if `klass` or None. - """ - - if type(data) is klass: - return data - if hasattr(data, "owner"): - return get_owner(data.owner, klass) - return None +from cudf.core.buffer.buffer import ( + Buffer, + BufferOwner, + get_owner, + get_ptr_and_size, +) def as_exposure_tracked_buffer( - data, exposed: bool, subclass: Optional[Type[T]] = None + data, exposed: bool, subclass: Optional[Type[ExposureTrackedBuffer]] = None ) -> BufferSlice: """Factory function to wrap `data` in a slice of an exposure tracked buffer @@ -89,27 +58,26 @@ def as_exposure_tracked_buffer( if not hasattr(data, "__cuda_array_interface__"): if exposed: raise ValueError("cannot created exposed host memory") - return cast( - BufferSlice, ExposureTrackedBuffer._from_host_memory(data)[:] - ) + tracked_buf = ExposureTrackedBuffer._from_host_memory(data) + return BufferSlice(owner=tracked_buf, offset=0, size=tracked_buf.size) owner = get_owner(data, subclass or ExposureTrackedBuffer) - if owner is None: - return cast( - BufferSlice, - ExposureTrackedBuffer._from_device_memory(data, exposed=exposed)[ - : - ], - ) - - # At this point, we know that `data` is owned by a exposure tracked buffer - ptr, size = get_ptr_and_size(data.__cuda_array_interface__) - if size > 0 and owner._ptr == 0: - raise ValueError("Cannot create a non-empty slice of a null buffer") - return BufferSlice(base=owner, offset=ptr - owner._ptr, size=size) - - -class ExposureTrackedBuffer(Buffer): + if owner is not None: + # `data` is owned by an exposure tracked buffer + ptr, size = get_ptr_and_size(data.__cuda_array_interface__) + base_ptr = owner.get_ptr(mode="read") + if size > 0 and base_ptr == 0: + raise ValueError( + "Cannot create a non-empty slice of a null buffer" + ) + return BufferSlice(owner=owner, offset=ptr - base_ptr, size=size) + + # `data` is new device memory + owner = ExposureTrackedBuffer._from_device_memory(data, exposed=exposed) + return BufferSlice(owner=owner, offset=0, size=owner.size) + + +class ExposureTrackedBuffer(BufferOwner): """A Buffer that tracks its "expose" status. In order to implement copy-on-write and spillable buffers, we need the @@ -140,102 +108,49 @@ def mark_exposed(self) -> None: @classmethod def _from_device_memory(cls, data: Any, *, exposed: bool = False) -> Self: - """Create an exposure tracked buffer from device memory. - - No data is being copied. - - Parameters - ---------- - data : device-buffer-like - An object implementing the CUDA Array Interface. - exposed : bool, optional - Mark the buffer as permanently exposed. - - Returns - ------- - ExposureTrackedBuffer - Buffer representing the same device memory as `data` - """ ret = super()._from_device_memory(data) ret._exposed = exposed ret._slices = weakref.WeakSet() return ret - def _getitem(self, offset: int, size: int) -> BufferSlice: - return BufferSlice(base=self, offset=offset, size=size) - @property def __cuda_array_interface__(self) -> Mapping: self.mark_exposed() return super().__cuda_array_interface__ - def __repr__(self) -> str: - return ( - f"" - ) - -class BufferSlice(ExposureTrackedBuffer): +class BufferSlice(Buffer): """A slice (aka. a view) of a exposure tracked buffer. Parameters ---------- - base + owner The exposure tracked buffer this slice refers to. offset - The offset relative to the start memory of base (in bytes). + The offset relative to the start memory of owner (in bytes). size The size of the slice (in bytes) - passthrough_attributes - Name of attributes that are passed through to the base as-is. """ + _owner: ExposureTrackedBuffer + def __init__( self, - base: ExposureTrackedBuffer, + owner: ExposureTrackedBuffer, offset: int, size: int, - *, - passthrough_attributes: Container[str] = ("exposed",), ) -> None: - if size < 0: - raise ValueError("size cannot be negative") - if offset < 0: - raise ValueError("offset cannot be negative") - if offset + size > base.size: - raise ValueError( - "offset+size cannot be greater than the size of base" - ) - self._base = base - self._offset = offset - self._size = size - self._owner = base - self._passthrough_attributes = passthrough_attributes - base._slices.add(self) - - def __getattr__(self, name): - if name in self._passthrough_attributes: - return getattr(self._base, name) - raise AttributeError( - f"{self.__class__.__name__} object has no attribute {name}" - ) - - def _getitem(self, offset: int, size: int) -> BufferSlice: - return BufferSlice( - base=self._base, offset=offset + self._offset, size=size - ) + super().__init__(owner=owner, offset=offset, size=size) + self._owner._slices.add(self) + + @property + def exposed(self) -> bool: + return self._owner.exposed def get_ptr(self, *, mode: Literal["read", "write"]) -> int: if mode == "write" and cudf.get_option("copy_on_write"): self.make_single_owner_inplace() - return self._base.get_ptr(mode=mode) + self._offset - - def memoryview( - self, *, offset: int = 0, size: Optional[int] = None - ) -> memoryview: - return self._base.memoryview(offset=self._offset + offset, size=size) + return super().get_ptr(mode=mode) def copy(self, deep: bool = True) -> Self: """Return a copy of Buffer. @@ -261,15 +176,13 @@ def copy(self, deep: bool = True) -> Self: Returns ------- BufferSlice - A slice pointing to either a new or the existing base buffer - depending on the expose status of the base buffer and the + A slice pointing to either a new or the existing owner + depending on the expose status of the owner and the copy-on-write option (see above). """ if cudf.get_option("copy_on_write"): - base_copy = self._base.copy(deep=deep or self.exposed) - else: - base_copy = self._base.copy(deep=deep) - return cast(Self, base_copy[self._offset : self._offset + self._size]) + return super().copy(deep=deep or self.exposed) + return super().copy(deep=deep) @property def __cuda_array_interface__(self) -> Mapping: @@ -278,7 +191,7 @@ def __cuda_array_interface__(self) -> Mapping: return super().__cuda_array_interface__ def make_single_owner_inplace(self) -> None: - """Make sure this slice is the only one pointing to the base. + """Make sure this slice is the only one pointing to the owner. This is used by copy-on-write to trigger a deep copy when write access is detected. @@ -294,18 +207,11 @@ def make_single_owner_inplace(self) -> None: Buffer representing the same device memory as `data` """ - if len(self._base._slices) > 1: - # If this is not the only slice pointing to `self._base`, we - # point to a new deep copy of the base. + if len(self._owner._slices) > 1: + # If this is not the only slice pointing to `self._owner`, we + # point to a new deep copy of the owner. t = self.copy(deep=True) - self._base = t._base + self._owner = t._owner self._offset = t._offset self._size = t._size - self._owner = t._base - self._base._slices.add(self) - - def __repr__(self) -> str: - return ( - f"" - ) + self._owner._slices.add(self) diff --git a/python/cudf/cudf/core/buffer/spillable_buffer.py b/python/cudf/cudf/core/buffer/spillable_buffer.py index 84fb2044c62..87c132ef238 100644 --- a/python/cudf/cudf/core/buffer/spillable_buffer.py +++ b/python/cudf/cudf/core/buffer/spillable_buffer.py @@ -16,7 +16,9 @@ from cudf.core.buffer.buffer import ( Buffer, + BufferOwner, cuda_array_interface_wrapper, + get_owner, get_ptr_and_size, host_memory_allocation, ) @@ -26,32 +28,8 @@ from cudf.core.buffer.spill_manager import SpillManager -def get_spillable_owner(data) -> Optional[SpillableBuffer]: - """Get the spillable owner of `data`, if any exist - - Search through the stack of data owners in order to find an - owner of type `SpillableBuffer` (not subclasses). - - Parameters - ---------- - data : buffer-like or array-like - A buffer-like or array-like object that represent C-contiguous memory. - - Return - ------ - SpillableBuffer or None - The owner of `data` if spillable or None. - """ - - if type(data) is SpillableBuffer: - return data - if hasattr(data, "owner"): - return get_spillable_owner(data.owner) - return None - - -def as_spillable_buffer(data, exposed: bool) -> SpillableBuffer: - """Factory function to wrap `data` in a SpillableBuffer object. +def as_spillable_buffer(data, exposed: bool) -> SpillableBufferSlice: + """Factory function to wrap `data` in a SpillableBufferSlice object. If `data` isn't a buffer already, a new buffer that points to the memory of `data` is created. If `data` represents host memory, it is copied to a new @@ -76,7 +54,7 @@ def as_spillable_buffer(data, exposed: bool) -> SpillableBuffer: Return ------ - SpillableBuffer + SpillableBufferSlice A spillabe buffer instance that represents the device memory of `data`. """ @@ -85,25 +63,41 @@ def as_spillable_buffer(data, exposed: bool) -> SpillableBuffer: if not hasattr(data, "__cuda_array_interface__"): if exposed: raise ValueError("cannot created exposed host memory") - return SpillableBuffer._from_host_memory(data) + tracked_buf = SpillableBuffer._from_host_memory(data) + return SpillableBufferSlice( + owner=tracked_buf, offset=0, size=tracked_buf.size + ) - spillable_owner = get_spillable_owner(data) - if spillable_owner is None: - return SpillableBuffer._from_device_memory(data, exposed=exposed) + if not hasattr(data, "__cuda_array_interface__"): + if exposed: + raise ValueError("cannot created exposed host memory") + spillabe_buf = SpillableBuffer._from_host_memory(data) + return SpillableBufferSlice( + owner=spillabe_buf, offset=0, size=spillabe_buf.size + ) - if not spillable_owner.exposed and get_spill_lock() is None: - raise ValueError( - "A owning spillable buffer must " - "either be exposed or spilled locked." + owner = get_owner(data, SpillableBuffer) + if owner is not None: + if not owner.exposed and get_spill_lock() is None: + raise ValueError( + "A owning spillable buffer must " + "either be exposed or spilled locked." + ) + # `data` is owned by an spillable buffer, which is exposed or + # spilled locked. + ptr, size = get_ptr_and_size(data.__cuda_array_interface__) + base_ptr = owner.get_ptr(mode="read") + if size > 0 and owner._ptr == 0: + raise ValueError( + "Cannot create a non-empty slice of a null buffer" + ) + return SpillableBufferSlice( + owner=owner, offset=ptr - base_ptr, size=size ) - # At this point, we know that `data` is owned by a spillable buffer, - # which is exposed or spilled locked. - ptr, size = get_ptr_and_size(data.__cuda_array_interface__) - base_ptr = spillable_owner.memory_info()[0] - return SpillableBufferSlice( - spillable_owner, offset=ptr - base_ptr, size=size - ) + # `data` is new device memory + owner = SpillableBuffer._from_device_memory(data, exposed=exposed) + return SpillableBufferSlice(owner=owner, offset=0, size=owner.size) class SpillLock: @@ -140,7 +134,7 @@ def __getitem__(self, i): raise IndexError("tuple index out of range") -class SpillableBuffer(Buffer): +class SpillableBuffer(BufferOwner): """A Buffer that supports spilling memory off the GPU to avoid OOMs. This buffer supports spilling the represented data to host memory. @@ -265,11 +259,6 @@ def _from_host_memory(cls, data: Any) -> Self: def is_spilled(self) -> bool: return self._ptr_desc["type"] != "gpu" - def copy(self, deep: bool = True) -> Self: - spill_lock = SpillLock() - self.spill_lock(spill_lock=spill_lock) - return super().copy(deep=deep) - def spill(self, target: str = "cpu") -> None: """Spill or un-spill this buffer in-place @@ -451,52 +440,6 @@ def memoryview( ) return ret - def _getitem(self, offset: int, size: int) -> SpillableBufferSlice: - return SpillableBufferSlice(base=self, offset=offset, size=size) - - def serialize(self) -> Tuple[dict, list]: - """Serialize the Buffer - - Normally, we would use `[self]` as the frames. This would work but - also mean that `self` becomes exposed permanently if the frames are - later accessed through `__cuda_array_interface__`, which is exactly - what libraries like Dask+UCX would do when communicating! - - The sound solution is to modify Dask et al. so that they access the - frames through `.get_ptr()` and holds on to the `spill_lock` until - the frame has been transferred. However, until this adaptation we - use a hack where the frame is a `Buffer` with a `spill_lock` as the - owner, which makes `self` unspillable while the frame is alive but - doesn't expose `self` when `__cuda_array_interface__` is accessed. - - Warning, this hack means that the returned frame must be copied before - given to `.deserialize()`, otherwise we would have a `Buffer` pointing - to memory already owned by an existing `SpillableBuffer`. - """ - header: Dict[Any, Any] - frames: List[Buffer | memoryview] - with self.lock: - header = {} - header["type-serialized"] = pickle.dumps(self.__class__) - header["frame_count"] = 1 - if self.is_spilled: - frames = [self.memoryview()] - else: - # TODO: Use `frames=[self]` instead of this hack, see doc above - spill_lock = SpillLock() - self.spill_lock(spill_lock) - ptr, size, _ = self.memory_info() - frames = [ - Buffer._from_device_memory( - cuda_array_interface_wrapper( - ptr=ptr, - size=size, - owner=(self._owner, spill_lock), - ) - ) - ] - return header, frames - def __repr__(self) -> str: if self._ptr_desc["type"] != "gpu": ptr_info = str(self._ptr_desc) @@ -510,89 +453,112 @@ def __repr__(self) -> str: ) -class SpillableBufferSlice(SpillableBuffer): +class SpillableBufferSlice(Buffer): """A slice of a spillable buffer This buffer applies the slicing and then delegates all - operations to its base buffer. + operations to its owning buffer. Parameters ---------- - base : SpillableBuffer - The base of the view + owner : SpillableBuffer + The owner of the view offset : int - Memory offset into the base buffer + Memory offset into the owning buffer size : int Size of the view (in bytes) """ - def __init__(self, base: SpillableBuffer, offset: int, size: int) -> None: + _owner: SpillableBuffer + + def __init__(self, owner: SpillableBuffer, offset: int, size: int) -> None: if size < 0: raise ValueError("size cannot be negative") if offset < 0: raise ValueError("offset cannot be negative") - if offset + size > base.size: + if offset + size > owner.size: raise ValueError( - "offset+size cannot be greater than the size of base" + "offset+size cannot be greater than the size of the owner" ) - self._base = base + self._owner = owner self._offset = offset self._size = size - self._owner = base - self.lock = base.lock - - def get_ptr(self, *, mode: Literal["read", "write"]) -> int: - """ - A passthrough method to `SpillableBuffer.get_ptr` - with factoring in the `offset`. - """ - return self._base.get_ptr(mode=mode) + self._offset - - def _getitem(self, offset: int, size: int) -> SpillableBufferSlice: - return SpillableBufferSlice( - base=self._base, offset=offset + self._offset, size=size - ) - - @classmethod - def deserialize(cls, header: dict, frames: list): - # TODO: because of the hack in `SpillableBuffer.serialize()` where - # frames are of type `Buffer`, we always deserialize as if they are - # `SpillableBuffer`. In the future, we should be able to - # deserialize into `SpillableBufferSlice` when the frames hasn't been - # copied. - return SpillableBuffer.deserialize(header, frames) - - def memoryview( - self, *, offset: int = 0, size: Optional[int] = None - ) -> memoryview: - size = self._size if size is None else size - return self._base.memoryview(offset=self._offset + offset, size=size) - - def __repr__(self) -> str: - return ( - f" None: - return self._base.spill(target=target) + return self._owner.spill(target=target) @property def is_spilled(self) -> bool: - return self._base.is_spilled + return self._owner.is_spilled @property def exposed(self) -> bool: - return self._base.exposed + return self._owner.exposed @property def spillable(self) -> bool: - return self._base.spillable + return self._owner.spillable def spill_lock(self, spill_lock: SpillLock) -> None: - self._base.spill_lock(spill_lock=spill_lock) + self._owner.spill_lock(spill_lock=spill_lock) def memory_info(self) -> Tuple[int, int, str]: - (ptr, _, device_type) = self._base.memory_info() + (ptr, _, device_type) = self._owner.memory_info() return (ptr + self._offset, self.nbytes, device_type) + + def mark_exposed(self) -> None: + self._owner.mark_exposed() + + def serialize(self) -> Tuple[dict, list]: + """Serialize the Buffer + + Normally, we would use `[self]` as the frames. This would work but + also mean that `self` becomes exposed permanently if the frames are + later accessed through `__cuda_array_interface__`, which is exactly + what libraries like Dask+UCX would do when communicating! + + The sound solution is to modify Dask et al. so that they access the + frames through `.get_ptr()` and holds on to the `spill_lock` until + the frame has been transferred. However, until this adaptation we + use a hack where the frame is a `Buffer` with a `spill_lock` as the + owner, which makes `self` unspillable while the frame is alive but + doesn't expose `self` when `__cuda_array_interface__` is accessed. + + Warning, this hack means that the returned frame must be copied before + given to `.deserialize()`, otherwise we would have a `Buffer` pointing + to memory already owned by an existing `SpillableBuffer`. + """ + header: Dict[str, Any] = {} + frames: List[BufferOwner | memoryview] + with self.lock: + header["type-serialized"] = pickle.dumps(self.__class__) + header["owner-type-serialized"] = pickle.dumps(type(self._owner)) + header["frame_count"] = 1 + if self.is_spilled: + frames = [self.memoryview()] + else: + # TODO: Use `frames=[self]` instead of this hack, see doc above + spill_lock = SpillLock() + self.spill_lock(spill_lock) + ptr, size, _ = self.memory_info() + frames = [ + BufferOwner._from_device_memory( + cuda_array_interface_wrapper( + ptr=ptr, + size=size, + owner=(self._owner, spill_lock), + ) + ) + ] + return header, frames + + @property + def __cuda_array_interface__(self) -> dict: + return { + "data": DelayedPointerTuple(self), + "shape": (self.size,), + "strides": None, + "typestr": "|u1", + "version": 0, + } diff --git a/python/cudf/cudf/core/buffer/utils.py b/python/cudf/cudf/core/buffer/utils.py index 373be99ec96..8ed17178044 100644 --- a/python/cudf/cudf/core/buffer/utils.py +++ b/python/cudf/cudf/core/buffer/utils.py @@ -6,7 +6,11 @@ from contextlib import ContextDecorator from typing import Any, Dict, Optional, Tuple, Union -from cudf.core.buffer.buffer import Buffer, cuda_array_interface_wrapper +from cudf.core.buffer.buffer import ( + Buffer, + BufferOwner, + cuda_array_interface_wrapper, +) from cudf.core.buffer.exposure_tracked_buffer import as_exposure_tracked_buffer from cudf.core.buffer.spill_manager import get_global_manager from cudf.core.buffer.spillable_buffer import SpillLock, as_spillable_buffer @@ -78,8 +82,14 @@ def as_buffer( if get_global_manager() is not None: return as_spillable_buffer(data, exposed=exposed) if hasattr(data, "__cuda_array_interface__"): - return Buffer._from_device_memory(data) - return Buffer._from_host_memory(data) + owner = BufferOwner._from_device_memory(data) + else: + owner = BufferOwner._from_host_memory(data) + return Buffer( + owner=owner, + offset=0, + size=owner.size, + ) _thread_spill_locks: Dict[int, Tuple[Optional[SpillLock], int]] = {} diff --git a/python/cudf/cudf/tests/test_copying.py b/python/cudf/cudf/tests/test_copying.py index 085774e9dbc..4151aa76753 100644 --- a/python/cudf/cudf/tests/test_copying.py +++ b/python/cudf/cudf/tests/test_copying.py @@ -113,11 +113,8 @@ def test_series_setitem_partial_slice_cow_on(): assert_eq(new_copy, cudf.Series([1, 2, 300, 300, 5])) new_slice = actual[2:] - # TODO: when COW and spilling has been unified, find a clean way to - # test this without accessing the internal attributes _base and _ptr assert ( - new_slice._column.base_data._base._ptr - == actual._column.base_data._base._ptr + new_slice._column.base_data.owner == actual._column.base_data.owner ) new_slice[0:2] = 10 assert_eq(new_slice, cudf.Series([10, 10, 5], index=[2, 3, 4])) @@ -134,9 +131,11 @@ def test_series_setitem_partial_slice_cow_off(): assert_eq(new_copy, cudf.Series([1, 2, 300, 300, 5])) new_slice = actual[2:] - assert ( - new_slice._column.base_data._ptr == actual._column.base_data._ptr - ) + # Since COW is off, a slice should point to the same memory + ptr1 = new_slice._column.base_data.get_ptr(mode="read") + ptr2 = actual._column.base_data.get_ptr(mode="read") + assert ptr1 == ptr2 + new_slice[0:2] = 10 assert_eq(new_slice, cudf.Series([10, 10, 5], index=[2, 3, 4])) assert_eq(actual, cudf.Series([1, 2, 10, 10, 5])) diff --git a/python/cudf/cudf/tests/test_spilling.py b/python/cudf/cudf/tests/test_spilling.py index 88ce908aa5f..55efa1e63a4 100644 --- a/python/cudf/cudf/tests/test_spilling.py +++ b/python/cudf/cudf/tests/test_spilling.py @@ -21,7 +21,7 @@ import cudf.options from cudf.core.abc import Serializable from cudf.core.buffer import ( - Buffer, + BufferOwner, acquire_spill_lock, as_buffer, get_spill_lock, @@ -71,17 +71,17 @@ def single_column_df(target="gpu") -> cudf.DataFrame: return ret -def single_column_df_data(df: cudf.DataFrame) -> SpillableBuffer: +def single_column_df_data(df: cudf.DataFrame) -> SpillableBufferSlice: """Access `.data` of the column of a standard dataframe""" ret = df._data._data["a"].data - assert isinstance(ret, SpillableBuffer) + assert isinstance(ret, SpillableBufferSlice) return ret -def single_column_df_base_data(df: cudf.DataFrame) -> SpillableBuffer: +def single_column_df_base_data(df: cudf.DataFrame) -> SpillableBufferSlice: """Access `.base_data` of the column of a standard dataframe""" ret = df._data._data["a"].base_data - assert isinstance(ret, SpillableBuffer) + assert isinstance(ret, SpillableBufferSlice) return ret @@ -117,7 +117,7 @@ def manager(request): def test_spillable_buffer(manager: SpillManager): buf = as_buffer(data=rmm.DeviceBuffer(size=10), exposed=False) - assert isinstance(buf, SpillableBuffer) + assert isinstance(buf, SpillableBufferSlice) assert buf.spillable buf.mark_exposed() assert buf.exposed @@ -196,10 +196,10 @@ def test_creations(manager: SpillManager): def test_spillable_df_groupby(manager: SpillManager): df = cudf.DataFrame({"a": [1, 1, 1]}) gb = df.groupby("a") - assert len(single_column_df_base_data(df)._spill_locks) == 0 + assert len(single_column_df_base_data(df).owner._spill_locks) == 0 gb._groupby # `gb._groupby`, which is cached on `gb`, holds a spill lock - assert len(single_column_df_base_data(df)._spill_locks) == 1 + assert len(single_column_df_base_data(df).owner._spill_locks) == 1 assert not single_column_df_data(df).spillable del gb assert single_column_df_data(df).spillable @@ -375,7 +375,7 @@ def test_get_ptr(manager: SpillManager, target): mem = np.empty(10, dtype="u1") buf = as_buffer(data=mem, exposed=False) assert buf.spillable - assert len(buf._spill_locks) == 0 + assert len(buf.owner._spill_locks) == 0 with acquire_spill_lock(): buf.get_ptr(mode="read") assert not buf.spillable @@ -439,7 +439,7 @@ def test_serialize_device(manager, target, view): header, frames = df1.device_serialize() assert len(frames) == 1 if target == "gpu": - assert isinstance(frames[0], Buffer) + assert isinstance(frames[0], BufferOwner) assert not single_column_df_data(df1).is_spilled assert not single_column_df_data(df1).spillable frames[0] = cupy.array(frames[0], copy=True) @@ -497,9 +497,9 @@ def test_serialize_cuda_dataframe(manager: SpillManager): df1, serializers=("cuda",), on_error="raise" ) buf: SpillableBufferSlice = single_column_df_data(df1) - assert len(buf._base._spill_locks) == 1 + assert len(buf.owner._spill_locks) == 1 assert len(frames) == 1 - assert isinstance(frames[0], Buffer) + assert isinstance(frames[0], BufferOwner) assert frames[0].get_ptr(mode="read") == buf.get_ptr(mode="read") frames[0] = cupy.array(frames[0], copy=True) @@ -542,8 +542,9 @@ def test_df_transpose(manager: SpillManager): def test_as_buffer_of_spillable_buffer(manager: SpillManager): data = cupy.arange(10, dtype="u1") b1 = as_buffer(data, exposed=False) - assert isinstance(b1, SpillableBuffer) - assert b1.owner is data + assert isinstance(b1, SpillableBufferSlice) + assert isinstance(b1.owner, SpillableBuffer) + assert b1.owner.owner is data b2 = as_buffer(b1) assert b1 is b2 @@ -558,7 +559,7 @@ def test_as_buffer_of_spillable_buffer(manager: SpillManager): with acquire_spill_lock(): b3 = as_buffer(b1.get_ptr(mode="read"), size=b1.size, owner=b1) assert isinstance(b3, SpillableBufferSlice) - assert b3.owner is b1 + assert b3.owner is b1.owner b4 = as_buffer( b1.get_ptr(mode="write") + data.itemsize, @@ -566,12 +567,12 @@ def test_as_buffer_of_spillable_buffer(manager: SpillManager): owner=b3, ) assert isinstance(b4, SpillableBufferSlice) - assert b4.owner is b1 + assert b4.owner is b1.owner assert all(cupy.array(b4.memoryview()) == data[1:]) b5 = as_buffer(b4.get_ptr(mode="write"), size=b4.size - 1, owner=b4) assert isinstance(b5, SpillableBufferSlice) - assert b5.owner is b1 + assert b5.owner is b1.owner assert all(cupy.array(b5.memoryview()) == data[1:-1]) @@ -593,7 +594,7 @@ def test_memoryview_slice(manager: SpillManager, dtype): def test_statistics(manager: SpillManager): assert len(manager.statistics.spill_totals) == 0 - buf: SpillableBuffer = as_buffer( + buf: SpillableBufferSlice = as_buffer( data=rmm.DeviceBuffer(size=10), exposed=False ) buf.spill(target="cpu") @@ -618,7 +619,7 @@ def test_statistics(manager: SpillManager): def test_statistics_expose(manager: SpillManager): assert len(manager.statistics.spill_totals) == 0 - buffers: List[SpillableBuffer] = [ + buffers: List[SpillableBufferSlice] = [ as_buffer(data=rmm.DeviceBuffer(size=10), exposed=False) for _ in range(10) ] @@ -644,7 +645,7 @@ def test_statistics_expose(manager: SpillManager): assert stat.spilled_nbytes == 0 # Create and spill 10 new buffers - buffers: List[SpillableBuffer] = [ + buffers: List[SpillableBufferSlice] = [ as_buffer(data=rmm.DeviceBuffer(size=10), exposed=False) for _ in range(10) ] From 6201b6ab86d2e52a89b91ce635a26d6b4ddd4fe9 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Wed, 2 Aug 2023 13:57:02 +0200 Subject: [PATCH 02/22] rename ExposureTrackedBuffer => ExposureTrackedBufferOwner --- python/cudf/cudf/core/abc.py | 2 +- python/cudf/cudf/core/buffer/buffer.py | 2 +- .../core/buffer/exposure_tracked_buffer.py | 62 +++++++++---------- 3 files changed, 32 insertions(+), 34 deletions(-) diff --git a/python/cudf/cudf/core/abc.py b/python/cudf/cudf/core/abc.py index a833663ce00..b47515da2b3 100644 --- a/python/cudf/cudf/core/abc.py +++ b/python/cudf/cudf/core/abc.py @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2022, NVIDIA CORPORATION. +# Copyright (c) 2020-2023, NVIDIA CORPORATION. """Common abstract base classes for cudf.""" import pickle diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index e6f285dd1a8..991cc023fe3 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -350,7 +350,7 @@ def copy(self, deep: bool = True) -> Self: Returns ------- - BufferSlice + ExposureTrackedBuffer A slice pointing to either a new or the existing owner depending on the expose status of the owner and the copy-on-write option (see above). diff --git a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py index e90a64fc906..f67622ab227 100644 --- a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py +++ b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py @@ -3,7 +3,7 @@ from __future__ import annotations import weakref -from typing import Any, Literal, Mapping, Optional, Type +from typing import Any, Literal, Mapping from typing_extensions import Self @@ -17,17 +17,11 @@ def as_exposure_tracked_buffer( - data, exposed: bool, subclass: Optional[Type[ExposureTrackedBuffer]] = None -) -> BufferSlice: + data, + exposed: bool, +) -> ExposureTrackedBuffer: """Factory function to wrap `data` in a slice of an exposure tracked buffer - If `subclass` is None, a new ExposureTrackedBuffer that points to the - memory of `data` is created and a BufferSlice that points to all of the - new ExposureTrackedBuffer is returned. - - If `subclass` is not None, a new `subclass` is created instead. Still, - a BufferSlice that points to all of the new `subclass` is returned - It is illegal for an exposure tracked buffer to own another exposure tracked buffer. When representing the same memory, we should have a single exposure tracked buffer and multiple buffer slices. @@ -35,7 +29,7 @@ def as_exposure_tracked_buffer( Developer Notes --------------- This function always returns slices thus all buffers in cudf will use - `BufferSlice` when copy-on-write is enabled. The slices implement + `ExposureTrackedBuffer` when copy-on-write is enabled. The slices implement copy-on-write by trigging deep copies when write access is detected and multiple slices points to the same exposure tracked buffer. @@ -45,23 +39,23 @@ def as_exposure_tracked_buffer( A buffer-like or array-like object that represents C-contiguous memory. exposed Mark the buffer as permanently exposed. - subclass - If not None, a subclass of ExposureTrackedBuffer to wrap `data`. Return ------ - BufferSlice - A buffer slice that points to a ExposureTrackedBuffer (or `subclass`), + ExposureTrackedBuffer + A buffer slice that points to a ExposureTrackedBufferOwner, which in turn wraps `data`. """ if not hasattr(data, "__cuda_array_interface__"): if exposed: raise ValueError("cannot created exposed host memory") - tracked_buf = ExposureTrackedBuffer._from_host_memory(data) - return BufferSlice(owner=tracked_buf, offset=0, size=tracked_buf.size) + tracked_buf = ExposureTrackedBufferOwner._from_host_memory(data) + return ExposureTrackedBuffer( + owner=tracked_buf, offset=0, size=tracked_buf.size + ) - owner = get_owner(data, subclass or ExposureTrackedBuffer) + owner = get_owner(data, ExposureTrackedBufferOwner) if owner is not None: # `data` is owned by an exposure tracked buffer ptr, size = get_ptr_and_size(data.__cuda_array_interface__) @@ -70,21 +64,25 @@ def as_exposure_tracked_buffer( raise ValueError( "Cannot create a non-empty slice of a null buffer" ) - return BufferSlice(owner=owner, offset=ptr - base_ptr, size=size) + return ExposureTrackedBuffer( + owner=owner, offset=ptr - base_ptr, size=size + ) # `data` is new device memory - owner = ExposureTrackedBuffer._from_device_memory(data, exposed=exposed) - return BufferSlice(owner=owner, offset=0, size=owner.size) + owner = ExposureTrackedBufferOwner._from_device_memory( + data, exposed=exposed + ) + return ExposureTrackedBuffer(owner=owner, offset=0, size=owner.size) -class ExposureTrackedBuffer(BufferOwner): +class ExposureTrackedBufferOwner(BufferOwner): """A Buffer that tracks its "expose" status. In order to implement copy-on-write and spillable buffers, we need the ability to detect external access to the underlying memory. We say that the buffer has been exposed if the device pointer (integer or void*) has - been accessed outside of ExposureTrackedBuffer. In this case, we have no - control over knowing if the data is being modified by a third-party. + been accessed outside of ExposureTrackedBufferOwner. In this case, we have + no control over knowing if the data is being modified by a third-party. Attributes ---------- @@ -92,11 +90,11 @@ class ExposureTrackedBuffer(BufferOwner): The current exposure status of the buffer. Notice, once the exposure status becomes True, it should never change back. _slices - The set of BufferSlice instances that point to this buffer. + The set of ExposureTrackedBuffer instances that point to this buffer. """ _exposed: bool - _slices: weakref.WeakSet[BufferSlice] + _slices: weakref.WeakSet[ExposureTrackedBuffer] @property def exposed(self) -> bool: @@ -119,24 +117,24 @@ def __cuda_array_interface__(self) -> Mapping: return super().__cuda_array_interface__ -class BufferSlice(Buffer): - """A slice (aka. a view) of a exposure tracked buffer. +class ExposureTrackedBuffer(Buffer): + """An exposure tracked buffer. Parameters ---------- owner - The exposure tracked buffer this slice refers to. + The owning exposure tracked buffer this refers to. offset The offset relative to the start memory of owner (in bytes). size The size of the slice (in bytes) """ - _owner: ExposureTrackedBuffer + _owner: ExposureTrackedBufferOwner def __init__( self, - owner: ExposureTrackedBuffer, + owner: ExposureTrackedBufferOwner, offset: int, size: int, ) -> None: @@ -175,7 +173,7 @@ def copy(self, deep: bool = True) -> Self: Returns ------- - BufferSlice + ExposureTrackedBuffer A slice pointing to either a new or the existing owner depending on the expose status of the owner and the copy-on-write option (see above). From 06ab42ed9a81238e424a546b9fd563d810c6aab2 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Wed, 2 Aug 2023 14:10:35 +0200 Subject: [PATCH 03/22] rename SpillableBuffer => SpillableBufferOwner --- python/cudf/cudf/core/buffer/spill_manager.py | 12 ++--- .../cudf/cudf/core/buffer/spillable_buffer.py | 53 ++++++++++--------- python/cudf/cudf/tests/test_spilling.py | 30 +++++------ 3 files changed, 48 insertions(+), 47 deletions(-) diff --git a/python/cudf/cudf/core/buffer/spill_manager.py b/python/cudf/cudf/core/buffer/spill_manager.py index f056a0fd592..2fb0c646017 100644 --- a/python/cudf/cudf/core/buffer/spill_manager.py +++ b/python/cudf/cudf/core/buffer/spill_manager.py @@ -15,7 +15,7 @@ import rmm.mr -from cudf.core.buffer.spillable_buffer import SpillableBuffer +from cudf.core.buffer.spillable_buffer import SpillableBufferOwner from cudf.options import get_option from cudf.utils.string import format_bytes @@ -122,7 +122,7 @@ def log_spill(self, src: str, dst: str, nbytes: int, time: float) -> None: total_time + time, ) - def log_expose(self, buf: SpillableBuffer) -> None: + def log_expose(self, buf: SpillableBufferOwner) -> None: """Log an expose event We track logged exposes by grouping them by their traceback such @@ -218,7 +218,7 @@ class SpillManager: SpillStatistics for the different levels. """ - _buffers: weakref.WeakValueDictionary[int, SpillableBuffer] + _buffers: weakref.WeakValueDictionary[int, SpillableBufferOwner] statistics: SpillStatistics def __init__( @@ -292,14 +292,14 @@ def _out_of_memory_handle(self, nbytes: int, *, retry_once=True) -> bool: ) return False # Since we didn't find anything to spill, we give up - def add(self, buffer: SpillableBuffer) -> None: + def add(self, buffer: SpillableBufferOwner) -> None: """Add buffer to the set of managed buffers The manager keeps a weak reference to the buffer Parameters ---------- - buffer : SpillableBuffer + buffer : SpillableBufferOwner The buffer to manage """ if buffer.size > 0 and not buffer.exposed: @@ -310,7 +310,7 @@ def add(self, buffer: SpillableBuffer) -> None: def buffers( self, order_by_access_time: bool = False - ) -> Tuple[SpillableBuffer, ...]: + ) -> Tuple[SpillableBufferOwner, ...]: """Get all managed buffers Parameters diff --git a/python/cudf/cudf/core/buffer/spillable_buffer.py b/python/cudf/cudf/core/buffer/spillable_buffer.py index 87c132ef238..c6df067956a 100644 --- a/python/cudf/cudf/core/buffer/spillable_buffer.py +++ b/python/cudf/cudf/core/buffer/spillable_buffer.py @@ -28,8 +28,8 @@ from cudf.core.buffer.spill_manager import SpillManager -def as_spillable_buffer(data, exposed: bool) -> SpillableBufferSlice: - """Factory function to wrap `data` in a SpillableBufferSlice object. +def as_spillable_buffer(data, exposed: bool) -> SpillableBuffer: + """Factory function to wrap `data` in a SpillableBuffer object. If `data` isn't a buffer already, a new buffer that points to the memory of `data` is created. If `data` represents host memory, it is copied to a new @@ -54,7 +54,7 @@ def as_spillable_buffer(data, exposed: bool) -> SpillableBufferSlice: Return ------ - SpillableBufferSlice + SpillableBuffer A spillabe buffer instance that represents the device memory of `data`. """ @@ -63,20 +63,20 @@ def as_spillable_buffer(data, exposed: bool) -> SpillableBufferSlice: if not hasattr(data, "__cuda_array_interface__"): if exposed: raise ValueError("cannot created exposed host memory") - tracked_buf = SpillableBuffer._from_host_memory(data) - return SpillableBufferSlice( + tracked_buf = SpillableBufferOwner._from_host_memory(data) + return SpillableBuffer( owner=tracked_buf, offset=0, size=tracked_buf.size ) if not hasattr(data, "__cuda_array_interface__"): if exposed: raise ValueError("cannot created exposed host memory") - spillabe_buf = SpillableBuffer._from_host_memory(data) - return SpillableBufferSlice( + spillabe_buf = SpillableBufferOwner._from_host_memory(data) + return SpillableBuffer( owner=spillabe_buf, offset=0, size=spillabe_buf.size ) - owner = get_owner(data, SpillableBuffer) + owner = get_owner(data, SpillableBufferOwner) if owner is not None: if not owner.exposed and get_spill_lock() is None: raise ValueError( @@ -91,13 +91,11 @@ def as_spillable_buffer(data, exposed: bool) -> SpillableBufferSlice: raise ValueError( "Cannot create a non-empty slice of a null buffer" ) - return SpillableBufferSlice( - owner=owner, offset=ptr - base_ptr, size=size - ) + return SpillableBuffer(owner=owner, offset=ptr - base_ptr, size=size) # `data` is new device memory - owner = SpillableBuffer._from_device_memory(data, exposed=exposed) - return SpillableBufferSlice(owner=owner, offset=0, size=owner.size) + owner = SpillableBufferOwner._from_device_memory(data, exposed=exposed) + return SpillableBuffer(owner=owner, offset=0, size=owner.size) class SpillLock: @@ -134,7 +132,7 @@ def __getitem__(self, i): raise IndexError("tuple index out of range") -class SpillableBuffer(BufferOwner): +class SpillableBufferOwner(BufferOwner): """A Buffer that supports spilling memory off the GPU to avoid OOMs. This buffer supports spilling the represented data to host memory. @@ -143,9 +141,9 @@ class SpillableBuffer(BufferOwner): device memory usage see `cudf.core.buffer.spill_manager.SpillManager`. Unspill is triggered automatically when accessing the data of the buffer. - The buffer might not be spillable, which is based on the "expose" status - of the buffer. We say that the buffer has been exposed if the device - pointer (integer or void*) has been accessed outside of SpillableBuffer. + The buffer might not be spillable, which is based on the "expose" status of + the buffer. We say that the buffer has been exposed if the device pointer + (integer or void*) has been accessed outside of SpillableBufferOwner. In this case, we cannot invalidate the device pointer by moving the data to host. @@ -153,7 +151,8 @@ class SpillableBuffer(BufferOwner): property. To avoid this, one can use `.get_ptr()` instead, which support exposing the buffer temporarily. - Use the factory function `as_buffer` to create a SpillableBuffer instance. + Use the factory function `as_buffer` to create a SpillableBufferOwner + instance. """ lock: RLock @@ -209,7 +208,7 @@ def _from_device_memory(cls, data: Any, *, exposed: bool = False) -> Self: Returns ------- - SpillableBuffer + SpillableBufferOwner Buffer representing the same device memory as `data` """ ret = super()._from_device_memory(data) @@ -234,7 +233,7 @@ def _from_host_memory(cls, data: Any) -> Self: Returns ------- - SpillableBuffer + SpillableBufferOwner Buffer representing a copy of `data`. """ @@ -446,14 +445,14 @@ def __repr__(self) -> str: else: ptr_info = str(hex(self._ptr)) return ( - f"" ) -class SpillableBufferSlice(Buffer): +class SpillableBuffer(Buffer): """A slice of a spillable buffer This buffer applies the slicing and then delegates all @@ -461,7 +460,7 @@ class SpillableBufferSlice(Buffer): Parameters ---------- - owner : SpillableBuffer + owner : SpillableBufferOwner The owner of the view offset : int Memory offset into the owning buffer @@ -469,9 +468,11 @@ class SpillableBufferSlice(Buffer): Size of the view (in bytes) """ - _owner: SpillableBuffer + _owner: SpillableBufferOwner - def __init__(self, owner: SpillableBuffer, offset: int, size: int) -> None: + def __init__( + self, owner: SpillableBufferOwner, offset: int, size: int + ) -> None: if size < 0: raise ValueError("size cannot be negative") if offset < 0: @@ -527,7 +528,7 @@ def serialize(self) -> Tuple[dict, list]: Warning, this hack means that the returned frame must be copied before given to `.deserialize()`, otherwise we would have a `Buffer` pointing - to memory already owned by an existing `SpillableBuffer`. + to memory already owned by an existing `SpillableBufferOwner`. """ header: Dict[str, Any] = {} frames: List[BufferOwner | memoryview] diff --git a/python/cudf/cudf/tests/test_spilling.py b/python/cudf/cudf/tests/test_spilling.py index 55efa1e63a4..39a8d4c8b72 100644 --- a/python/cudf/cudf/tests/test_spilling.py +++ b/python/cudf/cudf/tests/test_spilling.py @@ -34,7 +34,7 @@ ) from cudf.core.buffer.spillable_buffer import ( SpillableBuffer, - SpillableBufferSlice, + SpillableBufferOwner, SpillLock, ) from cudf.testing._utils import assert_eq @@ -71,17 +71,17 @@ def single_column_df(target="gpu") -> cudf.DataFrame: return ret -def single_column_df_data(df: cudf.DataFrame) -> SpillableBufferSlice: +def single_column_df_data(df: cudf.DataFrame) -> SpillableBuffer: """Access `.data` of the column of a standard dataframe""" ret = df._data._data["a"].data - assert isinstance(ret, SpillableBufferSlice) + assert isinstance(ret, SpillableBuffer) return ret -def single_column_df_base_data(df: cudf.DataFrame) -> SpillableBufferSlice: +def single_column_df_base_data(df: cudf.DataFrame) -> SpillableBuffer: """Access `.base_data` of the column of a standard dataframe""" ret = df._data._data["a"].base_data - assert isinstance(ret, SpillableBufferSlice) + assert isinstance(ret, SpillableBuffer) return ret @@ -117,7 +117,7 @@ def manager(request): def test_spillable_buffer(manager: SpillManager): buf = as_buffer(data=rmm.DeviceBuffer(size=10), exposed=False) - assert isinstance(buf, SpillableBufferSlice) + assert isinstance(buf, SpillableBuffer) assert buf.spillable buf.mark_exposed() assert buf.exposed @@ -496,7 +496,7 @@ def test_serialize_cuda_dataframe(manager: SpillManager): header, frames = protocol.serialize( df1, serializers=("cuda",), on_error="raise" ) - buf: SpillableBufferSlice = single_column_df_data(df1) + buf: SpillableBuffer = single_column_df_data(df1) assert len(buf.owner._spill_locks) == 1 assert len(frames) == 1 assert isinstance(frames[0], BufferOwner) @@ -542,8 +542,8 @@ def test_df_transpose(manager: SpillManager): def test_as_buffer_of_spillable_buffer(manager: SpillManager): data = cupy.arange(10, dtype="u1") b1 = as_buffer(data, exposed=False) - assert isinstance(b1, SpillableBufferSlice) - assert isinstance(b1.owner, SpillableBuffer) + assert isinstance(b1, SpillableBuffer) + assert isinstance(b1.owner, SpillableBufferOwner) assert b1.owner.owner is data b2 = as_buffer(b1) assert b1 is b2 @@ -558,7 +558,7 @@ def test_as_buffer_of_spillable_buffer(manager: SpillManager): with acquire_spill_lock(): b3 = as_buffer(b1.get_ptr(mode="read"), size=b1.size, owner=b1) - assert isinstance(b3, SpillableBufferSlice) + assert isinstance(b3, SpillableBuffer) assert b3.owner is b1.owner b4 = as_buffer( @@ -566,12 +566,12 @@ def test_as_buffer_of_spillable_buffer(manager: SpillManager): size=b1.size - data.itemsize, owner=b3, ) - assert isinstance(b4, SpillableBufferSlice) + assert isinstance(b4, SpillableBuffer) assert b4.owner is b1.owner assert all(cupy.array(b4.memoryview()) == data[1:]) b5 = as_buffer(b4.get_ptr(mode="write"), size=b4.size - 1, owner=b4) - assert isinstance(b5, SpillableBufferSlice) + assert isinstance(b5, SpillableBuffer) assert b5.owner is b1.owner assert all(cupy.array(b5.memoryview()) == data[1:-1]) @@ -594,7 +594,7 @@ def test_memoryview_slice(manager: SpillManager, dtype): def test_statistics(manager: SpillManager): assert len(manager.statistics.spill_totals) == 0 - buf: SpillableBufferSlice = as_buffer( + buf: SpillableBuffer = as_buffer( data=rmm.DeviceBuffer(size=10), exposed=False ) buf.spill(target="cpu") @@ -619,7 +619,7 @@ def test_statistics(manager: SpillManager): def test_statistics_expose(manager: SpillManager): assert len(manager.statistics.spill_totals) == 0 - buffers: List[SpillableBufferSlice] = [ + buffers: List[SpillableBuffer] = [ as_buffer(data=rmm.DeviceBuffer(size=10), exposed=False) for _ in range(10) ] @@ -645,7 +645,7 @@ def test_statistics_expose(manager: SpillManager): assert stat.spilled_nbytes == 0 # Create and spill 10 new buffers - buffers: List[SpillableBufferSlice] = [ + buffers: List[SpillableBuffer] = [ as_buffer(data=rmm.DeviceBuffer(size=10), exposed=False) for _ in range(10) ] From fa076ef2ab4098409821f2492ab85947cc00100e Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Wed, 2 Aug 2023 16:24:25 +0200 Subject: [PATCH 04/22] fix test_buffer_creation_from_any --- python/cudf/cudf/tests/test_buffer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf/cudf/tests/test_buffer.py b/python/cudf/cudf/tests/test_buffer.py index 1c9e7475080..2d561af5a95 100644 --- a/python/cudf/cudf/tests/test_buffer.py +++ b/python/cudf/cudf/tests/test_buffer.py @@ -64,7 +64,7 @@ def test_buffer_creation_from_any(): assert isinstance(b, Buffer) assert ary.data.ptr == b.get_ptr(mode="read") assert ary.nbytes == b.size - assert b.owner.owner is ary + assert b.owner.owner.owner is ary @pytest.mark.parametrize( From b37425ce6061044a20152cf0cd1a539acbcc8415 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 3 Aug 2023 08:55:33 +0200 Subject: [PATCH 05/22] doc --- python/cudf/cudf/core/buffer/buffer.py | 58 +++++++++++++------------- 1 file changed, 29 insertions(+), 29 deletions(-) diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index 991cc023fe3..770645051ae 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -127,7 +127,17 @@ def cuda_array_interface_wrapper( class BufferOwner(Serializable): - """A owning buffer represents device memory.""" + """A owning buffer that represents device memory. + + This class isn't meant to be used throughout cuDF. Instead, it + standardizes data owning by wrapping any data object that + represents device memory. Multiple `Buffer` instances, which are + the ones used throughout cuDF, can then refer to the same + `BufferOwner` instance. + + Use `_from_device_memory` and `_from_host_memory` to create + a new instance from either device or host memory respectively. + """ _ptr: int _size: int @@ -135,7 +145,7 @@ class BufferOwner(Serializable): @classmethod def _from_device_memory(cls, data: Any) -> Self: - """Create a Buffer from an object exposing `__cuda_array_interface__`. + """Create from an object exposing `__cuda_array_interface__`. No data is being copied. @@ -146,8 +156,8 @@ def _from_device_memory(cls, data: Any) -> Self: Returns ------- - Buffer - Buffer representing the same device memory as `data` + BufferOwner + BufferOwner wrapping `data` """ # Bypass `__init__` and initialize attributes manually @@ -166,7 +176,7 @@ def _from_device_memory(cls, data: Any) -> Self: @classmethod def _from_host_memory(cls, data: Any) -> Self: - """Create a Buffer from a buffer or array like object + """Create an owner from a buffer or array like object Data must implement `__array_interface__`, the buffer protocol, and/or be convertible to a buffer object using `numpy.array()` @@ -182,8 +192,8 @@ def _from_host_memory(cls, data: Any) -> Self: Returns ------- - Buffer - Buffer representing a copy of `data`. + BufferOwner + BufferOwner wrapping a device copy of `data`. """ # Convert to numpy array, this will not copy data in most cases. @@ -265,7 +275,7 @@ def __repr__(self) -> str: class Buffer(Serializable): - """A Buffer represents memory. + """A buffer that represents a slice or view of a `BufferOwner`. Use the factory function `as_buffer` to create a Buffer instance. """ @@ -330,30 +340,19 @@ def memoryview( def copy(self, deep: bool = True) -> Self: """Return a copy of Buffer. - What actually happens when `deep == False` depends on the - "copy_on_write" option. When copy-on-write is enabled, a shallow copy - becomes a deep copy if the buffer has been exposed. This is because we - have no control over knowing if the data is being modified when the - buffer has been exposed to third-party. - Parameters ---------- deep : bool, default True - The semantics when copy-on-write is disabled: - - If deep=True, returns a deep copy of the underlying data. - - If deep=False, returns a shallow copy of the Buffer pointing - to the same underlying data. - The semantics when copy-on-write is enabled: - - From the users perspective, always a deep copy of the - underlying data. However, the data isn't actually copied - until someone writers to the returned buffer. + - If deep=True, returns a deep copy of the underlying data. + - If deep=False, returns a new `Buffer` instance that refers + to the same `BufferOwner` as this one. Thus, no device + data are being copied. Returns ------- - ExposureTrackedBuffer - A slice pointing to either a new or the existing owner - depending on the expose status of the owner and the - copy-on-write option (see above). + Buffer + A new buffer that either refers to either a new or an existing + `BufferOwner` depending on the `deep` argument (see above). """ # When doing a shallow copy, we just return a new slice @@ -362,7 +361,7 @@ def copy(self, deep: bool = True) -> Self: owner=self._owner, offset=self._offset, size=self._size ) - # Otherwise, we create a new copy of the memory this slice represents + # Otherwise, we create a new copy of the memory owner = self._owner._from_device_memory( rmm.DeviceBuffer( ptr=self._owner.get_ptr(mode="read") + self._offset, @@ -385,14 +384,15 @@ def __cuda_array_interface__(self) -> Mapping: def serialize(self) -> Tuple[dict, list]: """Serialize the buffer into header and frames. - The frames can be a mixture of memoryview and Buffer objects. + The frames can be a mixture of memoryview, Buffer, and BufferOwner + objects. Returns ------- Tuple[dict, List] The first element of the returned tuple is a dict containing any serializable metadata required to reconstruct the object. The - second element is a list containing Buffers and memoryviews. + second element is a list containing single frame. """ header: Dict[str, Any] = {} header["type-serialized"] = pickle.dumps(type(self)) From d6fc8b47c25f9f0f09b0404cbc967719960b5b9b Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 3 Aug 2023 09:49:01 +0200 Subject: [PATCH 06/22] clean up inits --- python/cudf/cudf/core/buffer/buffer.py | 14 ++++++-- .../core/buffer/exposure_tracked_buffer.py | 16 ++++----- .../cudf/cudf/core/buffer/spillable_buffer.py | 34 +++---------------- python/cudf/cudf/core/buffer/utils.py | 9 ++--- 4 files changed, 27 insertions(+), 46 deletions(-) diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index 770645051ae..a2b29a00c95 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -278,15 +278,25 @@ class Buffer(Serializable): """A buffer that represents a slice or view of a `BufferOwner`. Use the factory function `as_buffer` to create a Buffer instance. + + Parameters + ---------- + owner + The owning exposure buffer this refers to. + offset + The offset relative to the start memory of owner (in bytes). + size + The size of the buffer (in bytes). If None, use the size of owner. """ def __init__( self, *, owner: BufferOwner, - offset: int, - size: int, + offset: int = 0, + size: Optional[int] = None, ) -> None: + size = owner.size if size is None else size if size < 0: raise ValueError("size cannot be negative") if offset < 0: diff --git a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py index f67622ab227..c8a93607d7d 100644 --- a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py +++ b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py @@ -3,7 +3,7 @@ from __future__ import annotations import weakref -from typing import Any, Literal, Mapping +from typing import Any, Literal, Mapping, Optional from typing_extensions import Self @@ -50,9 +50,8 @@ def as_exposure_tracked_buffer( if not hasattr(data, "__cuda_array_interface__"): if exposed: raise ValueError("cannot created exposed host memory") - tracked_buf = ExposureTrackedBufferOwner._from_host_memory(data) return ExposureTrackedBuffer( - owner=tracked_buf, offset=0, size=tracked_buf.size + owner=ExposureTrackedBufferOwner._from_host_memory(data) ) owner = get_owner(data, ExposureTrackedBufferOwner) @@ -69,10 +68,11 @@ def as_exposure_tracked_buffer( ) # `data` is new device memory - owner = ExposureTrackedBufferOwner._from_device_memory( - data, exposed=exposed + return ExposureTrackedBuffer( + owner=ExposureTrackedBufferOwner._from_device_memory( + data, exposed=exposed + ) ) - return ExposureTrackedBuffer(owner=owner, offset=0, size=owner.size) class ExposureTrackedBufferOwner(BufferOwner): @@ -135,8 +135,8 @@ class ExposureTrackedBuffer(Buffer): def __init__( self, owner: ExposureTrackedBufferOwner, - offset: int, - size: int, + offset: int = 0, + size: Optional[int] = None, ) -> None: super().__init__(owner=owner, offset=offset, size=size) self._owner._slices.add(self) diff --git a/python/cudf/cudf/core/buffer/spillable_buffer.py b/python/cudf/cudf/core/buffer/spillable_buffer.py index c6df067956a..08f8fefbf89 100644 --- a/python/cudf/cudf/core/buffer/spillable_buffer.py +++ b/python/cudf/cudf/core/buffer/spillable_buffer.py @@ -63,17 +63,8 @@ def as_spillable_buffer(data, exposed: bool) -> SpillableBuffer: if not hasattr(data, "__cuda_array_interface__"): if exposed: raise ValueError("cannot created exposed host memory") - tracked_buf = SpillableBufferOwner._from_host_memory(data) return SpillableBuffer( - owner=tracked_buf, offset=0, size=tracked_buf.size - ) - - if not hasattr(data, "__cuda_array_interface__"): - if exposed: - raise ValueError("cannot created exposed host memory") - spillabe_buf = SpillableBufferOwner._from_host_memory(data) - return SpillableBuffer( - owner=spillabe_buf, offset=0, size=spillabe_buf.size + owner=SpillableBufferOwner._from_host_memory(data) ) owner = get_owner(data, SpillableBufferOwner) @@ -94,8 +85,9 @@ def as_spillable_buffer(data, exposed: bool) -> SpillableBuffer: return SpillableBuffer(owner=owner, offset=ptr - base_ptr, size=size) # `data` is new device memory - owner = SpillableBufferOwner._from_device_memory(data, exposed=exposed) - return SpillableBuffer(owner=owner, offset=0, size=owner.size) + return SpillableBuffer( + owner=SpillableBufferOwner._from_device_memory(data, exposed=exposed) + ) class SpillLock: @@ -470,22 +462,6 @@ class SpillableBuffer(Buffer): _owner: SpillableBufferOwner - def __init__( - self, owner: SpillableBufferOwner, offset: int, size: int - ) -> None: - if size < 0: - raise ValueError("size cannot be negative") - if offset < 0: - raise ValueError("offset cannot be negative") - if offset + size > owner.size: - raise ValueError( - "offset+size cannot be greater than the size of the owner" - ) - self._owner = owner - self._offset = offset - self._size = size - self.lock = owner.lock - def spill(self, target: str = "cpu") -> None: return self._owner.spill(target=target) @@ -532,7 +508,7 @@ def serialize(self) -> Tuple[dict, list]: """ header: Dict[str, Any] = {} frames: List[BufferOwner | memoryview] - with self.lock: + with self._owner.lock: header["type-serialized"] = pickle.dumps(self.__class__) header["owner-type-serialized"] = pickle.dumps(type(self._owner)) header["frame_count"] = 1 diff --git a/python/cudf/cudf/core/buffer/utils.py b/python/cudf/cudf/core/buffer/utils.py index 8ed17178044..e3ab1b7d3b0 100644 --- a/python/cudf/cudf/core/buffer/utils.py +++ b/python/cudf/cudf/core/buffer/utils.py @@ -82,14 +82,9 @@ def as_buffer( if get_global_manager() is not None: return as_spillable_buffer(data, exposed=exposed) if hasattr(data, "__cuda_array_interface__"): - owner = BufferOwner._from_device_memory(data) + return Buffer(owner=BufferOwner._from_device_memory(data)) else: - owner = BufferOwner._from_host_memory(data) - return Buffer( - owner=owner, - offset=0, - size=owner.size, - ) + return Buffer(owner=BufferOwner._from_host_memory(data)) _thread_spill_locks: Dict[int, Tuple[Optional[SpillLock], int]] = {} From 3b499708aa5ff6699d0b1f43e961e0676b603cfd Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 3 Aug 2023 10:33:04 +0200 Subject: [PATCH 07/22] unify as_spillable_buffer, as_exposure_tracked_buffer, and as_buffer --- .../source/developer_guide/library_design.md | 14 ++-- python/cudf/cudf/core/buffer/buffer.py | 39 ++++------ .../core/buffer/exposure_tracked_buffer.py | 70 +---------------- .../cudf/cudf/core/buffer/spillable_buffer.py | 73 ++---------------- python/cudf/cudf/core/buffer/utils.py | 75 ++++++++++++++++--- 5 files changed, 96 insertions(+), 175 deletions(-) diff --git a/docs/cudf/source/developer_guide/library_design.md b/docs/cudf/source/developer_guide/library_design.md index 016c2c1d281..1c85e71128c 100644 --- a/docs/cudf/source/developer_guide/library_design.md +++ b/docs/cudf/source/developer_guide/library_design.md @@ -325,26 +325,26 @@ This section describes the internal implementation details of the copy-on-write It is recommended that developers familiarize themselves with [the user-facing documentation](copy-on-write-user-doc) of this functionality before reading through the internals below. -The core copy-on-write implementation relies on the factory function `as_exposure_tracked_buffer` and the two classes `ExposureTrackedBuffer` and `BufferSlice`. +The core copy-on-write implementation relies on the two classes `ExposureTrackedBufferOwner` and `ExposureTrackedBuffer`. -An `ExposureTrackedBuffer` is a subclass of the regular `Buffer` that tracks internal and external references to its underlying memory. Internal references are tracked by maintaining [weak references](https://docs.python.org/3/library/weakref.html) to every `BufferSlice` of the underlying memory. External references are tracked through "exposure" status of the underlying memory. A buffer is considered exposed if the device pointer (integer or void*) has been handed out to a library outside of cudf. In this case, we have no way of knowing if the data are being modified by a third party. +An `ExposureTrackedBufferOwner` is a subclass of the `BufferOwner` that tracks internal and external references to its underlying memory. Internal references are tracked by maintaining [weak references](https://docs.python.org/3/library/weakref.html) to every `ExposureTrackedBuffer` of the underlying memory. External references are tracked through "exposure" status of the underlying memory. A buffer is considered exposed if the device pointer (integer or void*) has been handed out to a library outside of cudf. In this case, we have no way of knowing if the data are being modified by a third party. -`BufferSlice` is a subclass of `ExposureTrackedBuffer` that represents a _slice_ of the memory underlying a exposure tracked buffer. +`ExposureTrackedBuffer` is a subclass of `Buffer` that represents a _slice_ of the memory underlying an exposure tracked buffer. -When the cudf option `"copy_on_write"` is `True`, `as_buffer` calls `as_exposure_tracked_buffer`, which always returns a `BufferSlice`. It is then the slices that determine whether or not to make a copy when a write operation is performed on a `Column` (see below). If multiple slices point to the same underlying memory, then a copy must be made whenever a modification is attempted. +When the cudf option `"copy_on_write"` is `True`, `as_buffer` returns a `ExposureTrackedBuffer`. It is this class that determine whether or not to make a copy when a write operation is performed on a `Column` (see below). If multiple slices point to the same underlying memory, then a copy must be made whenever a modification is attempted. ### Eager copies when exposing to third-party libraries -If a `Column`/`BufferSlice` is exposed to a third-party library via `__cuda_array_interface__`, we are no longer able to track whether or not modification of the buffer has occurred. Hence whenever +If a `Column`/`ExposureTrackedBuffer` is exposed to a third-party library via `__cuda_array_interface__`, we are no longer able to track whether or not modification of the buffer has occurred. Hence whenever someone accesses data through the `__cuda_array_interface__`, we eagerly trigger the copy by calling -`.make_single_owner_inplace` which ensures a true copy of underlying data is made and that the slice is the sole owner. Any future copy requests must also trigger a true physical copy (since we cannot track the lifetime of the third-party object). To handle this we also mark the `Column`/`BufferSlice` as exposed thus indicating that any future shallow-copy requests will trigger a true physical copy rather than a copy-on-write shallow copy. +`.make_single_owner_inplace` which ensures a true copy of underlying data is made and that the slice is the sole owner. Any future copy requests must also trigger a true physical copy (since we cannot track the lifetime of the third-party object). To handle this we also mark the `Column`/`ExposureTrackedBuffer` as exposed thus indicating that any future shallow-copy requests will trigger a true physical copy rather than a copy-on-write shallow copy. ### Obtaining a read-only object A read-only object can be quite useful for operations that will not mutate the data. This can be achieved by calling `.get_ptr(mode="read")`, and using `cuda_array_interface_wrapper` to wrap a `__cuda_array_interface__` object around it. -This will not trigger a deep copy even if multiple `BufferSlice` points to the same `ExposureTrackedBuffer`. This API should only be used when the lifetime of the proxy object is restricted to cudf's internal code execution. Handing this out to external libraries or user-facing APIs will lead to untracked references and undefined copy-on-write behavior. We currently use this API for device to host +This will not trigger a deep copy even if multiple `ExposureTrackedBuffer` points to the same `ExposureTrackedBufferOwner`. This API should only be used when the lifetime of the proxy object is restricted to cudf's internal code execution. Handing this out to external libraries or user-facing APIs will lead to untracked references and undefined copy-on-write behavior. We currently use this API for device to host copies like in `ColumnBase.data_array_view(mode="read")` which is used for `Column.values_host`. diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index a2b29a00c95..f651eb35bd6 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -5,17 +5,7 @@ import math import pickle from types import SimpleNamespace -from typing import ( - Any, - Dict, - Literal, - Mapping, - Optional, - Sequence, - Tuple, - Type, - TypeVar, -) +from typing import Any, Dict, Literal, Mapping, Optional, Sequence, Tuple import numpy from typing_extensions import Self @@ -26,14 +16,12 @@ from cudf.core.abc import Serializable from cudf.utils.string import format_bytes -T = TypeVar("T") - -def get_owner(data, klass: Type[T]) -> Optional[T]: +def get_buffer_owner(data) -> Optional[BufferOwner]: """Get the owner of `data`, if any exist Search through the stack of data owners in order to find an - owner of type `klass` (not subclasses). + owner BufferOwner (incl. subclasses). Parameters ---------- @@ -42,14 +30,14 @@ def get_owner(data, klass: Type[T]) -> Optional[T]: Return ------ - klass or None - The owner of `data` if `klass` or None. + BufferOwner or None + The owner of `data` if found otherwise None. """ - if type(data) is klass: + if isinstance(data, BufferOwner): return data if hasattr(data, "owner"): - return get_owner(data.owner, klass) + return get_buffer_owner(data.owner) return None @@ -144,7 +132,7 @@ class BufferOwner(Serializable): _owner: object @classmethod - def _from_device_memory(cls, data: Any) -> Self: + def _from_device_memory(cls, data: Any, exposed: bool) -> Self: """Create from an object exposing `__cuda_array_interface__`. No data is being copied. @@ -153,6 +141,10 @@ def _from_device_memory(cls, data: Any) -> Self: ---------- data : device-buffer-like An object implementing the CUDA Array Interface. + exposed : bool + Mark the buffer as permanently exposed. This is used by + ExposureTrackedBuffer to determine when a deep copy is required + and by SpillableBuffer to mark the buffer unspillable. Returns ------- @@ -203,7 +195,7 @@ def _from_host_memory(cls, data: Any) -> Self: # Copy to device memory buf = rmm.DeviceBuffer(ptr=ptr, size=size) # Create from device memory - return cls._from_device_memory(buf) + return cls._from_device_memory(buf, exposed=False) @property def size(self) -> int: @@ -376,7 +368,8 @@ def copy(self, deep: bool = True) -> Self: rmm.DeviceBuffer( ptr=self._owner.get_ptr(mode="read") + self._offset, size=self.size, - ) + ), + exposed=False, ) return self.__class__(owner=owner, offset=0, size=owner.size) @@ -435,7 +428,7 @@ def deserialize(cls, header: dict, frames: list) -> Self: owner_type: BufferOwner = pickle.loads(header["owner-type-serialized"]) if hasattr(frame, "__cuda_array_interface__"): - owner = owner_type._from_device_memory(frame) + owner = owner_type._from_device_memory(frame, exposed=False) else: owner = owner_type._from_host_memory(frame) return cls( diff --git a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py index c8a93607d7d..4f4a56470e8 100644 --- a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py +++ b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py @@ -8,71 +8,7 @@ from typing_extensions import Self import cudf -from cudf.core.buffer.buffer import ( - Buffer, - BufferOwner, - get_owner, - get_ptr_and_size, -) - - -def as_exposure_tracked_buffer( - data, - exposed: bool, -) -> ExposureTrackedBuffer: - """Factory function to wrap `data` in a slice of an exposure tracked buffer - - It is illegal for an exposure tracked buffer to own another exposure - tracked buffer. When representing the same memory, we should have a single - exposure tracked buffer and multiple buffer slices. - - Developer Notes - --------------- - This function always returns slices thus all buffers in cudf will use - `ExposureTrackedBuffer` when copy-on-write is enabled. The slices implement - copy-on-write by trigging deep copies when write access is detected - and multiple slices points to the same exposure tracked buffer. - - Parameters - ---------- - data : buffer-like or array-like - A buffer-like or array-like object that represents C-contiguous memory. - exposed - Mark the buffer as permanently exposed. - - Return - ------ - ExposureTrackedBuffer - A buffer slice that points to a ExposureTrackedBufferOwner, - which in turn wraps `data`. - """ - - if not hasattr(data, "__cuda_array_interface__"): - if exposed: - raise ValueError("cannot created exposed host memory") - return ExposureTrackedBuffer( - owner=ExposureTrackedBufferOwner._from_host_memory(data) - ) - - owner = get_owner(data, ExposureTrackedBufferOwner) - if owner is not None: - # `data` is owned by an exposure tracked buffer - ptr, size = get_ptr_and_size(data.__cuda_array_interface__) - base_ptr = owner.get_ptr(mode="read") - if size > 0 and base_ptr == 0: - raise ValueError( - "Cannot create a non-empty slice of a null buffer" - ) - return ExposureTrackedBuffer( - owner=owner, offset=ptr - base_ptr, size=size - ) - - # `data` is new device memory - return ExposureTrackedBuffer( - owner=ExposureTrackedBufferOwner._from_device_memory( - data, exposed=exposed - ) - ) +from cudf.core.buffer.buffer import Buffer, BufferOwner class ExposureTrackedBufferOwner(BufferOwner): @@ -105,8 +41,8 @@ def mark_exposed(self) -> None: self._exposed = True @classmethod - def _from_device_memory(cls, data: Any, *, exposed: bool = False) -> Self: - ret = super()._from_device_memory(data) + def _from_device_memory(cls, data: Any, exposed: bool) -> Self: + ret = super()._from_device_memory(data, exposed=exposed) ret._exposed = exposed ret._slices = weakref.WeakSet() return ret diff --git a/python/cudf/cudf/core/buffer/spillable_buffer.py b/python/cudf/cudf/core/buffer/spillable_buffer.py index 08f8fefbf89..ddee018d415 100644 --- a/python/cudf/cudf/core/buffer/spillable_buffer.py +++ b/python/cudf/cudf/core/buffer/spillable_buffer.py @@ -18,8 +18,6 @@ Buffer, BufferOwner, cuda_array_interface_wrapper, - get_owner, - get_ptr_and_size, host_memory_allocation, ) from cudf.utils.string import format_bytes @@ -28,68 +26,6 @@ from cudf.core.buffer.spill_manager import SpillManager -def as_spillable_buffer(data, exposed: bool) -> SpillableBuffer: - """Factory function to wrap `data` in a SpillableBuffer object. - - If `data` isn't a buffer already, a new buffer that points to the memory of - `data` is created. If `data` represents host memory, it is copied to a new - `rmm.DeviceBuffer` device allocation. Otherwise, the memory of `data` is - **not** copied, instead the new buffer keeps a reference to `data` in order - to retain its lifetime. - - If `data` is owned by a spillable buffer, a "slice" of the buffer is - returned. In this case, the spillable buffer must either be "exposed" or - spilled locked (called within an acquire_spill_lock context). This is to - guarantee that the memory of `data` isn't spilled before this function gets - to calculate the offset of the new slice. - - It is illegal for a spillable buffer to own another spillable buffer. - - Parameters - ---------- - data : buffer-like or array-like - A buffer-like or array-like object that represent C-contiguous memory. - exposed : bool, optional - Mark the buffer as permanently exposed (unspillable). - - Return - ------ - SpillableBuffer - A spillabe buffer instance that represents the device memory of `data`. - """ - - from cudf.core.buffer.utils import get_spill_lock - - if not hasattr(data, "__cuda_array_interface__"): - if exposed: - raise ValueError("cannot created exposed host memory") - return SpillableBuffer( - owner=SpillableBufferOwner._from_host_memory(data) - ) - - owner = get_owner(data, SpillableBufferOwner) - if owner is not None: - if not owner.exposed and get_spill_lock() is None: - raise ValueError( - "A owning spillable buffer must " - "either be exposed or spilled locked." - ) - # `data` is owned by an spillable buffer, which is exposed or - # spilled locked. - ptr, size = get_ptr_and_size(data.__cuda_array_interface__) - base_ptr = owner.get_ptr(mode="read") - if size > 0 and owner._ptr == 0: - raise ValueError( - "Cannot create a non-empty slice of a null buffer" - ) - return SpillableBuffer(owner=owner, offset=ptr - base_ptr, size=size) - - # `data` is new device memory - return SpillableBuffer( - owner=SpillableBufferOwner._from_device_memory(data, exposed=exposed) - ) - - class SpillLock: pass @@ -186,7 +122,7 @@ def _finalize_init(self, ptr_desc: Dict[str, Any], exposed: bool) -> None: self._manager.add(self) @classmethod - def _from_device_memory(cls, data: Any, *, exposed: bool = False) -> Self: + def _from_device_memory(cls, data: Any, exposed: bool) -> Self: """Create a spillabe buffer from device memory. No data is being copied. @@ -195,7 +131,7 @@ def _from_device_memory(cls, data: Any, *, exposed: bool = False) -> Self: ---------- data : device-buffer-like An object implementing the CUDA Array Interface. - exposed : bool, optional + exposed : bool Mark the buffer as permanently exposed (unspillable). Returns @@ -203,7 +139,7 @@ def _from_device_memory(cls, data: Any, *, exposed: bool = False) -> Self: SpillableBufferOwner Buffer representing the same device memory as `data` """ - ret = super()._from_device_memory(data) + ret = super()._from_device_memory(data, exposed=exposed) ret._finalize_init(ptr_desc={"type": "gpu"}, exposed=exposed) return ret @@ -525,7 +461,8 @@ def serialize(self) -> Tuple[dict, list]: ptr=ptr, size=size, owner=(self._owner, spill_lock), - ) + ), + exposed=False, ) ] return header, frames diff --git a/python/cudf/cudf/core/buffer/utils.py b/python/cudf/cudf/core/buffer/utils.py index e3ab1b7d3b0..3e2dc11e993 100644 --- a/python/cudf/cudf/core/buffer/utils.py +++ b/python/cudf/cudf/core/buffer/utils.py @@ -4,16 +4,25 @@ import threading from contextlib import ContextDecorator -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any, Dict, Optional, Tuple, Type, Union from cudf.core.buffer.buffer import ( Buffer, BufferOwner, cuda_array_interface_wrapper, + get_buffer_owner, + get_ptr_and_size, +) +from cudf.core.buffer.exposure_tracked_buffer import ( + ExposureTrackedBuffer, + ExposureTrackedBufferOwner, ) -from cudf.core.buffer.exposure_tracked_buffer import as_exposure_tracked_buffer from cudf.core.buffer.spill_manager import get_global_manager -from cudf.core.buffer.spillable_buffer import SpillLock, as_spillable_buffer +from cudf.core.buffer.spillable_buffer import ( + SpillableBuffer, + SpillableBufferOwner, + SpillLock, +) from cudf.options import get_option @@ -34,7 +43,18 @@ def as_buffer( If `data` is an integer, it is assumed to point to device memory. - Raises ValueError if data isn't C-contiguous. + Raises ValueError if `data` isn't C-contiguous. + + If copy-on-write is enabled, a ExposureTrackedBuffer that refers to a + ExposureTrackedBufferOwner is returned. + + If spilling is enabled, a SpillableBuffer that refers ot a + SpillableBufferOwner is returned. If `data` is owned by a spillable buffer, + it must either be "exposed" or spilled locked (called within an + acquire_spill_lock context). This is to guarantee that the memory of `data` + isn't spilled before this function gets to calculate the offset of the new + SpillableBuffer. + Parameters ---------- @@ -77,14 +97,49 @@ def as_buffer( "`data` is a buffer-like or array-like object" ) - if get_option("copy_on_write"): - return as_exposure_tracked_buffer(data, exposed=exposed) + # Find the buffer types to return based on the current config + owner_class: Type[BufferOwner] + buffer_class: Type[Buffer] if get_global_manager() is not None: - return as_spillable_buffer(data, exposed=exposed) - if hasattr(data, "__cuda_array_interface__"): - return Buffer(owner=BufferOwner._from_device_memory(data)) + owner_class = SpillableBufferOwner + buffer_class = SpillableBuffer + elif get_option("copy_on_write"): + owner_class = ExposureTrackedBufferOwner + buffer_class = ExposureTrackedBuffer else: - return Buffer(owner=BufferOwner._from_host_memory(data)) + owner_class = BufferOwner + buffer_class = Buffer + + # Handle host memory, + if not hasattr(data, "__cuda_array_interface__"): + if exposed: + raise ValueError("cannot created exposed host memory") + return buffer_class(owner=owner_class._from_host_memory(data)) + + # Check if `data` is owned by a known class + owner = get_buffer_owner(data) + if owner is None: # `data` is new device memory + return buffer_class( + owner=owner_class._from_device_memory(data, exposed=exposed) + ) + + # At this point, we know that `data` is owned by a known class, which + # should be the same class as specified by the current config (see above) + assert owner.__class__ is owner_class + if ( + isinstance(owner, SpillableBufferOwner) + and not owner.exposed + and get_spill_lock() is None + ): + raise ValueError( + "A owning spillable buffer must " + "either be exposed or spilled locked." + ) + ptr, size = get_ptr_and_size(data.__cuda_array_interface__) + base_ptr = owner.get_ptr(mode="read") + if size > 0 and base_ptr == 0: + raise ValueError("Cannot create a non-empty slice of a null buffer") + return buffer_class(owner=owner, offset=ptr - base_ptr, size=size) _thread_spill_locks: Dict[int, Tuple[Optional[SpillLock], int]] = {} From 31966705d86796c591ff01a5cffb4c2e0e5dff5b Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 23 Nov 2023 16:20:09 +0100 Subject: [PATCH 08/22] Apply suggestions from code review Co-authored-by: Lawrence Mitchell --- docs/cudf/source/developer_guide/library_design.md | 6 +++--- python/cudf/cudf/core/buffer/buffer.py | 12 ++++++++---- python/cudf/cudf/core/buffer/utils.py | 6 +++--- 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/docs/cudf/source/developer_guide/library_design.md b/docs/cudf/source/developer_guide/library_design.md index 1c85e71128c..39d56559add 100644 --- a/docs/cudf/source/developer_guide/library_design.md +++ b/docs/cudf/source/developer_guide/library_design.md @@ -327,11 +327,11 @@ below. The core copy-on-write implementation relies on the two classes `ExposureTrackedBufferOwner` and `ExposureTrackedBuffer`. -An `ExposureTrackedBufferOwner` is a subclass of the `BufferOwner` that tracks internal and external references to its underlying memory. Internal references are tracked by maintaining [weak references](https://docs.python.org/3/library/weakref.html) to every `ExposureTrackedBuffer` of the underlying memory. External references are tracked through "exposure" status of the underlying memory. A buffer is considered exposed if the device pointer (integer or void*) has been handed out to a library outside of cudf. In this case, we have no way of knowing if the data are being modified by a third party. +An `ExposureTrackedBufferOwner` is a subclass of the `BufferOwner` that tracks internal and external references to its underlying memory. Internal references are tracked by maintaining [weak references](https://docs.python.org/3/library/weakref.html) to every `ExposureTrackedBuffer` of the underlying memory. External references are tracked through "exposure" status of the underlying memory. A buffer is considered exposed if the device pointer (integer or void*) has been handed out to a library outside of cudf. In this case, we have no way of knowing if the data are being modified by a third party. `ExposureTrackedBuffer` is a subclass of `Buffer` that represents a _slice_ of the memory underlying an exposure tracked buffer. -When the cudf option `"copy_on_write"` is `True`, `as_buffer` returns a `ExposureTrackedBuffer`. It is this class that determine whether or not to make a copy when a write operation is performed on a `Column` (see below). If multiple slices point to the same underlying memory, then a copy must be made whenever a modification is attempted. +When the cudf option `"copy_on_write"` is `True`, `as_buffer` returns a `ExposureTrackedBuffer`. It is this class that determines whether or not to make a copy when a write operation is performed on a `Column` (see below). If multiple slices point to the same underlying memory, then a copy must be made whenever a modification is attempted. ### Eager copies when exposing to third-party libraries @@ -344,7 +344,7 @@ someone accesses data through the `__cuda_array_interface__`, we eagerly trigger A read-only object can be quite useful for operations that will not mutate the data. This can be achieved by calling `.get_ptr(mode="read")`, and using `cuda_array_interface_wrapper` to wrap a `__cuda_array_interface__` object around it. -This will not trigger a deep copy even if multiple `ExposureTrackedBuffer` points to the same `ExposureTrackedBufferOwner`. This API should only be used when the lifetime of the proxy object is restricted to cudf's internal code execution. Handing this out to external libraries or user-facing APIs will lead to untracked references and undefined copy-on-write behavior. We currently use this API for device to host +This will not trigger a deep copy even if multiple `ExposureTrackedBuffer`s point to the same `ExposureTrackedBufferOwner`. This API should only be used when the lifetime of the proxy object is restricted to cudf's internal code execution. Handing this out to external libraries or user-facing APIs will lead to untracked references and undefined copy-on-write behavior. We currently use this API for device to host copies like in `ColumnBase.data_array_view(mode="read")` which is used for `Column.values_host`. diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index f651eb35bd6..0a4e73d5cbc 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -18,7 +18,7 @@ def get_buffer_owner(data) -> Optional[BufferOwner]: - """Get the owner of `data`, if any exist + """Get the owner of `data`, if one exists Search through the stack of data owners in order to find an owner BufferOwner (incl. subclasses). @@ -115,7 +115,7 @@ def cuda_array_interface_wrapper( class BufferOwner(Serializable): - """A owning buffer that represents device memory. + """An owning buffer that represents device memory. This class isn't meant to be used throughout cuDF. Instead, it standardizes data owning by wrapping any data object that @@ -133,7 +133,7 @@ class BufferOwner(Serializable): @classmethod def _from_device_memory(cls, data: Any, exposed: bool) -> Self: - """Create from an object exposing `__cuda_array_interface__`. + """Create from an object providing a `__cuda_array_interface__`. No data is being copied. @@ -262,7 +262,7 @@ def memoryview( def __repr__(self) -> str: return ( f"<{self.__class__.__name__} size={format_bytes(self._size)} " - f"ptr={hex(self._ptr)} owner={repr(self._owner)}>" + f"ptr=0x{self._ptr:019X_} owner={self._owner!r}>" ) @@ -271,6 +271,10 @@ class Buffer(Serializable): Use the factory function `as_buffer` to create a Buffer instance. + Note + ---- + This buffer is untyped, so all indexing and sizes are in bytes. + Parameters ---------- owner diff --git a/python/cudf/cudf/core/buffer/utils.py b/python/cudf/cudf/core/buffer/utils.py index 3e2dc11e993..0c3a32a1210 100644 --- a/python/cudf/cudf/core/buffer/utils.py +++ b/python/cudf/cudf/core/buffer/utils.py @@ -50,7 +50,7 @@ def as_buffer( If spilling is enabled, a SpillableBuffer that refers ot a SpillableBufferOwner is returned. If `data` is owned by a spillable buffer, - it must either be "exposed" or spilled locked (called within an + it must either be "exposed" or spill locked (called within an acquire_spill_lock context). This is to guarantee that the memory of `data` isn't spilled before this function gets to calculate the offset of the new SpillableBuffer. @@ -132,8 +132,8 @@ def as_buffer( and get_spill_lock() is None ): raise ValueError( - "A owning spillable buffer must " - "either be exposed or spilled locked." + "An owning spillable buffer must " + "either be exposed or spill locked." ) ptr, size = get_ptr_and_size(data.__cuda_array_interface__) base_ptr = owner.get_ptr(mode="read") From 1f82e1b54a85a3aac5603cfb20bb622a015c8593 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 23 Nov 2023 16:21:49 +0100 Subject: [PATCH 09/22] Update python/cudf/cudf/core/buffer/utils.py Co-authored-by: Lawrence Mitchell --- python/cudf/cudf/core/buffer/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf/cudf/core/buffer/utils.py b/python/cudf/cudf/core/buffer/utils.py index 0c3a32a1210..8ebbe4c4ecb 100644 --- a/python/cudf/cudf/core/buffer/utils.py +++ b/python/cudf/cudf/core/buffer/utils.py @@ -48,7 +48,7 @@ def as_buffer( If copy-on-write is enabled, a ExposureTrackedBuffer that refers to a ExposureTrackedBufferOwner is returned. - If spilling is enabled, a SpillableBuffer that refers ot a + If spilling is enabled, a SpillableBuffer that refers to a SpillableBufferOwner is returned. If `data` is owned by a spillable buffer, it must either be "exposed" or spill locked (called within an acquire_spill_lock context). This is to guarantee that the memory of `data` From d22370c00e5f9830037c43afd6ebb07eb227087f Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 24 Nov 2023 10:19:59 +0100 Subject: [PATCH 10/22] fix errror check --- python/cudf/cudf/core/buffer/buffer.py | 2 +- python/cudf/cudf/tests/test_spilling.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index 0a4e73d5cbc..fb3edcb8efb 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -262,7 +262,7 @@ def memoryview( def __repr__(self) -> str: return ( f"<{self.__class__.__name__} size={format_bytes(self._size)} " - f"ptr=0x{self._ptr:019X_} owner={self._owner!r}>" + f"ptr={hex(self._ptr)} owner={self._owner!r}>" ) diff --git a/python/cudf/cudf/tests/test_spilling.py b/python/cudf/cudf/tests/test_spilling.py index 39a8d4c8b72..499092d2d16 100644 --- a/python/cudf/cudf/tests/test_spilling.py +++ b/python/cudf/cudf/tests/test_spilling.py @@ -550,7 +550,7 @@ def test_as_buffer_of_spillable_buffer(manager: SpillManager): with pytest.raises( ValueError, - match="buffer must either be exposed or spilled locked", + match="owning spillable buffer must either be exposed or spill locked", ): # Use `memory_info` to access device point _without_ making # the buffer unspillable. From 3053e3a43d729e51f81eb960f1caf28a3f5689f1 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 24 Nov 2023 10:39:08 +0100 Subject: [PATCH 11/22] clean up __str__ and __repr__ --- python/cudf/cudf/core/buffer/buffer.py | 8 +++++++- python/cudf/cudf/core/buffer/spillable_buffer.py | 4 ++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index fb3edcb8efb..528b50cfd21 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -259,7 +259,7 @@ def memoryview( ) return memoryview(host_buf).toreadonly() - def __repr__(self) -> str: + def __str__(self) -> str: return ( f"<{self.__class__.__name__} size={format_bytes(self._size)} " f"ptr={hex(self._ptr)} owner={self._owner!r}>" @@ -442,6 +442,12 @@ def deserialize(cls, header: dict, frames: list) -> Self: ) def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(owner={self._owner!r}, " + f"offset={self._offset!r}, size={self._size!r})" + ) + + def __str__(self) -> str: return ( f"<{self.__class__.__name__} size={format_bytes(self._size)} " f"offset={format_bytes(self._offset)} of {self._owner}>" diff --git a/python/cudf/cudf/core/buffer/spillable_buffer.py b/python/cudf/cudf/core/buffer/spillable_buffer.py index 7d35bb20b7a..19b602749f1 100644 --- a/python/cudf/cudf/core/buffer/spillable_buffer.py +++ b/python/cudf/cudf/core/buffer/spillable_buffer.py @@ -381,13 +381,13 @@ def memoryview( ) return ret - def __repr__(self) -> str: + def __str__(self) -> str: if self._ptr_desc["type"] != "gpu": ptr_info = str(self._ptr_desc) else: ptr_info = str(hex(self._ptr)) return ( - f"" From 63382d9528315cb89d87264e26aefcdd51107177 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 24 Nov 2023 10:50:12 +0100 Subject: [PATCH 12/22] Update python/cudf/cudf/core/buffer/buffer.py Co-authored-by: Lawrence Mitchell --- python/cudf/cudf/core/buffer/buffer.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index 0a4e73d5cbc..df83b5523cd 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -150,6 +150,13 @@ def _from_device_memory(cls, data: Any, exposed: bool) -> Self: ------- BufferOwner BufferOwner wrapping `data` + + Raises + ------ + AttributeError + If data does not support the cuda array interface + ValueError + If the resulting buffer has negative size """ # Bypass `__init__` and initialize attributes manually From bd59f7ca4f7968a3e97eaa01a6779f8c45500638 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 24 Nov 2023 11:41:40 +0100 Subject: [PATCH 13/22] type hint clean up --- python/cudf/cudf/core/buffer/buffer.py | 4 ++-- python/cudf/cudf/core/buffer/spillable_buffer.py | 12 ------------ 2 files changed, 2 insertions(+), 14 deletions(-) diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index 0b913292e05..4a9218e768c 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -215,7 +215,7 @@ def nbytes(self) -> int: return self._size @property - def owner(self) -> Any: + def owner(self) -> object: """Object owning the memory of the buffer.""" return self._owner @@ -323,7 +323,7 @@ def nbytes(self) -> int: return self._size @property - def owner(self) -> Any: + def owner(self) -> BufferOwner: """Object owning the memory of the buffer.""" return self._owner diff --git a/python/cudf/cudf/core/buffer/spillable_buffer.py b/python/cudf/cudf/core/buffer/spillable_buffer.py index 19b602749f1..f2d6dbb2280 100644 --- a/python/cudf/cudf/core/buffer/spillable_buffer.py +++ b/python/cudf/cudf/core/buffer/spillable_buffer.py @@ -331,10 +331,6 @@ def memory_info(self) -> Tuple[int, int, str]: ).__array_interface__["data"][0] return (ptr, self.nbytes, self._ptr_desc["type"]) - @property - def owner(self) -> Any: - return self._owner - @property def exposed(self) -> bool: return self._exposed @@ -343,14 +339,6 @@ def exposed(self) -> bool: def spillable(self) -> bool: return not self._exposed and len(self._spill_locks) == 0 - @property - def size(self) -> int: - return self._size - - @property - def nbytes(self) -> int: - return self._size - @property def last_accessed(self) -> float: return self._last_accessed From 411590dc90e7794aac60b71704e63c31a1edfa90 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 24 Nov 2023 11:46:15 +0100 Subject: [PATCH 14/22] test_buffer_str --- python/cudf/cudf/tests/test_buffer.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/python/cudf/cudf/tests/test_buffer.py b/python/cudf/cudf/tests/test_buffer.py index 2d561af5a95..aa7d727fcbd 100644 --- a/python/cudf/cudf/tests/test_buffer.py +++ b/python/cudf/cudf/tests/test_buffer.py @@ -67,13 +67,20 @@ def test_buffer_creation_from_any(): assert b.owner.owner.owner is ary +@pytest.mark.parametrize("size", [10, 2**10 + 500, 2**20]) +def test_buffer_str(size): + ary = cp.arange(size, dtype="uint8") + buf = as_buffer(ary) + assert f"size={size}" in repr(buf) + + @pytest.mark.parametrize( "size,expect", [(10, "10B"), (2**10 + 500, "1.49KiB"), (2**20, "1MiB")] ) def test_buffer_repr(size, expect): ary = cp.arange(size, dtype="uint8") buf = as_buffer(ary) - assert f"size={expect}" in repr(buf) + assert f"size={expect}" in str(buf) @pytest.mark.parametrize( From f450e2898cb48c97bcf210c514153f506fd12fce Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 24 Nov 2023 12:12:45 +0100 Subject: [PATCH 15/22] BufferOwner: now has exposed and mark_exposed --- python/cudf/cudf/core/buffer/buffer.py | 22 ++++++++++++++++++ .../core/buffer/exposure_tracked_buffer.py | 19 +-------------- .../cudf/cudf/core/buffer/spillable_buffer.py | 23 ++++++------------- 3 files changed, 30 insertions(+), 34 deletions(-) diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index 4a9218e768c..d1542fb3218 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -130,6 +130,7 @@ class BufferOwner(Serializable): _ptr: int _size: int _owner: object + _exposed: bool @classmethod def _from_device_memory(cls, data: Any, exposed: bool) -> Self: @@ -162,6 +163,7 @@ def _from_device_memory(cls, data: Any, exposed: bool) -> Self: # Bypass `__init__` and initialize attributes manually ret = cls.__new__(cls) ret._owner = data + ret._exposed = exposed if isinstance(data, rmm.DeviceBuffer): # Common case shortcut ret._ptr = data.ptr ret._size = data.size @@ -230,6 +232,26 @@ def __cuda_array_interface__(self) -> Mapping: "version": 0, } + @property + def exposed(self) -> bool: + """The current exposure status of the buffer + + This is used by ExposureTrackedBuffer to determine when a deep copy + is required and by SpillableBuffer to mark the buffer unspillable. + """ + return self._exposed + + def mark_exposed(self) -> None: + """Mark the buffer as "exposed" permanently + + This is used by ExposureTrackedBuffer to determine when a deep copy + is required and by SpillableBuffer to mark the buffer unspillable. + + Notice, once the exposure status becomes True, it will never change + back. + """ + self._exposed = True + def get_ptr(self, *, mode: Literal["read", "write"]) -> int: """Device pointer to the start of the buffer. diff --git a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py index 4f4a56470e8..2dd9cba7178 100644 --- a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py +++ b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py @@ -19,31 +19,14 @@ class ExposureTrackedBufferOwner(BufferOwner): the buffer has been exposed if the device pointer (integer or void*) has been accessed outside of ExposureTrackedBufferOwner. In this case, we have no control over knowing if the data is being modified by a third-party. - - Attributes - ---------- - _exposed - The current exposure status of the buffer. Notice, once the exposure - status becomes True, it should never change back. - _slices - The set of ExposureTrackedBuffer instances that point to this buffer. """ - _exposed: bool + # The set of ExposureTrackedBuffer instances that point to this buffer. _slices: weakref.WeakSet[ExposureTrackedBuffer] - @property - def exposed(self) -> bool: - return self._exposed - - def mark_exposed(self) -> None: - """Mark the buffer as "exposed" permanently""" - self._exposed = True - @classmethod def _from_device_memory(cls, data: Any, exposed: bool) -> Self: ret = super()._from_device_memory(data, exposed=exposed) - ret._exposed = exposed ret._slices = weakref.WeakSet() return ret diff --git a/python/cudf/cudf/core/buffer/spillable_buffer.py b/python/cudf/cudf/core/buffer/spillable_buffer.py index f2d6dbb2280..69b094ae890 100644 --- a/python/cudf/cudf/core/buffer/spillable_buffer.py +++ b/python/cudf/cudf/core/buffer/spillable_buffer.py @@ -88,10 +88,9 @@ class SpillableBufferOwner(BufferOwner): _spill_locks: weakref.WeakSet _last_accessed: float _ptr_desc: Dict[str, Any] - _exposed: bool _manager: SpillManager - def _finalize_init(self, ptr_desc: Dict[str, Any], exposed: bool) -> None: + def _finalize_init(self, ptr_desc: Dict[str, Any]) -> None: """Finish initialization of the spillable buffer This implements the common initialization that `_from_device_memory` @@ -101,8 +100,6 @@ def _finalize_init(self, ptr_desc: Dict[str, Any], exposed: bool) -> None: ---------- ptr_desc : dict Description of the memory. - exposed : bool, optional - Mark the buffer as permanently exposed (unspillable). """ from cudf.core.buffer.spill_manager import get_global_manager @@ -111,7 +108,6 @@ def _finalize_init(self, ptr_desc: Dict[str, Any], exposed: bool) -> None: self._spill_locks = weakref.WeakSet() self._last_accessed = time.monotonic() self._ptr_desc = ptr_desc - self._exposed = exposed manager = get_global_manager() if manager is None: raise ValueError( @@ -141,7 +137,7 @@ def _from_device_memory(cls, data: Any, exposed: bool) -> Self: Buffer representing the same device memory as `data` """ ret = super()._from_device_memory(data, exposed=exposed) - ret._finalize_init(ptr_desc={"type": "gpu"}, exposed=exposed) + ret._finalize_init(ptr_desc={"type": "gpu"}) return ret @classmethod @@ -178,9 +174,8 @@ def _from_host_memory(cls, data: Any) -> Self: ret._owner = None ret._ptr = 0 ret._size = data.nbytes - ret._finalize_init( - ptr_desc={"type": "cpu", "memoryview": data}, exposed=False - ) + ret._exposed = False + ret._finalize_init(ptr_desc={"type": "cpu", "memoryview": data}) return ret @property @@ -259,10 +254,10 @@ def mark_exposed(self) -> None: self._manager.spill_to_device_limit() with self.lock: - if not self._exposed: + if not self.exposed: self._manager.statistics.log_expose(self) self.spill(target="gpu") - self._exposed = True + super().mark_exposed() self._last_accessed = time.monotonic() def spill_lock(self, spill_lock: SpillLock) -> None: @@ -331,13 +326,9 @@ def memory_info(self) -> Tuple[int, int, str]: ).__array_interface__["data"][0] return (ptr, self.nbytes, self._ptr_desc["type"]) - @property - def exposed(self) -> bool: - return self._exposed - @property def spillable(self) -> bool: - return not self._exposed and len(self._spill_locks) == 0 + return not self.exposed and len(self._spill_locks) == 0 @property def last_accessed(self) -> float: From c423812ebbb462c74a059799e578292d673d95fb Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Wed, 13 Dec 2023 08:41:46 +0100 Subject: [PATCH 16/22] doc --- python/cudf/cudf/core/buffer/buffer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index d1542fb3218..8123a47b67b 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -17,7 +17,7 @@ from cudf.utils.string import format_bytes -def get_buffer_owner(data) -> Optional[BufferOwner]: +def get_buffer_owner(data: Any) -> Optional[BufferOwner]: """Get the owner of `data`, if one exists Search through the stack of data owners in order to find an @@ -26,7 +26,7 @@ def get_buffer_owner(data) -> Optional[BufferOwner]: Parameters ---------- data - The data object + The data object to search for a BufferOwner instance Return ------ From d81cd3bc894ec5af5e71892bbdc5a57fc21849f1 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 14 Dec 2023 11:03:46 +0100 Subject: [PATCH 17/22] merge `ExposureTrackedBufferOwner` and `BufferOwner` --- .../source/developer_guide/library_design.md | 4 +-- python/cudf/cudf/core/buffer/buffer.py | 16 ++++++++++ .../core/buffer/exposure_tracked_buffer.py | 32 ++----------------- python/cudf/cudf/core/buffer/utils.py | 10 ++---- 4 files changed, 24 insertions(+), 38 deletions(-) diff --git a/docs/cudf/source/developer_guide/library_design.md b/docs/cudf/source/developer_guide/library_design.md index 39d56559add..0b37de00f6b 100644 --- a/docs/cudf/source/developer_guide/library_design.md +++ b/docs/cudf/source/developer_guide/library_design.md @@ -325,9 +325,9 @@ This section describes the internal implementation details of the copy-on-write It is recommended that developers familiarize themselves with [the user-facing documentation](copy-on-write-user-doc) of this functionality before reading through the internals below. -The core copy-on-write implementation relies on the two classes `ExposureTrackedBufferOwner` and `ExposureTrackedBuffer`. +The core copy-on-write implementation relies on `ExposureTrackedBuffer` and the tracking features of `BufferOwner`. -An `ExposureTrackedBufferOwner` is a subclass of the `BufferOwner` that tracks internal and external references to its underlying memory. Internal references are tracked by maintaining [weak references](https://docs.python.org/3/library/weakref.html) to every `ExposureTrackedBuffer` of the underlying memory. External references are tracked through "exposure" status of the underlying memory. A buffer is considered exposed if the device pointer (integer or void*) has been handed out to a library outside of cudf. In this case, we have no way of knowing if the data are being modified by a third party. +`BufferOwner` tracks internal and external references to its underlying memory. Internal references are tracked by maintaining [weak references](https://docs.python.org/3/library/weakref.html) to every `ExposureTrackedBuffer` of the underlying memory. External references are tracked through "exposure" status of the underlying memory. A buffer is considered exposed if the device pointer (integer or void*) has been handed out to a library outside of cudf. In this case, we have no way of knowing if the data are being modified by a third party. `ExposureTrackedBuffer` is a subclass of `Buffer` that represents a _slice_ of the memory underlying an exposure tracked buffer. diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index 8123a47b67b..f53e2d92a05 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -4,6 +4,7 @@ import math import pickle +import weakref from types import SimpleNamespace from typing import Any, Dict, Literal, Mapping, Optional, Sequence, Tuple @@ -123,6 +124,12 @@ class BufferOwner(Serializable): the ones used throughout cuDF, can then refer to the same `BufferOwner` instance. + In order to implement copy-on-write and spillable buffers, we need the + ability to detect external access to the underlying memory. We say that + the buffer has been exposed if the device pointer (integer or void*) has + been accessed outside of BufferOwner. In this case, we have no control + over knowing if the data is being modified by a third-party. + Use `_from_device_memory` and `_from_host_memory` to create a new instance from either device or host memory respectively. """ @@ -131,6 +138,14 @@ class BufferOwner(Serializable): _size: int _owner: object _exposed: bool + # The set of buffers that point to this owner. + _slices: weakref.WeakSet[Buffer] + + def __init__(self): + raise ValueError( + f"do not create a {self.__class__} directly, please " + "use the factory function `cudf.core.buffer.as_buffer`" + ) @classmethod def _from_device_memory(cls, data: Any, exposed: bool) -> Self: @@ -164,6 +179,7 @@ def _from_device_memory(cls, data: Any, exposed: bool) -> Self: ret = cls.__new__(cls) ret._owner = data ret._exposed = exposed + ret._slices = weakref.WeakSet() if isinstance(data, rmm.DeviceBuffer): # Common case shortcut ret._ptr = data.ptr ret._size = data.size diff --git a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py index 2dd9cba7178..db42e7722e5 100644 --- a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py +++ b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py @@ -2,8 +2,7 @@ from __future__ import annotations -import weakref -from typing import Any, Literal, Mapping, Optional +from typing import Literal, Mapping, Optional from typing_extensions import Self @@ -11,31 +10,6 @@ from cudf.core.buffer.buffer import Buffer, BufferOwner -class ExposureTrackedBufferOwner(BufferOwner): - """A Buffer that tracks its "expose" status. - - In order to implement copy-on-write and spillable buffers, we need the - ability to detect external access to the underlying memory. We say that - the buffer has been exposed if the device pointer (integer or void*) has - been accessed outside of ExposureTrackedBufferOwner. In this case, we have - no control over knowing if the data is being modified by a third-party. - """ - - # The set of ExposureTrackedBuffer instances that point to this buffer. - _slices: weakref.WeakSet[ExposureTrackedBuffer] - - @classmethod - def _from_device_memory(cls, data: Any, exposed: bool) -> Self: - ret = super()._from_device_memory(data, exposed=exposed) - ret._slices = weakref.WeakSet() - return ret - - @property - def __cuda_array_interface__(self) -> Mapping: - self.mark_exposed() - return super().__cuda_array_interface__ - - class ExposureTrackedBuffer(Buffer): """An exposure tracked buffer. @@ -49,11 +23,11 @@ class ExposureTrackedBuffer(Buffer): The size of the slice (in bytes) """ - _owner: ExposureTrackedBufferOwner + _owner: BufferOwner def __init__( self, - owner: ExposureTrackedBufferOwner, + owner: BufferOwner, offset: int = 0, size: Optional[int] = None, ) -> None: diff --git a/python/cudf/cudf/core/buffer/utils.py b/python/cudf/cudf/core/buffer/utils.py index 8ebbe4c4ecb..22bc439931b 100644 --- a/python/cudf/cudf/core/buffer/utils.py +++ b/python/cudf/cudf/core/buffer/utils.py @@ -13,10 +13,7 @@ get_buffer_owner, get_ptr_and_size, ) -from cudf.core.buffer.exposure_tracked_buffer import ( - ExposureTrackedBuffer, - ExposureTrackedBufferOwner, -) +from cudf.core.buffer.exposure_tracked_buffer import ExposureTrackedBuffer from cudf.core.buffer.spill_manager import get_global_manager from cudf.core.buffer.spillable_buffer import ( SpillableBuffer, @@ -45,8 +42,7 @@ def as_buffer( Raises ValueError if `data` isn't C-contiguous. - If copy-on-write is enabled, a ExposureTrackedBuffer that refers to a - ExposureTrackedBufferOwner is returned. + If copy-on-write is enabled, a ExposureTrackedBuffer is returned. If spilling is enabled, a SpillableBuffer that refers to a SpillableBufferOwner is returned. If `data` is owned by a spillable buffer, @@ -104,7 +100,7 @@ def as_buffer( owner_class = SpillableBufferOwner buffer_class = SpillableBuffer elif get_option("copy_on_write"): - owner_class = ExposureTrackedBufferOwner + owner_class = BufferOwner buffer_class = ExposureTrackedBuffer else: owner_class = BufferOwner From 89b45fdf3d9caeb453f36e1ff97fbb6c6b5c3ae3 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 14 Dec 2023 11:07:38 +0100 Subject: [PATCH 18/22] spill serialize now returns a Buffer --- python/cudf/cudf/core/abc.py | 1 - python/cudf/cudf/core/buffer/buffer.py | 11 ----------- .../cudf/cudf/core/buffer/spillable_buffer.py | 18 ++++++++++-------- python/cudf/cudf/tests/test_spilling.py | 6 +++--- 4 files changed, 13 insertions(+), 23 deletions(-) diff --git a/python/cudf/cudf/core/abc.py b/python/cudf/cudf/core/abc.py index b47515da2b3..ea44a18ba2a 100644 --- a/python/cudf/cudf/core/abc.py +++ b/python/cudf/cudf/core/abc.py @@ -92,7 +92,6 @@ def device_serialize(self): isinstance( f, ( - cudf.core.buffer.BufferOwner, cudf.core.buffer.Buffer, memoryview, ), diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index f53e2d92a05..2c4a46a9804 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -237,17 +237,6 @@ def owner(self) -> object: """Object owning the memory of the buffer.""" return self._owner - @property - def __cuda_array_interface__(self) -> Mapping: - """Implementation of the CUDA Array Interface.""" - return { - "data": (self.get_ptr(mode="write"), False), - "shape": (self.size,), - "strides": None, - "typestr": "|u1", - "version": 0, - } - @property def exposed(self) -> bool: """The current exposure status of the buffer diff --git a/python/cudf/cudf/core/buffer/spillable_buffer.py b/python/cudf/cudf/core/buffer/spillable_buffer.py index 69b094ae890..0e333806882 100644 --- a/python/cudf/cudf/core/buffer/spillable_buffer.py +++ b/python/cudf/cudf/core/buffer/spillable_buffer.py @@ -436,7 +436,7 @@ def serialize(self) -> Tuple[dict, list]: to memory already owned by an existing `SpillableBufferOwner`. """ header: Dict[str, Any] = {} - frames: List[BufferOwner | memoryview] + frames: List[Buffer | memoryview] with self._owner.lock: header["type-serialized"] = pickle.dumps(self.__class__) header["owner-type-serialized"] = pickle.dumps(type(self._owner)) @@ -449,13 +449,15 @@ def serialize(self) -> Tuple[dict, list]: self.spill_lock(spill_lock) ptr, size, _ = self.memory_info() frames = [ - BufferOwner._from_device_memory( - cuda_array_interface_wrapper( - ptr=ptr, - size=size, - owner=(self._owner, spill_lock), - ), - exposed=False, + Buffer( + owner=BufferOwner._from_device_memory( + cuda_array_interface_wrapper( + ptr=ptr, + size=size, + owner=(self._owner, spill_lock), + ), + exposed=False, + ) ) ] return header, frames diff --git a/python/cudf/cudf/tests/test_spilling.py b/python/cudf/cudf/tests/test_spilling.py index 499092d2d16..66adce8dc06 100644 --- a/python/cudf/cudf/tests/test_spilling.py +++ b/python/cudf/cudf/tests/test_spilling.py @@ -21,7 +21,7 @@ import cudf.options from cudf.core.abc import Serializable from cudf.core.buffer import ( - BufferOwner, + Buffer, acquire_spill_lock, as_buffer, get_spill_lock, @@ -439,7 +439,7 @@ def test_serialize_device(manager, target, view): header, frames = df1.device_serialize() assert len(frames) == 1 if target == "gpu": - assert isinstance(frames[0], BufferOwner) + assert isinstance(frames[0], Buffer) assert not single_column_df_data(df1).is_spilled assert not single_column_df_data(df1).spillable frames[0] = cupy.array(frames[0], copy=True) @@ -499,7 +499,7 @@ def test_serialize_cuda_dataframe(manager: SpillManager): buf: SpillableBuffer = single_column_df_data(df1) assert len(buf.owner._spill_locks) == 1 assert len(frames) == 1 - assert isinstance(frames[0], BufferOwner) + assert isinstance(frames[0], Buffer) assert frames[0].get_ptr(mode="read") == buf.get_ptr(mode="read") frames[0] = cupy.array(frames[0], copy=True) From 5ffba0c3dc2631f7c7da26cc0703ea41d1682bbd Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 12 Jan 2024 08:22:57 +0100 Subject: [PATCH 19/22] doc Co-authored-by: Vyas Ramasubramani --- python/cudf/cudf/core/buffer/buffer.py | 2 +- python/cudf/cudf/core/buffer/utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index 2c4a46a9804..ef5362b4b4f 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -128,7 +128,7 @@ class BufferOwner(Serializable): ability to detect external access to the underlying memory. We say that the buffer has been exposed if the device pointer (integer or void*) has been accessed outside of BufferOwner. In this case, we have no control - over knowing if the data is being modified by a third-party. + over knowing if the data is being modified by a third party. Use `_from_device_memory` and `_from_host_memory` to create a new instance from either device or host memory respectively. diff --git a/python/cudf/cudf/core/buffer/utils.py b/python/cudf/cudf/core/buffer/utils.py index 22bc439931b..ebb2261baae 100644 --- a/python/cudf/cudf/core/buffer/utils.py +++ b/python/cudf/cudf/core/buffer/utils.py @@ -42,7 +42,7 @@ def as_buffer( Raises ValueError if `data` isn't C-contiguous. - If copy-on-write is enabled, a ExposureTrackedBuffer is returned. + If copy-on-write is enabled, an ExposureTrackedBuffer is returned. If spilling is enabled, a SpillableBuffer that refers to a SpillableBufferOwner is returned. If `data` is owned by a spillable buffer, From c7d63783aa5c6089ef916c70b59c547294bce504 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 12 Jan 2024 10:58:33 +0100 Subject: [PATCH 20/22] Buffer.memoryview(): removed the size and offset arguments --- python/cudf/cudf/core/buffer/buffer.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index ef5362b4b4f..92a61c4ffff 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -371,11 +371,8 @@ def __getitem__(self, key: slice) -> Self: def get_ptr(self, *, mode: Literal["read", "write"]) -> int: return self._owner.get_ptr(mode=mode) + self._offset - def memoryview( - self, *, offset: int = 0, size: Optional[int] = None - ) -> memoryview: - size = self._size if size is None else size - return self._owner.memoryview(offset=self._offset + offset, size=size) + def memoryview(self) -> memoryview: + return self._owner.memoryview(offset=self._offset, size=self._size) def copy(self, deep: bool = True) -> Self: """Return a copy of Buffer. From a00a12bb279d9f9827b301ae4cfaedb0fffe4ce2 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 12 Jan 2024 11:12:17 +0100 Subject: [PATCH 21/22] copyrights --- python/cudf/cudf/core/abc.py | 2 +- python/cudf/cudf/core/buffer/__init__.py | 2 +- python/cudf/cudf/core/buffer/buffer.py | 2 +- python/cudf/cudf/core/buffer/exposure_tracked_buffer.py | 2 +- python/cudf/cudf/core/buffer/spill_manager.py | 2 +- python/cudf/cudf/core/buffer/spillable_buffer.py | 2 +- python/cudf/cudf/core/buffer/utils.py | 2 +- python/cudf/cudf/tests/test_buffer.py | 2 +- python/cudf/cudf/tests/test_copying.py | 2 +- python/cudf/cudf/tests/test_spilling.py | 2 +- 10 files changed, 10 insertions(+), 10 deletions(-) diff --git a/python/cudf/cudf/core/abc.py b/python/cudf/cudf/core/abc.py index ea44a18ba2a..ce6bb83bc77 100644 --- a/python/cudf/cudf/core/abc.py +++ b/python/cudf/cudf/core/abc.py @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2023, NVIDIA CORPORATION. +# Copyright (c) 2020-2024, NVIDIA CORPORATION. """Common abstract base classes for cudf.""" import pickle diff --git a/python/cudf/cudf/core/buffer/__init__.py b/python/cudf/cudf/core/buffer/__init__.py index b5799c076f9..9b9774c12be 100644 --- a/python/cudf/cudf/core/buffer/__init__.py +++ b/python/cudf/cudf/core/buffer/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2023, NVIDIA CORPORATION. +# Copyright (c) 2022-2024, NVIDIA CORPORATION. from cudf.core.buffer.buffer import ( Buffer, diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index 92a61c4ffff..afaa45e989e 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2023, NVIDIA CORPORATION. +# Copyright (c) 2020-2024, NVIDIA CORPORATION. from __future__ import annotations diff --git a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py index db42e7722e5..4c08016adbb 100644 --- a/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py +++ b/python/cudf/cudf/core/buffer/exposure_tracked_buffer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2023, NVIDIA CORPORATION. +# Copyright (c) 2020-2024, NVIDIA CORPORATION. from __future__ import annotations diff --git a/python/cudf/cudf/core/buffer/spill_manager.py b/python/cudf/cudf/core/buffer/spill_manager.py index 9735ff42a46..3e654e01401 100644 --- a/python/cudf/cudf/core/buffer/spill_manager.py +++ b/python/cudf/cudf/core/buffer/spill_manager.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2023, NVIDIA CORPORATION. +# Copyright (c) 2022-2024, NVIDIA CORPORATION. from __future__ import annotations diff --git a/python/cudf/cudf/core/buffer/spillable_buffer.py b/python/cudf/cudf/core/buffer/spillable_buffer.py index 0e333806882..aeac4b76e58 100644 --- a/python/cudf/cudf/core/buffer/spillable_buffer.py +++ b/python/cudf/cudf/core/buffer/spillable_buffer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2023, NVIDIA CORPORATION. +# Copyright (c) 2022-2024, NVIDIA CORPORATION. from __future__ import annotations diff --git a/python/cudf/cudf/core/buffer/utils.py b/python/cudf/cudf/core/buffer/utils.py index ebb2261baae..3d65e5c38ab 100644 --- a/python/cudf/cudf/core/buffer/utils.py +++ b/python/cudf/cudf/core/buffer/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2023, NVIDIA CORPORATION. +# Copyright (c) 2022-2024, NVIDIA CORPORATION. from __future__ import annotations diff --git a/python/cudf/cudf/tests/test_buffer.py b/python/cudf/cudf/tests/test_buffer.py index aa7d727fcbd..03637e05eae 100644 --- a/python/cudf/cudf/tests/test_buffer.py +++ b/python/cudf/cudf/tests/test_buffer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2023, NVIDIA CORPORATION. +# Copyright (c) 2020-2024, NVIDIA CORPORATION. import cupy as cp import pytest diff --git a/python/cudf/cudf/tests/test_copying.py b/python/cudf/cudf/tests/test_copying.py index 4151aa76753..e737a73e86b 100644 --- a/python/cudf/cudf/tests/test_copying.py +++ b/python/cudf/cudf/tests/test_copying.py @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2023, NVIDIA CORPORATION. +# Copyright (c) 2020-2024, NVIDIA CORPORATION. import cupy as cp import numpy as np diff --git a/python/cudf/cudf/tests/test_spilling.py b/python/cudf/cudf/tests/test_spilling.py index 66adce8dc06..7e66a7ab4ba 100644 --- a/python/cudf/cudf/tests/test_spilling.py +++ b/python/cudf/cudf/tests/test_spilling.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2023, NVIDIA CORPORATION. +# Copyright (c) 2022-2024, NVIDIA CORPORATION. import importlib import random From 6a9597adec73ab11e16045e5ba1b9655a9142e79 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 12 Jan 2024 13:16:52 +0100 Subject: [PATCH 22/22] moved get_buffer_owner to utils.py --- python/cudf/cudf/core/buffer/buffer.py | 24 ------------------------ python/cudf/cudf/core/buffer/utils.py | 25 ++++++++++++++++++++++++- 2 files changed, 24 insertions(+), 25 deletions(-) diff --git a/python/cudf/cudf/core/buffer/buffer.py b/python/cudf/cudf/core/buffer/buffer.py index afaa45e989e..8d278c9c065 100644 --- a/python/cudf/cudf/core/buffer/buffer.py +++ b/python/cudf/cudf/core/buffer/buffer.py @@ -18,30 +18,6 @@ from cudf.utils.string import format_bytes -def get_buffer_owner(data: Any) -> Optional[BufferOwner]: - """Get the owner of `data`, if one exists - - Search through the stack of data owners in order to find an - owner BufferOwner (incl. subclasses). - - Parameters - ---------- - data - The data object to search for a BufferOwner instance - - Return - ------ - BufferOwner or None - The owner of `data` if found otherwise None. - """ - - if isinstance(data, BufferOwner): - return data - if hasattr(data, "owner"): - return get_buffer_owner(data.owner) - return None - - def host_memory_allocation(nbytes: int) -> memoryview: """Allocate host memory using NumPy diff --git a/python/cudf/cudf/core/buffer/utils.py b/python/cudf/cudf/core/buffer/utils.py index 3d65e5c38ab..c2ec7effd13 100644 --- a/python/cudf/cudf/core/buffer/utils.py +++ b/python/cudf/cudf/core/buffer/utils.py @@ -10,7 +10,6 @@ Buffer, BufferOwner, cuda_array_interface_wrapper, - get_buffer_owner, get_ptr_and_size, ) from cudf.core.buffer.exposure_tracked_buffer import ExposureTrackedBuffer @@ -23,6 +22,30 @@ from cudf.options import get_option +def get_buffer_owner(data: Any) -> Optional[BufferOwner]: + """Get the owner of `data`, if one exists + + Search through the stack of data owners in order to find an + owner BufferOwner (incl. subclasses). + + Parameters + ---------- + data + The data object to search for a BufferOwner instance + + Return + ------ + BufferOwner or None + The owner of `data` if found otherwise None. + """ + + if isinstance(data, BufferOwner): + return data + if hasattr(data, "owner"): + return get_buffer_owner(data.owner) + return None + + def as_buffer( data: Union[int, Any], *,