Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

IVF-PQ Python wrappers #970

Merged
merged 38 commits into from
Nov 15, 2022

Conversation

tfeher
Copy link
Contributor

@tfeher tfeher commented Nov 1, 2022

This PR adds python wrappers to IVF-PQ.

@tfeher tfeher force-pushed the fea_ivf_pq_python_wrapper branch from cf40569 to acdc3c4 Compare November 3, 2022 01:38
@tfeher tfeher marked this pull request as ready for review November 3, 2022 01:46
@tfeher tfeher requested review from a team as code owners November 3, 2022 01:46
@tfeher
Copy link
Contributor Author

tfeher commented Nov 3, 2022

There are a few items that still need to be addressed, but otherwise it is ready for the first review @cjnolet

  • handle codebook_kind parameter
  • handle lut_dtype and internal_distance_dtype search params
  • Add more test parameters combinations
  • Add xfail tests

@achirkin
Copy link
Contributor

achirkin commented Nov 7, 2022

What would you think about re-exporting default values of C++ params structs? Like we did a while ago in cuml/linearSVM (rapidsai/cuml#4381).
This would require a bit of extra work to make it nice, e.g. partially port the WithReexportedParams metaclass definition to allow inserting default values into python docstrings. But I think it still makes sense in the long run. I suspect we may want to change the default parameter values in future (at least those that affect performance, but don't affect the data layout/search results).

@tfeher tfeher added enhancement New feature or request non-breaking Non-breaking change labels Nov 7, 2022
@tfeher
Copy link
Contributor Author

tfeher commented Nov 7, 2022

@achirkin Thanks for the suggestion! yes it occurred to me to return to this idea. I would dedicate a follow up PR for this task.

@tfeher
Copy link
Contributor Author

tfeher commented Nov 7, 2022

There is one issue pending: forwarding the C++ exceptions from build and search to Python. Using the standard method leads to the following error message:

pylibraft/neighbors/ivf_pq.cxx:7155:45: error: no matching function for call to 'raft::neighbors::ivf_pq::index<long unsigned int>::index()'
 7155 |   raft::neighbors::ivf_pq::index<uint64_t>  __pyx_t_20;

This happens because Cython tries define a variable for the return value, and the return value type (index) does not have a default constructor. I will look into a workaround.

Copy link
Member

@cjnolet cjnolet left a comment

Choose a reason for hiding this comment

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

Thanks for adding this feature to pylibraft! I think the feature mostly looks ready to go, but I'd like to keep pylibraft's APIs consistent and make sure they all directly correspond to the c++ APIs they are wrapping.

python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
@cjnolet
Copy link
Member

cjnolet commented Nov 8, 2022

Closes #998

@cjnolet cjnolet added improvement Improvement / enhancement to an existing function and removed enhancement New feature or request labels Nov 10, 2022
Copy link
Member

@cjnolet cjnolet left a comment

Choose a reason for hiding this comment

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

This looks really good overall. I'm excited to have this in the python API!

cpp/CMakeLists.txt Outdated Show resolved Hide resolved
cpp/CMakeLists.txt Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
@tfeher tfeher force-pushed the fea_ivf_pq_python_wrapper branch from 2d6da94 to 1b5ff59 Compare November 11, 2022 15:52
Copy link
Contributor Author

@tfeher tfeher left a comment

Choose a reason for hiding this comment

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

Thanks @cjnolet for the review! I have addressed the issues. There is a problem with the extend test, I am looking into it.

python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
@tfeher
Copy link
Contributor Author

tfeher commented Nov 11, 2022

Fixed the problem with extend.

Copy link
Member

@cjnolet cjnolet left a comment

Choose a reason for hiding this comment

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

Looking great! I definitely think this is going to make it into 22.12, just really one little thing- hiding the cuda data types behind np.dtype for API consistency

python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
Copy link
Contributor Author

@tfeher tfeher left a comment

Choose a reason for hiding this comment

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

Thanks for the reviews, I have updated the PR!

python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
python/pylibraft/pylibraft/neighbors/ivf_pq.pyx Outdated Show resolved Hide resolved
Copy link
Member

@cjnolet cjnolet left a comment

Choose a reason for hiding this comment

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

LGTM. Thanks again @tfeher!

@cjnolet
Copy link
Member

cjnolet commented Nov 15, 2022

@gpucibot merge

@rapids-bot rapids-bot bot merged commit 355f693 into rapidsai:branch-22.12 Nov 15, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
CMake cpp improvement Improvement / enhancement to an existing function non-breaking Non-breaking change python
Projects
Development

Successfully merging this pull request may close these issues.

3 participants