[FEA] Add host-streaming support to single-GPU KMeans - #8248
Conversation
tarang-jain
added
feature request
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds ChangesHost-Resident KMeans Streaming
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/src/kmeans/kmeans_fit.cu`:
- Around line 36-41: The current ptr_is_host logic treats any
cudaPointerGetAttributes failure as a host pointer; instead, change the handling
so that only a successful cudaPointerGetAttributes that yields attr.type ==
cudaMemoryTypeHost classifies the pointer as host. In ptr_is_host
(kmeans_fit.cu) check if err == cudaSuccess and then return true only when
attr.type == cudaMemoryTypeHost (or equivalent field for host memory); if
cudaPointerGetAttributes returns an error, clear the sticky error with
cudaGetLastError() but do NOT assume host—return false or propagate the error so
the CUDA-path dispatch is not chosen incorrectly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3c495c39-bb68-4250-93e1-1c3e8cbe3b1e
📒 Files selected for processing (7)
cpp/include/cuml/cluster/kmeans.hppcpp/include/cuml/cluster/kmeans_params.hppcpp/src/kmeans/kmeans_fit.cucpp/src/kmeans/kmeans_params.cpppython/cuml/cuml/cluster/cpp/kmeans.pxdpython/cuml/cuml/cluster/kmeans.pyxpython/cuml/tests/test_kmeans.py
csadorf
left a comment
csadorf
left a comment
There was a problem hiding this comment.
Just a quick superficial early review.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/src/kmeans/kmeans_fit.cu (1)
38-50:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winCRITICAL:
sample_weightresidency is inferred fromXand can build the wrong RAFT view.Line 38 picks the host/device path from
Xalone, then Lines 47-50 and Lines 66-69 unconditionally wrapsample_weightas a host or device vector in that same branch. A hostXwith devicesample_weight(or the reverse) will hand cuVS the wrong view type, which can turn into invalid reads or a crash.As per coding guidelines, avoid host/device confusion and validate raw-pointer inputs before building RAFT views.
Also applies to: 66-69
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/kmeans/kmeans_fit.cu` around lines 38 - 50, The code incorrectly infers sample_weight residency from X and builds RAFT views for sample_weight using X's path; instead, check the residency of sample_weight independently (e.g. call ML::is_device_or_managed_type(sample_weight) or equivalent) and construct the matching RAFT view (raft::make_host_vector_view or raft::make_device_vector_view) for sample_weight only when sample_weight != nullptr; apply the same fix to the equivalent branch around the device-path code (the block creating sw and the one at lines 66-69), and add a short validation/error path when the pointer residency cannot be determined to avoid handing cuVS a mismatched view.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@cpp/src/kmeans/kmeans_fit.cu`:
- Around line 38-50: The code incorrectly infers sample_weight residency from X
and builds RAFT views for sample_weight using X's path; instead, check the
residency of sample_weight independently (e.g. call
ML::is_device_or_managed_type(sample_weight) or equivalent) and construct the
matching RAFT view (raft::make_host_vector_view or
raft::make_device_vector_view) for sample_weight only when sample_weight !=
nullptr; apply the same fix to the equivalent branch around the device-path code
(the block creating sw and the one at lines 66-69), and add a short
validation/error path when the pointer residency cannot be determined to avoid
handing cuVS a mismatched view.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 36ed9256-4d1a-4edf-8a38-b718fecd659e
📒 Files selected for processing (3)
cpp/include/cuml/cluster/kmeans.hppcpp/src/kmeans/kmeans_fit.cupython/cuml/cuml/cluster/kmeans.pyx
csadorf
changed the title
csadorf
left a comment
csadorf
left a comment
There was a problem hiding this comment.
Looks good generally, just a few minor change requests.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/src/kmeans/kmeans_fit.cu`:
- Around line 90-96: The new host/device dispatch is only implemented in the
float,int overload, while the double and int64_t overloads still reference the
removed fit_impl path and must be updated to call fit_impl_device or
fit_impl_host based on ML::is_device_or_managed_type(X). Also ensure the
Python-facing int64_t overloads are routed through the same host/device
branching so host-resident inputs do not bypass the new split; update the
affected kmeans_fit overloads to use the same dispatch logic consistently.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: fe1a74d1-5516-4478-b8fb-485550890c66
📒 Files selected for processing (3)
cpp/src/kmeans/kmeans_fit.cupython/cuml/cuml/cluster/kmeans.pyxpython/cuml/tests/test_kmeans.py
💤 Files with no reviewable changes (1)
- python/cuml/tests/test_kmeans.py
🚧 Files skipped from review as they are similar to previous changes (1)
- python/cuml/cuml/cluster/kmeans.pyx
| int& n_iter) | ||
| { | ||
| fit_impl(handle, params, X, n_samples, n_features, sample_weight, centroids, inertia, n_iter); | ||
| if (ML::is_device_or_managed_type(X)) { |
There was a problem hiding this comment.
is_device_or_managed_type calls on memory_type to determine the memory type. But, the memory_type function always returns att.type despite it being undefined when cudaErrorInvalidValue. The err == cudaErrorInvalidValue branch should probably return cudaMemoryTypeUnregistered.
There was a problem hiding this comment.
In that case, if it turns out to be unregistered (i.e. errors out, the host kmeans fit path would still be called here. How should we handle that?
There was a problem hiding this comment.
By unregistered I mean pageable host memory, not an invalid pointer. There is something a bit odd about this function. It seems to accept a fail from cudaPointerGetAttributes, but still returns the undefined att.type in this case.
The cases should be:
cudaSuccess+cudaMemoryTypeDeviceorManaged→ use device fit.cudaSuccess+cudaMemoryTypeHostorUnregistered→ use host fit.cudaErrorInvalidValue→att.typeis undefined and should be treated as an error. Unless explicitly supporting legacy CUDA (pre CUDA 11) where valid unregistered host memory produced this error. Up to you, but RAPIDS does not support CUDA 11 anymore.
There was a problem hiding this comment.
In that case, that function just needs to be modified to throw an error when we get cudaErrorInvalidValue
|
@tarang-jain Is this ready for another review iteration? I have the impression that some of the open threads have not yet been resolved. |
|
@csadorf ready for next round of reviews |
viclafargue
left a comment
viclafargue
left a comment
There was a problem hiding this comment.
Great work, thanks! Just some minor comments.
|
|
||
| common_kwargs = dict( | ||
| n_clusters=n_clusters, | ||
| init=init, |
There was a problem hiding this comment.
Here the initial centroids are systematically provided. Could we check that everything works fine too when using the default initialization of both paths?
| int& n_iter) | ||
| { | ||
| fit_impl(handle, params, X, n_samples, n_features, sample_weight, centroids, inertia, n_iter); | ||
| if (ML::is_device_or_managed_type(X)) { |
There was a problem hiding this comment.
By unregistered I mean pageable host memory, not an invalid pointer. There is something a bit odd about this function. It seems to accept a fail from cudaPointerGetAttributes, but still returns the undefined att.type in this case.
The cases should be:
cudaSuccess+cudaMemoryTypeDeviceorManaged→ use device fit.cudaSuccess+cudaMemoryTypeHostorUnregistered→ use host fit.cudaErrorInvalidValue→att.typeis undefined and should be treated as an error. Unless explicitly supporting legacy CUDA (pre CUDA 11) where valid unregistered host memory produced this error. Up to you, but RAPIDS does not support CUDA 11 anymore.
jcrist
left a comment
jcrist
left a comment
There was a problem hiding this comment.
A few small nits, but this mostly looks good to me!
|
@tarang-jain Please address the compilation errors (it appears to be an incompatibility with cuVS). Please avoid force-pushes if possible, it makes it harder to (re)review. Thanks! |
Co-authored-by: Victor Lafargue <viclafargue@nvidia.com>
|
the compilation errors are intentional. We are waiting for NVIDIA/cuvs#2328 to be merged before merging this PR, so that cuml does not break when that happens. |
jcrist
left a comment
jcrist
left a comment
There was a problem hiding this comment.
A few small comments, but this mostly LGTM!
csadorf
left a comment
csadorf
left a comment
There was a problem hiding this comment.
Thanks for working through the remaining feedback. The device_buffer_samples behavior and documentation are now consistent. The remaining follow-up work on the KMeans++ default behavior divergence can be addressed as part of NVIDIA/cuvs#2359. LGTM.
|
/merge |
NVIDIA/cuvs#2066 allows passing host partitions directly to cuVS KMeans (MNMG). Merge this PR after NVIDIA/cuvs#2066 and #8248 Closes #8201 Authors: - Tarang Jain (https://github.com/tarang-jain) Approvers: - Simon Adorf (https://github.com/csadorf) URL: #8084
5 participants
Adds
device_buffer_samplesto single-GPU KMeans for host-resident input data.For host-resident inputs, positive values select host-streamed fitting, avoiding an up-front full-device copy of
X. Device-resident inputs continue to use the standard device path and ignoredevice_buffer_samples. The PR wiresdevice_buffer_samplesandinit_sizethrough C++ and Python, rejects positivedevice_buffer_samplesvalues for multi-GPU use, computes labels and inertia in chunks for the host path, and adds weighted and unweighted parity tests against the existing device path.Depends on NVIDIA/cuvs#2328 (merged).
Follow-up: NVIDIA/cuvs#2359
Closes #8202