[REVIEW] New Dataset API Clarifying Ownership#1846
Conversation
|
/ok to test 5447a4c |
|
/ok to test 17ab09d |
|
NB: I updated the label to |
@achirkin The problem w/ using mdspan/mdarray for this is that it's not carrying along the proper information to either the algorithms nor the user (which is why we created this specialized class for this in the first place!). Two immediate reasons why this API is necessary:
This new API solves both of these problems while leaving the control over the memory ownership entirely in the user's hands. We've discussed this for a long time. We've known this is needed for a long time. it's time to prioritize this and get it done. I agree that an anstract class might make more sense, but ultimately we should not be moving any owneship over to the algorithm (the user should maintain ownership over the class and underlying memory the entire time). |
|
The doc that outlines some of the API design choices can be found in slack. Let me know if there are any parts of the design that can be altered to better suit our users' needs. The following files are test case files I've added and can be ignored for now. They will be removed before the final merge with upstream repo:
|
|
/ok to test 70b6b58 |
achirkin
left a comment
There was a problem hiding this comment.
Sorry for being so late for the second round of reviews! I have a few suggestions for the dataset type hierarchy.
| /** Compressed dataset. */ | ||
| raft::device_matrix<uint8_t, index_type, raft::row_major> data; | ||
|
|
||
| vpq_dataset(raft::device_matrix<math_type, uint32_t, raft::row_major>&& vq_code_book, |
There was a problem hiding this comment.
This is not cagra specific. Please move this to the respective preprocessing/pq.hpp file. Each different quantizer sohuld provide their respective dataset implementation so that when imported, the dataset comes w/ the quantizer.
| * - IdxT: index type used for row counts (`n_rows()` return type). | ||
| */ | ||
| template <typename MatrixT, typename ViewT, typename DataT, typename IdxT> | ||
| struct dense_row_major_dataset_owning_storage { |
There was a problem hiding this comment.
Agreed- for a single dense contiguous owning matrix type, we should be using mdarray... nothing else. If we need some type of composite, we should define that upon use... but i'm not seeing an immediate need for that in this PR.
| cuvsDatasetPaddedView_t padded_dataset, | ||
| cuvsCagraIndex_t index); | ||
|
|
||
| /** |
There was a problem hiding this comment.
What is a host cagra index? CAGRA is a GPU algorithn. This naming could be described with more clarity. Does "host cagra index" just imply the underlying dataset is on host? Because there is a potential configuration for CAGRA where the graph is always on host too.
Also, in C, we tend to put the verb last- "Make" "Create" "Destroy" "Attach" should all be last.
I would
There was a problem hiding this comment.
Yes host cagra index just refers to an index where the attached dataset is on host but the index graph itself is on device. Such a dataset is not immediately searchable.
…ng 26.10 to 26.08 versions for release
… not correct. Use owning factory first and then create view from that.
Authors: - Zihao Wang (https://github.com/hhy3) - Anupam (https://github.com/aamijar) Approvers: - Tarang Jain (https://github.com/tarang-jain) URL: NVIDIA#2322
This PR removes the `libclang==20.1.8` pin from the `clang` dependency group in `dependencies.yaml`. ## Problem The explicit `libclang==20.1.8` pin causes the conda solver to hang when creating the combined RAPIDS devcontainer environment. Specifically, when cuvs is included alongside cudf (which depends on `matplotlib<3.11`), the solver cannot efficiently resolve the conflict between: - `libclang==20.1.8` → provides `libclang13 20.1.8` - `matplotlib<3.11` → pulls in `qt6-main` → requires `libclang13>=22.1.0` While the environment is technically solvable (older Qt builds accept `libclang13>=20.1.8`), the search space created by the 29 builds of `libclang 20.1.8` × multiple Qt versions × CUDA packages × rapidsai channel packages causes the solver to hang indefinitely. ## Why this is safe to remove 1. **`clang` and `clang-tools` do not depend on `libclang`.** The `clang` compiler depends on `libclang-cpp` (the C++ API), which is a separate package from `libclang` (the C API wrapper). 2. **The Rust bindgen workflow doesn't use the `clang` dependency group.** The `rust` file key does not include the `clang` group — bindgen finds `libclang.so` at runtime from whatever version is available in the environment (provided transitively via Qt/matplotlib's dependency on `libclang13`). 3. **Other RAPIDS repos don't pin `libclang`.** cudf, raft, and rmm only pin `clang` and `clang-tools` without `libclang` and do not experience this issue. 4. **Historical context:** `libclang` was originally added to the `rust` dependency group for bindgen support, then moved into the `clang` group during a version bump refactor. The `rust` file key no longer includes the `clang` group, making this entry vestigial. Authors: - Vyas Ramasubramani (https://github.com/vyasr) Approvers: - Gil Forsyth (https://github.com/gforsyth) URL: NVIDIA#2263
The sample code for cagra::build_knn_graph was using deprecated APIs. This PR updates it to use the current API. Authors: - tsuki (https://github.com/enp1s0) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: NVIDIA#2339
Forward-merge release/26.08 into main
…taset() function. Attach_dataset() takes in any index type as input (host vs device, standard vs padded) and a device padded device and always returns a device padded index which is searchable. It internally dispatches on every combination of input index and input dataset type allowed and does internal host to device or standard to padded conversion if needed. We introduce 2 helper functions: convert_host_to_device_index() and convert_standard_to_padded_index(). These 2 helpers along with the original update_dataset(), which only takes device_padded_index and device_padded_dataset() as input, are internally called within attach_dataset() within each of the dispatch cases.
| extern "C" cuvsError_t cuvsDatasetMakeDevicePadded(cuvsResources_t res, | ||
| DLManagedTensor* dataset_tensor, | ||
| cuvsDatasetPadded_t* padded_dataset) |
There was a problem hiding this comment.
All those new cuvsDataset... functions should be in a separate compilation unit (src/core/dataset.cpp) (Not blocking this PR though, it can be a follow-up)
| DLDataType dtype; | ||
| cuvsDatasetLayout_t layout; | ||
| } cuvsDatasetStandard; | ||
| typedef cuvsDatasetStandard* cuvsDatasetStandard_t; |
There was a problem hiding this comment.
The two structs are similar, can't there be just one struct definition and two typedefs? They already represent the same thing.
There was a problem hiding this comment.
Same with the view struct just below
| void (*destroy_addr)(void*); | ||
| cuvsDatasetViewKind_t kind; | ||
| DLDataType dtype; | ||
| cuvsDatasetLayout_t layout; |
There was a problem hiding this comment.
layout is redundant if kind is already present in the struct
| cuvsCagraIndexParams_t params, | ||
| DLManagedTensor* dataset, | ||
| cuvsCagraIndex_t index); | ||
| CUVS_EXPORT cuvsError_t cuvsCagraGetDatasetViewKind(DLManagedTensor* dataset, |
There was a problem hiding this comment.
Any reason this needs to be exposed to the public?
| CUVS_EXPORT cuvsError_t cuvsCagraDeserializePadded(cuvsResources_t res, | ||
| const char* filename, | ||
| cuvsCagraIndex_t index, | ||
| cuvsDatasetPadded_t* out_padded_dataset); |
There was a problem hiding this comment.
Question- why isn't the dataset type serialized with the cagra graph (in the even the user wants that serialized with the cagra graph)? This way we can just do "cuvsCagraDeserializeWithDataset" (which throws an error if the serialized file doesn't contain the dataset) and cuvsCagraDeserialize (which doesn't deserialize the dataset, even if it's in the file).
There was a problem hiding this comment.
Pre-allocating the dataset buffer gives us 2 advantages: (1) no hidden memory allocations, (2) deterministic dataset type.
If we are pre-allocating the dataset buffer, I think we can just parse the type from that directly so there wouldn't be a need to modify the serialize algo itself to serialize the dataset type along with the graph.
If the proposal is to remove the passing in of the dataset buffer entirely in favor of serializing the dataset type with the cagra graph (which more closely mirrors the original serialize function), we would lose the beenfit of advantage (1).
| typedef enum { | ||
| CUVS_DATASET_VIEW_KIND_DEVICE_PADDED = 0, | ||
| CUVS_DATASET_VIEW_KIND_HOST_PADDED = 1, | ||
| CUVS_DATASET_VIEW_KIND_DEVICE_STANDARD = 2, |
There was a problem hiding this comment.
Instead of view kind containing both host and device for each dataset type (this is going to quickly become unmaintainable), please just define the accessor type (HOST and DEVICE) here and reuse the dataset layout enum above.
| DLDataType dtype; | ||
| cuvsDatasetLayout_t layout; | ||
| } cuvsDatasetPadded; | ||
| typedef cuvsDatasetPadded* cuvsDatasetPadded_t; |
There was a problem hiding this comment.
We discussed this offline- we should probably just make a single dataset struct and use the enums to parameterize it so that we know what kind to cast it to in the C++ layer.
… NEIGHBORS_ANN_CAGRA_tests
…update_dataset() that only updated standard index with standard dataset or padded index with padded dataset() was renamed update_device_dataset_same_layout(). New centralized cuvsCagraUpdateDataset at C API layer was created which internally dispatches to cuvsCagraUpdateDeviceDatasetSameLayout if index type is device_padded_index. For all other index types, it dispatches internally to cuvsCagraAttachDataset. This resolves the mutability issue we were previously experiencing where cuvsCagraAttachDataset guarantees original input index is immutable and returns a new index of new type while cuvsCagraUpdateDeviceDatasetSameLayout requires mutating the input dataset and makes no such index copy. This is not implemented at the C++ API layer yet, it only exists in the C API layer since C API doesn't make immutability guarantees while C++ API does and tries to keep C++ logic immutable to prevent thread conflicts. Downstream langauge wrappers have been updated to call centralized cuvsCagraUpdateDataset as the single sole entry point for all cases of updating/attaching a dataset
…ut from exports in C API layer and remove from downstream wrappers. They are not meant to be exposed to users/callers. Users/callers should use cuvsCagraUpdateDataset as the only C API function they interface with when updating/attaching a dataset
Forward-merge release/26.08 into main
| _deserialize_padded<float>(res, filename, index, out_padded_dataset); | ||
| index->dtype.code = kDLFloat; | ||
| } else if (dtype.kind == 'e' && dtype.itemsize == 2) { | ||
| _deserialize_padded<half>(res, filename, index, out_padded_dataset); | ||
| index->dtype.code = kDLFloat; | ||
| } else if (dtype.kind == 'i' && dtype.itemsize == 1) { | ||
| _deserialize_padded<int8_t>(res, filename, index, out_padded_dataset); | ||
| index->dtype.code = kDLInt; | ||
| } else if (dtype.kind == 'u' && dtype.itemsize == 1) { | ||
| _deserialize_padded<uint8_t>(res, filename, index, out_padded_dataset); | ||
| index->dtype.code = kDLUInt; |
There was a problem hiding this comment.
I think there's a bug here in the order of these ops. The _deserialize_padded copies the index's dtype into the handle it returns:
out->dtype = output_index->dtype;So the dtype should already be set when it's called. Therefore, we should flip the two lines for each dtype kind branch:
index->dtype.code = ...
_deserialize_padded...Forward-merge release/26.08 into main
…or MG path. Factored out distribute_padded_dataset() shared GPU sharding logic and exposed both attach_dataset() and update_device_dataset_same_layout() C++ API functions for the MG path. Expose new MG centralized update dataset in python wrapper. Other downstream languages java, rust, and go for some reason don't have any CAGRA MG functions exposed. Will need to expose this in a future commit.
|
/ok to test d377781 |
Overview
Addressing #1574 and #1571.
Replaced strided_dataset with padded_dataset class. Added support all the way up to CAGRA code.
Old class structure (Classes + Inheritance):
New Class Structure (ContainerType Tags + Composition):
Inheritance is removed entirely and all dataset types are on the same level of the inheritance tree.
3 Levels:
Ownership
The index and cagra::build / cagra::index do not own raw vector storage, they only take views.
The old code had a type-erased std::unique_ptr<dataset_view<...>>, i.e. non-owning view handles. The new code uses templates on the index type which determines the type of dataset_view the index holds.
ACE v.s. non-ACE paths on Host
ACE path copies datasets that can't entirely fit in CPU memory in chunks onto GPU memory by calling make_padded_dataset. This is 1x memory on CPU and 1x memory on GPU.
Return types:
Used mainly to maintain lifetime of dataset.
cuvs_cagra_c_api_lifetime_holder
It is a single C++ struct in cagra.cpp that groups the real cagra::index with any extra heap-owned things the C API had to create so the index’s non-owning views stay valid.
Miscellaneous: Extend Serialize Deserialize
Will fill in later
Factories:
Places where make_padded_dataset/view are called internally (not by user):
Host non-ACE path
Tiered CAGRA
Ownership in Downstream Functions:
Improvements:
Breaking Changes for Dataset API:
The following functions are removed since index no longer owns the dataset, index only takes views:
Removed old functions that took mdspan or derivatives of mdspan.
4 cases where index previously owned dataset [all deprecated paths]:
2 edge case build() paths when attach_dataset_on_build == true and a successful dense attach:
Compression Param:
Merge:
These paths have since been removed.
Attach Dataset
Compressed Dataset
Merged Dataset
Deserialize
Helpers
How to attach a compressed dataset onto an uncompressed index?
How to attach a searchable device dataset onto an index built with host build?
a. Utilizes map of host dataset type to device dataset type counterpart
TODOs:
Recent Updates:
Future PRs:
PR#2: Add Support for Compressed Datasets
PR#3: Migrate Rest of Algorithms to use Dataset API