Skip to content

Don't import numba when dascore is imported - #853

Merged
d-chambers merged 1 commit into
devfrom
lazy-numba-jit
Aug 10, 2026
Merged

Don't import numba when dascore is imported#853
d-chambers merged 1 commit into
devfrom
lazy-numba-jit

Conversation

@d-chambers

@d-chambers d-chambers commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

Description

import dascore spends about 190 ms importing numba, even for users who never touch a jit'ed function.

maybe_numba_jit imports numba and builds the dispatcher when it decorates, not when the decorated function is called. Patch binds the transform functions as class attributes, so dascore.transform — and with it taup.py and kurtosis.py — is part of the eager import graph, and decorating their kernels drags numba in at import time. The decorator's docstring already promises that "users who don't use the function should be unaffected"; that held for the warning and the ImportError, but not for the import cost.

This moves the decorated kernels into private _taup_kernels.py and _kurtosis_kernels.py modules, which tau_p() and kurtosis() import inside the function body. dascore/utils/jit.py is deliberately untouched, so maybe_numba_jit still returns a real numba Dispatcher and kernels that call other kernels from nopython mode (_windowed_kurtosis calls _moving_sum) keep working. The kernel bodies are a pure move.

Median import dascore, interleaved A/B over 12 runs each:

median
dev 1082 ms
this branch 896 ms
saved 186 ms (~17%)

I first tried making maybe_numba_jit itself lazy. That breaks the nopython kernel-to-kernel call, because numba resolves globals at compile time and would find a python wrapper instead of a dispatcher; working around it meant republishing into module globals, which introduced a race and left .signatures/.py_func broken on any alias. Moving the kernels avoids all of that.

Two notes for reviewers:

  • The private kernel names (_jit_taup_uniform, _moving_sum, ...) are no longer importable from dascore.transform.taup / dascore.transform.kurtosis. In-repo tests are updated. They are underscore-private, so no shims were added.
  • This interacts with Enable PLC0415 so function-level imports need an explicit noqa #850. Once that lands and enables PLC0415, the two deferred imports in tau_p and kurtosis will need # noqa: PLC0415. They are omitted here because the rule is not enabled on dev yet and ruff would flag them as unused directives. Whichever PR merges second needs that one-line follow-up.

Changelog

  • changed: import dascore no longer imports numba (~190 ms), because the jit'ed tau_p and kurtosis kernels move to private modules imported when those functions are called.

Checklist

I have (if applicable):

  • referenced the GitHub issue this PR closes.
  • documented the new feature with docstrings and/or appropriate doc page.
  • included tests. See testing guidelines.
  • added the "ready_for_review" tag once the PR is ready to be reviewed.

@d-chambers d-chambers added the ready_for_review PR is ready for review label Aug 10, 2026
@coderabbitai

coderabbitai Bot commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

Warning

Review limit reached

@d-chambers, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 9 minutes

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: bd95a8f6-342b-4370-ae7e-76d4707120bb

📥 Commits

Reviewing files that changed from the base of the PR and between db17511 and 44db63c.

📒 Files selected for processing (7)
  • dascore/transform/_kurtosis_kernels.py
  • dascore/transform/_taup_kernels.py
  • dascore/transform/kurtosis.py
  • dascore/transform/taup.py
  • tests/test_imports.py
  • tests/test_transform/test_kurtosis.py
  • tests/test_transform/test_tau_p.py

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@d-chambers d-chambers added the benchmark Run the benchmark suite label Aug 10, 2026
maybe_numba_jit imports numba and builds the dispatcher when it decorates,
so every jit'ed kernel pulled numba in at import time. Patch binds the
transform functions as class attributes, which makes dascore.transform (and
therefore taup and kurtosis) part of the eager import graph, so importing
numba cost every user ~190 ms whether or not they ever called tau_p or
kurtosis.

Move the decorated kernels into private _taup_kernels and _kurtosis_kernels
modules and import them from inside tau_p and kurtosis instead. The kernels
are unchanged and maybe_numba_jit still returns a real numba dispatcher, so
kernels calling other kernels in nopython mode keep working.

Median `import dascore` drops from 1082 ms to 896 ms.
@codspeed-hq

codspeed-hq Bot commented Aug 10, 2026

Copy link
Copy Markdown

Merging this PR will not alter performance

✅ 79 untouched benchmarks


Comparing lazy-numba-jit (44db63c) with dev (db17511)

Open in CodSpeed

@codecov

codecov Bot commented Aug 10, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 100.00%. Comparing base (db17511) to head (44db63c).

Additional details and impacted files
@@            Coverage Diff            @@
##               dev      #853   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files          170       172    +2     
  Lines        18374     18380    +6     
=========================================
+ Hits         18374     18380    +6     
Flag Coverage Δ
network 48.86% <16.00%> (+<0.01%) ⬆️
unittests 100.00% <100.00%> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 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.

@d-chambers
d-chambers merged commit 7865baa into dev Aug 10, 2026
28 checks passed
@d-chambers
d-chambers deleted the lazy-numba-jit branch August 10, 2026 13:12
@d-chambers d-chambers removed the ready_for_review PR is ready for review label Aug 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

benchmark Run the benchmark suite

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant