Don't import numba when dascore is imported - #853
Conversation
|
Warning Review limit reached
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 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 configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (7)
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. Comment |
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.
9a9c37a to
44db63c
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. 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
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Description
import dascorespends about 190 ms importing numba, even for users who never touch a jit'ed function.maybe_numba_jitimports numba and builds the dispatcher when it decorates, not when the decorated function is called.Patchbinds the transform functions as class attributes, sodascore.transform— and with ittaup.pyandkurtosis.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 theImportError, but not for the import cost.This moves the decorated kernels into private
_taup_kernels.pyand_kurtosis_kernels.pymodules, whichtau_p()andkurtosis()import inside the function body.dascore/utils/jit.pyis deliberately untouched, somaybe_numba_jitstill returns a real numbaDispatcherand kernels that call other kernels from nopython mode (_windowed_kurtosiscalls_moving_sum) keep working. The kernel bodies are a pure move.Median
import dascore, interleaved A/B over 12 runs each:devI first tried making
maybe_numba_jititself 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_funcbroken on any alias. Moving the kernels avoids all of that.Two notes for reviewers:
_jit_taup_uniform,_moving_sum, ...) are no longer importable fromdascore.transform.taup/dascore.transform.kurtosis. In-repo tests are updated. They are underscore-private, so no shims were added.PLC0415, the two deferred imports intau_pandkurtosiswill need# noqa: PLC0415. They are omitted here because the rule is not enabled ondevyet and ruff would flag them as unused directives. Whichever PR merges second needs that one-line follow-up.Changelog
import dascoreno longer imports numba (~190 ms), because the jit'edtau_pandkurtosiskernels move to private modules imported when those functions are called.Checklist
I have (if applicable):