Apply new validation to cuml.neighbors - #8019
Conversation
viclafargue
added
improvement
|
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughRefactors neighbors code to use centralized input validators and raw CuPy buffers for C++ interop (using ChangesKernelDensity & tests
KNN / NearestNeighbors refactor
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 the current code and only fix it if needed.
Inline comments:
In `@python/cuml/cuml/neighbors/nearest_neighbors.pyx`:
- Around line 598-609: In __setstate__, the code may pass a None X_cp (derived
from self._fit_X) into RBCIndex.build or ApproxIndex.build causing
AttributeError; before calling RBCIndex.build(...) or ApproxIndex.build(...),
check that X_cp is not None (i.e. self._fit_X was present) and either skip index
building or raise a clear ValueError/TypeError indicating the pickled state is
missing _fit_X; reference the symbols __setstate__, self._fit_X, X_cp,
RBCIndex.build, ApproxIndex.build, fit_method and effective_metric_ when adding
the guard and error message so the failure is informative and prevents calling
.data on None.
🪄 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: 2a4b80dd-ad6b-4635-b66c-0cd08589a77f
📒 Files selected for processing (6)
python/cuml/cuml/neighbors/kernel_density.pypython/cuml/cuml/neighbors/kneighbors_classifier.pyxpython/cuml/cuml/neighbors/kneighbors_regressor.pyxpython/cuml/cuml/neighbors/nearest_neighbors.pyxpython/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-list.yamlpython/cuml/tests/test_kernel_density.py
💤 Files with no reviewable changes (1)
- python/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-list.yaml
| fit_X = getattr(self, "_fit_X", None) | ||
| X_cp = fit_X.to_output("cupy") if fit_X is not None else None | ||
| if fit_method == "rbc": | ||
| self._index = RBCIndex.build(X, self.effective_metric_) | ||
| self._index = RBCIndex.build(X_cp, self.effective_metric_) | ||
| else: | ||
| self._index = ApproxIndex.build( | ||
| X, | ||
| X_cp, | ||
| self.effective_metric_, | ||
| fit_method, | ||
| params=self.algo_params, | ||
| p=self._effective_p, | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check if there are any guards or guarantees that _fit_X is never None when __setstate__ runs
rg -n -A5 '__setstate__|_fit_X.*None' python/cuml/cuml/neighbors/nearest_neighbors.pyxRepository: rapidsai/cuml
Length of output: 854
🏁 Script executed:
rg -n "class RBCIndex|class ApproxIndex" python/cuml/cuml/neighbors/Repository: rapidsai/cuml
Length of output: 208
🏁 Script executed:
rg -n "def build" python/cuml/cuml/neighbors/ -A3 | head -40Repository: rapidsai/cuml
Length of output: 771
🏁 Script executed:
# Look for where RBCIndex and ApproxIndex are imported/defined
rg -n "RBCIndex|ApproxIndex" python/cuml/cuml/neighbors/nearest_neighbors.pyx | head -20Repository: rapidsai/cuml
Length of output: 839
🏁 Script executed:
sed -n '278,310p' python/cuml/cuml/neighbors/nearest_neighbors.pyxRepository: rapidsai/cuml
Length of output: 1094
🏁 Script executed:
sed -n '417,455p' python/cuml/cuml/neighbors/nearest_neighbors.pyxRepository: rapidsai/cuml
Length of output: 1353
Add None check before building index in __setstate__.
If a pickle file is corrupted and missing _fit_X, the code sets X_cp = None but then passes it to RBCIndex.build() or ApproxIndex.build(), both of which immediately access X.data.ptr without a None guard, causing an AttributeError.
Check that X_cp is not None before calling either build method, or raise an informative error if the pickled state is invalid.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/cuml/cuml/neighbors/nearest_neighbors.pyx` around lines 598 - 609, In
__setstate__, the code may pass a None X_cp (derived from self._fit_X) into
RBCIndex.build or ApproxIndex.build causing AttributeError; before calling
RBCIndex.build(...) or ApproxIndex.build(...), check that X_cp is not None (i.e.
self._fit_X was present) and either skip index building or raise a clear
ValueError/TypeError indicating the pickled state is missing _fit_X; reference
the symbols __setstate__, self._fit_X, X_cp, RBCIndex.build, ApproxIndex.build,
fit_method and effective_metric_ when adding the guard and error message so the
failure is informative and prevents calling .data on None.
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
python/cuml/cuml/neighbors/nearest_neighbors.pyx (1)
598-605:⚠️ Potential issue | 🔴 CriticalGuard
__setstate__against missing_fit_Xbefore rebuilding indices.Line 599 can produce
X_cp = None, but Lines 601/603-604 still callRBCIndex.build/ApproxIndex.build, which dereferenceX.data.ptrand crash with a non-actionable error.💡 Proposed fix
if (fit_method := state.get("_fit_method")) in ("rbc", "ivfpq", "ivfflat"): # TODO: These index types currently aren't pickleable. For now we # recreate them on load. fit_X = getattr(self, "_fit_X", None) X_cp = cp.asarray(fit_X) if fit_X is not None else None + if X_cp is None: + raise ValueError( + "Invalid pickled NearestNeighbors state: missing `_fit_X` " + f"required to rebuild index for fit_method={fit_method!r}." + ) if fit_method == "rbc": self._index = RBCIndex.build(X_cp, self.effective_metric_) else: self._index = ApproxIndex.build( X_cp,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/cuml/cuml/neighbors/nearest_neighbors.pyx` around lines 598 - 605, In __setstate__, avoid calling RBCIndex.build or ApproxIndex.build when self._fit_X is missing or None: check getattr(self, "_fit_X", None) (or X_cp after cp.asarray) and only call RBCIndex.build or ApproxIndex.build with X_cp and self.effective_metric_ when X_cp is not None; otherwise set self._index = None (or skip rebuilding) so RBCIndex.build/ApproxIndex.build are never invoked on a None/empty array that would dereference X.data.ptr.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@python/cuml/cuml/neighbors/nearest_neighbors.pyx`:
- Around line 879-888: The code in nearest_neighbors.pyx recomputes distances as
squared L2 and returns them (variables distances, indices), which breaks
semantics for euclidean/l2/minkowski p=2; modify the return path to take the
square root of distances whenever the metric is "euclidean" or "l2" or when
metric == "minkowski" and p == 2 before converting to contiguous cp.float32
(i.e., replace returning squared distances with cp.sqrt(distances) in those
cases), ensuring you still call cp.ascontiguousarray on the final distances and
indices; locate the block that computes self_diff, distances, and correct_order
and add the conditional sqrt there.
- Around line 924-927: The sparse kneighbors indices are allocated as int32
(indices_cp) but _drop_self_edges invokes swap_kernel which expects long long
int*; to fix, allocate indices_cp as np.int64 (order="C") and change the C
pointer variable to the matching C type (e.g., cdef long long int* indices_ptr =
<long long int *><uintptr_t>indices_cp.data.ptr) so swap_kernel receives 64-bit
indices before any later conversion; update any related declarations that
reference indices_ptr (and any calls into swap_kernel) to use the new long long
pointer type to prevent the unsafe int32→int64 mismatch.
---
Duplicate comments:
In `@python/cuml/cuml/neighbors/nearest_neighbors.pyx`:
- Around line 598-605: In __setstate__, avoid calling RBCIndex.build or
ApproxIndex.build when self._fit_X is missing or None: check getattr(self,
"_fit_X", None) (or X_cp after cp.asarray) and only call RBCIndex.build or
ApproxIndex.build with X_cp and self.effective_metric_ when X_cp is not None;
otherwise set self._index = None (or skip rebuilding) so
RBCIndex.build/ApproxIndex.build are never invoked on a None/empty array that
would dereference X.data.ptr.
🪄 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: f13e8ffe-cc2b-4914-8be5-f4d05d23fc17
📒 Files selected for processing (2)
python/cuml/cuml/neighbors/nearest_neighbors.pyxpython/cuml/tests/test_sklearn_compatibility.py
💤 Files with no reviewable changes (1)
- python/cuml/tests/test_sklearn_compatibility.py
| self_diff = X[indices] - X[:, cp.newaxis, :] | ||
| distances = cp.sum(self_diff * self_diff, axis=2) | ||
| correct_order = cp.argsort(distances, axis=1) | ||
|
|
||
| self_diff = X_cp[indices_cp] - X_cp[:, cp.newaxis, :] | ||
| distances_cp = cp.sum(self_diff * self_diff, axis=2) | ||
| correct_order = cp.argsort(distances_cp, axis=1) | ||
| distances = cp.take_along_axis(distances, correct_order, axis=1) | ||
| indices = cp.take_along_axis(indices, correct_order, axis=1) | ||
|
|
||
| distances_cp = cp.take_along_axis(distances_cp, correct_order, axis=1) | ||
| indices_cp = cp.take_along_axis(indices_cp, correct_order, axis=1) | ||
|
|
||
| distances = CumlArray( | ||
| data=cp.ascontiguousarray(distances_cp, dtype=cp.float32), index=index | ||
| ) | ||
| indices = CumlArray( | ||
| data=cp.ascontiguousarray(indices_cp, dtype=cp.int64), index=index | ||
| ) | ||
| distances = cp.ascontiguousarray(distances, dtype=cp.float32) | ||
| indices = cp.ascontiguousarray(indices, dtype=cp.int64) | ||
| return distances, indices |
There was a problem hiding this comment.
Two-pass precision returns squared distances for euclidean/l2 metrics.
At Lines 879-888, distances are recomputed as squared L2 and returned directly. For euclidean/l2 (and minkowski with p=2), this changes output semantics from distance to squared distance.
💡 Proposed fix
- self_diff = X[indices] - X[:, cp.newaxis, :]
- distances = cp.sum(self_diff * self_diff, axis=2)
- correct_order = cp.argsort(distances, axis=1)
+ self_diff = X[indices] - X[:, cp.newaxis, :]
+ sq_distances = cp.sum(self_diff * self_diff, axis=2)
+ correct_order = cp.argsort(sq_distances, axis=1)
- distances = cp.take_along_axis(distances, correct_order, axis=1)
+ sq_distances = cp.take_along_axis(sq_distances, correct_order, axis=1)
indices = cp.take_along_axis(indices, correct_order, axis=1)
- distances = cp.ascontiguousarray(distances, dtype=cp.float32)
+ if metric == DistanceType.L2Expanded:
+ distances = sq_distances
+ else:
+ distances = cp.sqrt(sq_distances)
+ distances = cp.ascontiguousarray(distances, dtype=cp.float32)
indices = cp.ascontiguousarray(indices, dtype=cp.int64)
return distances, indices🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/cuml/cuml/neighbors/nearest_neighbors.pyx` around lines 879 - 888, The
code in nearest_neighbors.pyx recomputes distances as squared L2 and returns
them (variables distances, indices), which breaks semantics for
euclidean/l2/minkowski p=2; modify the return path to take the square root of
distances whenever the metric is "euclidean" or "l2" or when metric ==
"minkowski" and p == 2 before converting to contiguous cp.float32 (i.e., replace
returning squared distances with cp.sqrt(distances) in those cases), ensuring
you still call cp.ascontiguousarray on the final distances and indices; locate
the block that computes self_diff, distances, and correct_order and add the
conditional sqrt there.
| indices_cp = cp.empty((X_n_rows, n_neighbors), dtype=np.int32, order="C") | ||
| distances_cp = cp.empty((X_n_rows, n_neighbors), dtype=np.float32, order="C") | ||
| cdef int* indices_ptr = <int *><uintptr_t>indices_cp.data.ptr | ||
| cdef float* distances_ptr = <float *><uintptr_t>distances_cp.data.ptr |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify the type mismatch is present in this file:
# - swap_kernel expects long long int* (int64)
# - sparse kneighbors allocates int32 indices
# - _drop_self_edges is called for training-data queries
rg -n -C3 'swap_kernel\(long long int\* I|indices_cp = cp.empty\(\(X_n_rows, n_neighbors\), dtype=np.int32|_drop_self_edges\(' python/cuml/cuml/neighbors/nearest_neighbors.pyxRepository: rapidsai/cuml
Length of output: 1239
🏁 Script executed:
sed -n '198,230p' python/cuml/cuml/neighbors/nearest_neighbors.pyxRepository: rapidsai/cuml
Length of output: 1177
🏁 Script executed:
sed -n '900,935p' python/cuml/cuml/neighbors/nearest_neighbors.pyxRepository: rapidsai/cuml
Length of output: 1559
🏁 Script executed:
sed -n '860,930p' python/cuml/cuml/neighbors/nearest_neighbors.pyx | grep -E "def |cdef |return |_drop_self_edges"Repository: rapidsai/cuml
Length of output: 1510
🏁 Script executed:
sed -n '900,1000p' python/cuml/cuml/neighbors/nearest_neighbors.pyx | tail -100Repository: rapidsai/cuml
Length of output: 3711
🏁 Script executed:
sed -n '770,795p' python/cuml/cuml/neighbors/nearest_neighbors.pyxRepository: rapidsai/cuml
Length of output: 1111
Fix type mismatch in sparse kneighbors indices passed to CUDA kernel.
Lines 924-927 allocate sparse indices_cp as int32, but _drop_self_edges calls swap_kernel (declared with long long int* at line 172) while indices are still int32. The conversion to int64 at line 229 happens after swap_kernel runs, risking memory corruption when querying training data (use_training_data=True).
Proposed fix
- indices_cp = cp.empty((X_n_rows, n_neighbors), dtype=np.int32, order="C")
+ # C++ sparse knn writes int32, then upcast for downstream parity
+ indices_cp = cp.empty((X_n_rows, n_neighbors), dtype=np.int32, order="C")
distances_cp = cp.empty((X_n_rows, n_neighbors), dtype=np.float32, order="C")
cdef int* indices_ptr = <int *><uintptr_t>indices_cp.data.ptr
cdef float* distances_ptr = <float *><uintptr_t>distances_cp.data.ptr
@@
- return distances_cp, indices_cp
+ indices_cp = cp.ascontiguousarray(indices_cp, dtype=np.int64)
+ return distances_cp, indices_cp🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/cuml/cuml/neighbors/nearest_neighbors.pyx` around lines 924 - 927, The
sparse kneighbors indices are allocated as int32 (indices_cp) but
_drop_self_edges invokes swap_kernel which expects long long int*; to fix,
allocate indices_cp as np.int64 (order="C") and change the C pointer variable to
the matching C type (e.g., cdef long long int* indices_ptr = <long long int
*><uintptr_t>indices_cp.data.ptr) so swap_kernel receives 64-bit indices before
any later conversion; update any related declarations that reference indices_ptr
(and any calls into swap_kernel) to use the new long long pointer type to
prevent the unsafe int32→int64 mismatch.
jcrist
left a comment
jcrist
left a comment
There was a problem hiding this comment.
Thanks for working on this! Just a couple fixups, but nice to see all the xfails this removes!
| raise ValueError(f"kernel={self.kernel!r} is not supported") | ||
|
|
||
| self._X, n_rows, n_cols, _ = input_to_cupy_array( | ||
| check_features(self, X, reset=True) |
There was a problem hiding this comment.
You could use check_inputs here instead of check_features + check_array + check_sample_weight + check_consistent_length. A bit shorter code, and ensures things are handled consistently across our estimators.
There was a problem hiding this comment.
The check_inputs function does not have a sample_weight_ensure_non_negative argument making it impossible to merge everything into a single operation. But, I could merge check_features + check_array.
There was a problem hiding this comment.
You could also use check_inputs, and call check_non_negative on sample weights after too. Could also add the plumbing for ensure_sample_weight_non_negative if ya want. Up to you.
There was a problem hiding this comment.
Thanks for pointing out the existence of this function! Just updated the code with it.
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 `@python/cuml/cuml/neighbors/kernel_density.py`:
- Around line 303-332: The code currently assigns learned attributes (self._X,
self.bandwidth_, self._sample_weight) before all validations complete; instead,
call check_inputs(...) and compute bandwidth_ into local variables (e.g.,
X_valid, bandwidth_val), call check_sample_weight(...) and
check_consistent_length(X_valid, sample_weight_valid), and only after all
validations succeed assign self._X = X_valid, self.bandwidth_ = bandwidth_val,
and self._sample_weight = sample_weight_valid; reference the functions
check_inputs, check_sample_weight, check_consistent_length and the attributes
self._X, self.bandwidth_, self._sample_weight when making the change.
🪄 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: 2f32b8ab-b63a-4378-b437-e1a13cba5e68
📒 Files selected for processing (3)
python/cuml/cuml/neighbors/kernel_density.pypython/cuml/cuml/neighbors/kneighbors_classifier.pyxpython/cuml/cuml/neighbors/kneighbors_regressor.pyx
🚧 Files skipped from review as they are similar to previous changes (1)
- python/cuml/cuml/neighbors/kneighbors_regressor.pyx
| index=X.index, | ||
| ) | ||
| distances_cp = cp.empty((X.shape[0], n_neighbors), dtype=np.float32, order="C") | ||
| indices_cp = cp.empty((X.shape[0], n_neighbors), dtype=np.int64, order="C") |
There was a problem hiding this comment.
nit: adding a _cp suffix is noisy and shouldn't be necessary in our new world of only one array type. Not a blocker for now, but please don't add more of these.
|
/merge |
4 participants
Closes #8001