Skip to content

Make UMAP callback pickleable#6402

Merged
rapids-bot[bot] merged 1 commit intorapidsai:branch-25.04from
jcrist:make-callback-pickleable
Mar 6, 2025
Merged

Make UMAP callback pickleable#6402
rapids-bot[bot] merged 1 commit intorapidsai:branch-25.04from
jcrist:make-callback-pickleable

Conversation

@jcrist
Copy link
Copy Markdown
Member

@jcrist jcrist commented Mar 6, 2025

The structure of this callback object should really be cleaned up (by moving the callback to the stack before the call to UMAP we can avoid having an object with pointers to itself and better manage the lifetime of the callback instance). This would also let us avoid implementing this callback in cython, which mucks up the default pickling implementation. As is though this is the easiest way to get things pickleable. Users will still need to define __reduce__ themselves if they add any state to an instance, but given that this code was added for an internal use case and hasn't seen any changes or expansion in 6 years I don't suspect there are many external users of this API. Can always clean this up later if a user asks about it.

Fixes #5282.

The structure of this callback object should really be cleaned up (by
moving the callback to the stack before the call to UMAP we can avoid
having an object with pointers to itself _and_ better manage the
lifetime of the callback instance). This would also let us avoid
implementing this callback in cython, which mucks up the default
pickling implementation. As is though this is the easiest way to get
things pickleable. Users will still need to define `__reduce__`
themselves if they add any state to an instance, but given that this
code was added for an internal use case and hasn't seen any changes or
expansion in 6 years I don't suspect there are many external users of
this API. Can always clean this up later if a user asks about it.
@jcrist jcrist requested a review from a team as a code owner March 6, 2025 17:31
@jcrist jcrist requested review from csadorf and wphicks March 6, 2025 17:31
@github-actions github-actions Bot added the Cython / Python Cython or Python issue label Mar 6, 2025
@jcrist jcrist added bug Something isn't working non-breaking Non-breaking change labels Mar 6, 2025
pass

def on_train_end(self, embeddings):
pass
Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Added default implementations of the callbacks.

if self.native_callback.pyCallbackClass == NULL:
raise ValueError(
"You need to call `super().__init__` in your callback."
)
Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

No need for this - __cinit__ is always called

assert cuml_trust > 0.9


def test_callback():
Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Moved this test to test_umap.py since it's more a test that UMAP calls the callback properly, not a test of the callback itself.

Copy link
Copy Markdown
Contributor

@csadorf csadorf left a comment

Choose a reason for hiding this comment

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

LGTM!

@jcrist
Copy link
Copy Markdown
Member Author

jcrist commented Mar 6, 2025

/merge

@rapids-bot rapids-bot Bot merged commit ee276d5 into rapidsai:branch-25.04 Mar 6, 2025
@jcrist jcrist deleted the make-callback-pickleable branch March 6, 2025 19:43
@jcrist jcrist mentioned this pull request Mar 6, 2025
9 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working Cython / Python Cython or Python issue non-breaking Non-breaking change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Can not serialize UMAP model when a custom callback is defined

2 participants