Skip to content

Unify Copy-On-Write and Spilling - #15436

Merged
rapids-bot[bot] merged 30 commits into
NVIDIA:branch-24.06from
madsbk:cow_and_spilling
Apr 20, 2024
Merged

Unify Copy-On-Write and Spilling#15436
rapids-bot[bot] merged 30 commits into
NVIDIA:branch-24.06from
madsbk:cow_and_spilling

Conversation

@madsbk

@madsbk madsbk commented Apr 2, 2024

Copy link
Copy Markdown
Contributor

Description

This is the final step to unify COW and spilling. Now, SpillableBuffer inherits from ExposureTrackedBuffer so the final class hierarchy becomes:

SpillableBufferOwner -> BufferOwner 
SpillableBuffer -> ExposureTrackedBuffer -> Buffer 

Additionally, spill-on-demand is now set globally using set_spill_on_demand_globally() instead of in the SpillManager constructor.

Checklist

  • I am familiar with the Contributing Guidelines.
  • New or existing tests cover these changes.
  • The documentation is up to date with these changes.

@github-actions github-actions Bot added the Python Affects Python cuDF API. label Apr 2, 2024
@madsbk madsbk added improvement Improvement / enhancement to an existing function non-breaking Non-breaking change and removed Python Affects Python cuDF API. labels Apr 2, 2024
@github-actions github-actions Bot added the Python Affects Python cuDF API. label Apr 3, 2024
Comment on lines +120 to +144
def __init__(
self,
ptr: int,
size: int,
owner: object,
exposed: bool,
):
"""Create a new buffer owner.

Do not use this directly, instead use `_from_device_memory` or
`_from_host_memory`.

Raises
------
ValueError
If size is negative
"""
if size < 0:
raise ValueError("size cannot be negative")

self._ptr = ptr
self._size = size
self._owner = owner
self._exposed = exposed
self._slices = weakref.WeakSet()

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Notice, I have removed the direct-use-of-Buffer guard. I think the four mandatory arguments are enough to prevent accidental use.

@madsbk
madsbk marked this pull request as ready for review April 3, 2024 12:44
@madsbk
madsbk requested a review from a team as a code owner April 3, 2024 12:44
@madsbk
madsbk requested review from bdice and brandon-b-miller April 3, 2024 12:44

@bdice bdice left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A few small suggestions but overall this looks good to me.

Comment thread python/cudf/cudf/core/buffer/spill_manager.py Outdated
Comment thread python/cudf/cudf/core/buffer/spill_manager.py Outdated
Comment thread python/cudf/cudf/core/buffer/spill_manager.py Outdated
Comment thread python/cudf/cudf/core/buffer/spill_manager.py Outdated
Comment thread python/cudf/cudf/core/buffer/spillable_buffer.py
Comment thread python/cudf/cudf/core/buffer/spillable_buffer.py Outdated
Co-authored-by: Bradley Dice <bdice@bradleydice.com>

@vyasr vyasr left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm excited to have this happen! Thanks for pushing through Mads.

We should also not forget about these two bits of work proposed from the last round:

I also wonder if there is a sensible way to unify the SpillableBufferOwner with the BufferOwner (current names, which will be changed in the future). Not in scope for this PR, but some of my prior concerns about the Liskov substitutability of different Buffer types still apply there (also CC @wence-).

Comment thread python/cudf/cudf/core/buffer/buffer.py Outdated
Comment thread python/cudf/cudf/tests/test_copying.py Outdated
Comment thread python/cudf/cudf/tests/test_spilling.py
Comment thread python/cudf/cudf/core/buffer/spill_manager.py

@wence- wence- left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor

@property
def exposed(self) -> bool:
return self._owner.exposed
self.owner._slices.add(self)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

question: should we assert isinstance(owner, BufferOwner) here? Or is it not necessary because we implicitly require it via the _slices attribute?

@madsbk madsbk Apr 11, 2024

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess, we could add assert isinstance(owner, BufferOwner) in Buffer.__init__ but I am not sure if that is pythonic?

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

An assertion would basically be there to protect against developer error. There's nothing unpythonic about it, that's definitely what the functionality is there for. The main question I would have is, how would we end up with something other than a BufferOwner here, and how bad would the failure mode be? Unless the assertion makes it markedly easier to debug the code it doesn't add a whole lot of value IMHO.

@madsbk madsbk Apr 17, 2024

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it is arguable a little unpythonic to prevent duck-typing by using isinstance to force a specify type (as opposed to a general ABC/protocol type).
Anyways, in this case as_buffer() is the only one creating ExposureTrackedBuffer and it explicitly sets the owner to SpillableBufferOwner | SpillableBuffer | BufferOwner so I don't think a type error here is likely.

Comment thread python/cudf/cudf/core/buffer/exposure_tracked_buffer.py Outdated
Comment thread python/cudf/cudf/core/buffer/spillable_buffer.py
Comment thread python/cudf/cudf/core/buffer/spillable_buffer.py
Comment thread python/cudf/cudf/tests/test_spilling.py Outdated
madsbk and others added 3 commits April 11, 2024 15:26
Co-authored-by: Lawrence Mitchell <wence@gmx.li>
@madsbk

madsbk commented Apr 11, 2024

Copy link
Copy Markdown
Contributor Author

Regarding renaming:

`BufferOwner` -> `Buffer`
`Buffer` -> `BufferView`

I think we should do it in a follow up PR, it will touch a lot of files throughout cudf!

Regarding #13801 (comment):

We should reassess this evaluation after the next PR unifying COW and spilling. I don't think I agree with this statement, it implies leaking knowledge of cudf buffer internals to dask. Once we've finished the unification we should revisit whether there's a more API-friendly way of doing this. If not, we need to think about the appropriate generalization of our exposure semantics to generic CAI usage.

I agree, if dask or ucx should handle this, we need to publish a buffer API that includes this .get_ptr() semantic. However, the current approach might not be as bad as I first thought. As @wence- mention, pytorch is doing something very similar.

@madsbk
madsbk requested review from vyasr and wence- April 11, 2024 15:04
@vyasr

vyasr commented Apr 16, 2024

Copy link
Copy Markdown
Collaborator

Happy to address one or both of those in a follow-up, especially since the second one apparently requires further discussion (I haven't given it any more thought yet). For the first, we should try to turn that around quickly to minimize conflicts since as you say it will touch a lot of files. Once you have that PR ready (after this one merges) feel free to ping me aggressively for reviews.

@vyasr vyasr left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm happy with this when Lawrence is, and when Prem is satisfied with the c-o-w testing.

Comment thread python/cudf/cudf/core/buffer/spillable_buffer.py
@madsbk
madsbk requested a review from galipremsagar April 17, 2024 07:25

@galipremsagar galipremsagar left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good to gooo!!

@vyasr

vyasr commented Apr 20, 2024

Copy link
Copy Markdown
Collaborator

/merge

@rapids-bot
rapids-bot Bot merged commit 96903bb into NVIDIA:branch-24.06 Apr 20, 2024
@madsbk
madsbk deleted the cow_and_spilling branch May 1, 2025 07:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

improvement Improvement / enhancement to an existing function non-breaking Non-breaking change Python Affects Python cuDF API.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants