Skip to content

[ENH] Distance module n_jobs support#2545

Merged
chrisholder merged 217 commits intomainfrom
distance-module-n-jobs
Aug 7, 2025
Merged

[ENH] Distance module n_jobs support#2545
chrisholder merged 217 commits intomainfrom
distance-module-n-jobs

Conversation

@chrisholder
Copy link
Copy Markdown
Contributor

Reference Issues/PRs

#1886 #1797

What does this implement/fix? Explain your changes.

This PR adds parallelism to the distance module. Specifically to pairwise distance computations.

To do this the n_job parameter can now be specified to any pairwise distance (aside mpdist which was excluded due to complexity) and it will use that many threads.

This PR just updates the distance module. In a follow up PR I will update all our models that use pairwise_distance and n_jobs to make use of this.

Does your contribution introduce a new dependency? If yes, which one?

Any other comments?

PR checklist

For all contributions
  • I've added myself to the list of contributors. Alternatively, you can use the @all-contributors bot to do this for you after the PR has been merged.
  • The PR title starts with either [ENH], [MNT], [DOC], [BUG], [REF], [DEP] or [GOV] indicating whether the PR topic is related to enhancement, maintenance, documentation, bugs, refactoring, deprecation or governance.
For new estimators and functions
  • I've added the estimator/function to the online API documentation.
  • (OPTIONAL) I've added myself as a __maintainer__ at the top of relevant files and want to be contacted regarding its maintenance. Unmaintained files may be removed. This is for the full file, and you should not add yourself if you are just making minor changes or do not want to help maintain its contents.
For developers with write access
  • (OPTIONAL) I've updated aeon's CODEOWNERS to receive notifications about future changes to these files.

@aeon-actions-bot aeon-actions-bot bot added distances Distances package enhancement New feature, improvement request or other non-bug code enhancement labels Feb 18, 2025
@aeon-actions-bot
Copy link
Copy Markdown
Contributor

aeon-actions-bot bot commented Feb 18, 2025

Thank you for contributing to aeon

I have added the following labels to this PR based on the title: [ $\color{#FEF1BE}{\textsf{enhancement}}$ ].
I have added the following labels to this PR based on the changes made: [ $\color{#5209C9}{\textsf{distances}}$ ]. Feel free to change these if they do not properly represent the PR.

The Checks tab will show the status of our automated tests. You can click on individual test runs in the tab or "Details" in the panel below to see more information if there is a failure.

If our pre-commit code quality check fails, any trivial fixes will automatically be pushed to your PR unless it is a draft.

Don't hesitate to ask questions on the aeon Slack channel if you have any.

PR CI actions

These checkboxes will add labels to enable/disable CI functionality for this PR. This may not take effect immediately, and a new commit may be required to run the new configuration.

  • Run pre-commit checks for all files
  • Run mypy typecheck tests
  • Run all pytest tests and configurations
  • Run all notebook example tests
  • Run numba-disabled codecov tests
  • Stop automatic pre-commit fixes (always disabled for drafts)
  • Disable numba cache loading
  • Push an empty commit to re-run CI checks

Copy link
Copy Markdown
Member

@baraline baraline left a comment

Choose a reason for hiding this comment

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

LGTM! The only thing that would be missing from the PR to fix the issue is to include n_jobs in the distance params in our KNeighborsTimeSeriesClassifier and we would be good to go

baraline
baraline previously approved these changes Feb 18, 2025
Copy link
Copy Markdown
Member

@MatthewMiddlehurst MatthewMiddlehurst left a comment

Choose a reason for hiding this comment

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

Would be nice to benchmark and see if this actually speeds it up.

@baraline
Copy link
Copy Markdown
Member

baraline commented Mar 1, 2025

Can confirm that this does indeed provide some speed ups :

import numpy as np
from aeon.distances import euclidean_pairwise_distance

X = np.random.rand(1000, 1000)
Y = np.random.rand(1000, 1000)
euclidean_pairwise_distance(X, Y)

%timeit euclidean_pairwise_distance(X, Y)
%timeit euclidean_pairwise_distance(X, Y, n_jobs=-1)
199 ms ± 7.05 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
82.7 ms ± 1.19 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Seconds Matthew's comment to revert the number of threads back to previously setted value (using get_num_threads).

@chrisholder
Copy link
Copy Markdown
Contributor Author

With distance measures like Euclidean I found unless you have a huge number of very long time series threading is almost never going to be faster just due to the cost of spawning a thread. This is why I added the warning as I can't really see a scenario where running pairwise euclidean distance thread will ever achieve a meaningful speed up.

However, for something like DTW I get huge speed ups:

    X = make_example_3d_numpy(10000, 1, 100, return_y=False, random_state=42)
    start = time.time()
    pairwise_distance(X, method="dtw", n_jobs=1)
    print(f"1 Job Time: {time.time() - start}")

    start = time.time()
    pairwise_distance(X, method="dtw", n_jobs=14)
    print(f"14 Job Time: {time.time() - start}")
Time: 1620.196809053421
Time: 255.28647780418396

@chrisholder
Copy link
Copy Markdown
Contributor Author

I also decided to add the threading with a custom distance function.

@chrisholder
Copy link
Copy Markdown
Contributor Author

I followed your suggestion @MatthewMiddlehurst and defined a decorator to handle setting and resetting threading.

@chrisholder chrisholder added the no numba cache Disable numba cache loading on a PR label Mar 4, 2025
@chrisholder chrisholder added the no numba cache Disable numba cache loading on a PR label Jul 31, 2025
@aeon-actions-bot aeon-actions-bot bot removed the no numba cache Disable numba cache loading on a PR label Jul 31, 2025
Copy link
Copy Markdown
Contributor

@TonyBagnall TonyBagnall left a comment

Choose a reason for hiding this comment

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

LGTM but will wait to see if @baraline or @MatthewMiddlehurst want to comment

baraline
baraline previously approved these changes Aug 2, 2025
Copy link
Copy Markdown
Member

@baraline baraline left a comment

Choose a reason for hiding this comment

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

The use of the decorator and how you handle number of threads seems to be the good solution here, so I'm good with it, makes the handling of the n_jobs straightforward for those functions without having to mess with them.

TonyBagnall
TonyBagnall previously approved these changes Aug 4, 2025
Copy link
Copy Markdown
Contributor

@TonyBagnall TonyBagnall left a comment

Choose a reason for hiding this comment

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

lgtm

@chrisholder chrisholder dismissed stale reviews from TonyBagnall and baraline via 3d56e6c August 5, 2025 17:18
@chrisholder chrisholder added the no numba cache Disable numba cache loading on a PR label Aug 5, 2025
@chrisholder
Copy link
Copy Markdown
Contributor Author

chrisholder commented Aug 5, 2025

I've had to disable the CanonicalIntervalForestRegressor "check_regressor_against_expected_results" test. I believe this has started to fail due to instability with numba in the module. I'm not really sure at all how this PR is impacting this test my best guess is switching the number of threads is causing unstable behaviour in tests that occur after certain distance tests run.

I have also noticed that two other regressors RDST and RIST Regressor also have excluded tests in the aeon/testing/testing_config.py EXCLUDED_TESTS so I would assume this is a similar issue. I have had this problem both with numba caches on and off. I think this is a deeper issue with the regression module itself.

I could change the expected results as it seems to consistently now resolve to 0.14 instead of 0.16 (fourth result in the expected array) however, I don't feel comfortable changing this not knowing if this is a bug or not.

I don't think I can fix this myself so instead after this PR is merged I'll open an issue documenting the bug so we can resolve it in the future. Here is a CI run: https://github.com/aeon-toolkit/aeon/actions/runs/16759105314/job/47458932638?pr=2545

FAILED aeon/testing/tests/test_all_estimators.py::test_all_estimators[check_regressor_against_expected_results(estimator_class=CanonicalIntervalForestRegressor,data_name=CardanoSentiment,data_loader=load_cardano_sentiment)] - AssertionError: 
Arrays are not almost equal to 2 decimals
Failed to reproduce results for CanonicalIntervalForestRegressor on CardanoSentiment
Mismatched elements: 1 / 10 (10%)
Max absolute difference among violations: 0.02247
Max relative difference among violations: 0.1605
 ACTUAL: array([0.28, 0.15, 0.33, 0.16, 0.19, 0.22, 0.15, 0.22, 0.13, 0.2 ])
 DESIRED: array([0.28, 0.15, 0.33, 0.14, 0.19, 0.22, 0.15, 0.23, 0.12, 0.2 ])
= 1 failed, 5841 passed, 64 skipped, 25 xfailed, 12 xpassed, 44707 warnings in 902.86s (0:15:02) =

Copy link
Copy Markdown
Member

@MatthewMiddlehurst MatthewMiddlehurst left a comment

Choose a reason for hiding this comment

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

in testing config we have a variable called MULTITHREAD_TESTING, can you use that for tests which play with threads, It runs it on a workflow without pytest-xdist (the tool we use to thread testing)

@chrisholder
Copy link
Copy Markdown
Contributor Author

Ill make that change now

Copy link
Copy Markdown
Member

@MatthewMiddlehurst MatthewMiddlehurst 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.

@chrisholder chrisholder merged commit 2255f86 into main Aug 7, 2025
31 of 32 checks passed
@chrisholder chrisholder deleted the distance-module-n-jobs branch August 7, 2025 14:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

distances Distances package enhancement New feature, improvement request or other non-bug code enhancement no numba cache Disable numba cache loading on a PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.