Plumb metric and metric_kwds through to UMAP with nn_descent

[jcrist]

Previously we were erroneously missing this plumbing, leaving the metric and metric_arg as the defaults when nn_descent is used. In the long run we should redo how the parameters are passed when nn_descent is enabled to avoid this duplication (there's a few other uselessly exposed params like return_distances), but for now fixing the plumbing to be more correct seems fine.

On top of #6303, leaving as draft for now until that's merged.

Also re-enables tests on ARM, fixes #5441.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[cjnolet]

/ok to test

[jcrist]

Rebased on main and marked ready for review.

[jcrist]

The actual fix is here (plumbing through the metric options to nn_descent_params). In the long run we should redesign the C++ layer to remove duplicate options - for now just ensuring they're forwarded correctly seems sufficient.

Everything else here is a simplification of the current pre-existing code.

[viclafargue]

LGTM, but we would have to make sure that NN Descent is actually compatible with all the metrics we make available. cc @jinsolp

[jcrist]

A quick browse through the code I believe all metrics are handled fine by nn_descent. Just to be sure I've also updated an existing metrics test to be parametrized on build_algo as well; this test runs fine locally as well.

The behavior before was incorrect (build_algo="nn_descent" would silently ignore the user selected metric), so I consider this PR strictly an improvement even in the off chance some parameter combo will error after this change. I'm fairly confident this won't happen though.

Planning to merge once CI is green.

[jinsolp]

Right now NN Descent in cuvs supports L2Expanded, CosineExpanded and InnerProduct, (link)
but in Raft Nn Descent supports only L2Expanded (link)

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 67.34%. Comparing base (fbcaf12) to head (d275218).

Additional details and impacted files

@@               Coverage Diff                @@
##           branch-25.04    #6304      +/-   ##
================================================
- Coverage         68.50%   67.34%   -1.16%     
================================================
  Files               204      204              
  Lines             13194    13194              
================================================
- Hits               9039     8886     -153     
- Misses             4155     4308     +153     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[jcrist]

/merge (just realized victor's review wasn't an actual approval)

[jcrist]

but in Raft Nn Descent supports only L2Expanded

Ah! Thanks @jinsolp! I'm curious - right now after plumbing the metric through nn_descent properly here I'm not seeing an error (and the added test passes), but you're right that it doesn't look like the metrics are supported after all. I don't see any validation in the raft code flow that would error on an invalid metric - perhaps the right additional fix here then is:

• Add validation on the python side for nn_descent to error on unsupported metrics
• Keep plumbing the args through properly, even though they appear to be mostly ignored on the raft side for now
• Keep the test, but update it to check for errors on unsupported metrics

Then whenever we get around to moving to cuvs instead we can increase the supported metrics pretty straightforwardly.

[jinsolp]

Yes I think that would be great! : )