Skip to content

Commit

Permalink
unify as_spillable_buffer, as_exposure_tracked_buffer, and as_buffer
Browse files Browse the repository at this point in the history
  • Loading branch information
@madsbk
madsbk committed Aug 3, 2023
1 parent 4790117 commit 9d04a0a
Show file tree
Hide file tree
Showing 5 changed files with 96 additions and 175 deletions.
14 changes: 7 additions & 7 deletions docs/cudf/source/developer_guide/library_design.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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`.


Expand Down
39 changes: 16 additions & 23 deletions python/cudf/cudf/core/buffer/buffer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
----------
Expand All @@ -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


Expand Down Expand Up @@ -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.
Expand All @@ -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
-------
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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(
Expand Down
70 changes: 3 additions & 67 deletions python/cudf/cudf/core/buffer/exposure_tracked_buffer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down
73 changes: 5 additions & 68 deletions python/cudf/cudf/core/buffer/spillable_buffer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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.
Expand All @@ -195,15 +131,15 @@ 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
-------
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

Expand Down Expand Up @@ -525,7 +461,8 @@ def serialize(self) -> Tuple[dict, list]:
ptr=ptr,
size=size,
owner=(self._owner, spill_lock),
)
),
exposed=False,
)
]
return header, frames
Expand Down
Loading

0 comments on commit 9d04a0a

Please sign in to comment.