Add RAFT option to make_blobs by NIne-WIngEd · Pull Request #8482 · NVIDIA/cuml

NIne-WIngEd · 2026-08-16T02:45:27Z

Adds an opt-in RAFT path to cuml.datasets.make_blobs.

The existing CuPy path stays the default.

Added focused tests for seeded reproducibility and C/F layouts.

Fixes #7363.

copy-pr-bot · 2026-08-16T02:45:31Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

coderabbitai · 2026-08-16T02:47:24Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5c4e8ae1-4aa7-43b3-af9c-cb230866b218

📥 Commits

Reviewing files that changed from the base of the PR and between 0d3a802 and 8bcdb37.

📒 Files selected for processing (4)

python/cuml/cuml/datasets/CMakeLists.txt
python/cuml/cuml/datasets/_blobs.pyx
python/cuml/cuml/datasets/blobs.py
python/cuml/tests/test_make_blobs.py

Included review availability: Your plan includes up to 12 reviews per rolling hour; 10 remain after this review.

📝 Walkthrough

Summary by CodeRabbit

New Features
- Added an optional RAFT-backed execution path for make_blobs dataset generation.
- Supports generated or explicitly provided centers, configurable dtypes and memory layouts, reproducible random states, and optional center output.
- Added GPU-backed dataset generation with validation for supported inputs and clear errors for unsupported configurations.
Bug Fixes
- Improved consistency and parity between supported dataset-generation paths.

Walkthrough

Adds an optional RAFT-backed make_blobs path. The change includes a Cython GPU bridge, input validation, native dispatch, documentation, and tests for parity, reproducibility, layouts, dtypes, centers, and invalid inputs.

Changes

RAFT-backed make_blobs

Layer / File(s)	Summary
GPU bridge and build integration `python/cuml/cuml/datasets/CMakeLists.txt`, `python/cuml/cuml/datasets/_blobs.pyx`	Adds the Cython module to the build. The wrapper allocates CuPy outputs, dispatches float32 and float64 native overloads, synchronizes the handle, and returns generated data and labels.
RAFT API path and validation `python/cuml/cuml/datasets/blobs.py`, `python/cuml/tests/test_make_blobs.py`	Adds `use_raft=False`, validates RAFT-supported arguments, invokes native generation, documents restrictions, and tests parity, reproducibility, centers, layouts, dtypes, and validation errors.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Merge Risk: ⚪ Minimal · up to 8bcdb

The PR adds an opt-in RAFT path while preserving the existing default behavior; no actionable merge-blocking risk remains after normal checks and review.

Possibly related PRs

NVIDIA/cuml#8481: Implements the same RAFT-backed make_blobs path in the same dataset files and tests.

Suggested reviewers: dantegd

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely identifies the main change: adding a RAFT option to make_blobs.
Description check	✅ Passed	The description accurately covers the opt-in RAFT path, default CuPy behavior, tests, and linked issue.
Linked Issues check	✅ Passed	The changes implement the RAFT option, add a Python wrapper, and preserve the existing CuPy path as required by issue `#7363`.
Out of Scope Changes check	✅ Passed	The build configuration, implementation, documentation, tests, and copyright update support the stated objectives without unrelated code changes.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

Create stacked PR
Commit on current branch

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

csadorf

Thanks a lot for the contribution! I have a few change requests and questions before we can move forward.

csadorf · 2026-08-17T20:01:36Z

    return_centers=False,
    order="F",
    dtype="float32",
+    use_raft=False,


⚠️ Whether we use RAFT or CuPy for the actual implementation is an implementation detail and should not be exposed within the Python API. The routing should be purely internal. I ran some preliminary benchmarks and it seems that the RAFT implementation is generally much faster.

We should maintain the CuPy implementation only to overcome any present limitations or if there is a non-trivial regime where CuPy is generally faster (which does not seem to be the case).

Yep, makes sense. I removed use_raft from the public API and made the routing internal. RAFT is now the normal path, with the CuPy path kept only for compatibility cases that the native path can't preserve yet.

csadorf · 2026-08-17T20:05:52Z

+            <int64_t>n_centers,
+            row_c,
+            <const float*>ctr_p,
+            <const float*>0,


The native API already accepts per-cluster standard deviations through the cluster_std device pointer. Could we pass an array here instead of restricting the RAFT path to a scalar cluster_std? That would preserve the existing Python API and eliminate one fallback case.

I checked this before wiring it up. The RAFT API accepts the device pointer, but the current implementation looks like it indexes cluster_std using the row index rather than the cluster label, so an n_clusters-length array looks unsafe there. I kept sequence cluster_std on the CuPy path for now. Would you prefer that I handle the RAFT-side issue separately first?

csadorf · 2026-08-17T20:05:53Z

+
+    if return_centers and made_ctr:
+        raise ValueError(
+            "`return_centers=True` with generated centers is not supported "


Could we preserve return_centers=True by generating the centers in the Python compatibility layer and passing them explicitly to RAFT? That would allow generated-center calls to use RAFT without narrowing the existing API.

To be clear, we should only do that for this specific condition.

Yep, updated

csadorf · 2026-08-17T20:05:54Z

+
+    # Native labels are integers; Python make_blobs has historically returned
+    # them in the requested floating dtype, so keep that behavior here.
+    h.sync()


make_blobs is decorated with mlfunc, so its CuPy operations run on Stream.ptds, and the default RAFT handle also uses cudaStreamPerThread. Why do we need to sync here?

You're right. Since both are using PTDS here, the explicit sync isn't needed. Removed it.

csadorf · 2026-08-17T20:09:19Z

+    # Native labels are integers; Python make_blobs has historically returned
+    # them in the requested floating dtype, so keep that behavior here.


This in-line comment seems to be at the wrong spot.

Yep, fixed. I removed that comment along with the explicit sync, so there isn't a misplaced explanation there anymore.

Add RAFT path for make_blobs

8bcdb37

NIne-WIngEd requested review from a team as code owners August 16, 2026 02:45

NIne-WIngEd requested a review from csadorf August 16, 2026 02:45

github-actions Bot added Cython / Python Cython or Python issue CMake labels Aug 16, 2026

github-actions Bot assigned NIne-WIngEd Aug 16, 2026

csadorf requested changes Aug 17, 2026

View reviewed changes

		# Native labels are integers; Python make_blobs has historically returned
		# them in the requested floating dtype, so keep that behavior here.

Uh oh!

Conversation

NIne-WIngEd commented Aug 16, 2026

Uh oh!

copy-pr-bot Bot commented Aug 16, 2026

Uh oh!

coderabbitai Bot commented Aug 16, 2026

Summary by CodeRabbit

Walkthrough

Changes

Uh oh!

csadorf left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants