JP-4324: implement multiprocessing for Adaptive Trace Model

[mperrin]

Resolves JP-4324

This PR enhances the recently-added Adaptive Trace Model step by adding multiprocessing support, similar to that used in the jump and ramp fit steps. (i.e. parameter maximum_cores which can take integers, 'half', 'all', etc). Only the spline fitting part of the code becomes multi process, not the oversampling part (which is less computationally intensive in general).

Tasks

• If you have a specific reviewer in mind, tag them: @melanieclarke @drlaw1558
• add a build milestone, i.e. Build 12.0 (use the latest build if not sure)
• Does this PR change user-facing code / API? (if not, label with no-changelog-entry-needed)
  • write news fragment(s) in changes/: echo "changed something" > changes/<PR#>.<changetype>.rst (see changelog readme for instructions)
    • if your change breaks step-level or public API (as defined in the docs), also add a changes/<PR#>.breaking.rst news fragment
  • update or add relevant tests
  • update relevant docstrings and / or docs/ page
  • start a regression test and include a link to the running job (click here for instructions)
    • Do truth files need to be updated ("okified")?
      • after the reviewer has approved these changes, run okify_regtests to update the truth files
• if a JIRA ticket exists, make sure it is resolved properly

[melanieclarke]

This looks great! Thanks for contributing.

The docs build failure is unrelated; we should have a fix for it soon. I see you're iterating with ruff to get the style checks to pass, but see also the contributing docs for instructions if you want to get pre-commit set up to run locally.

We should also mention this step in the multiprocessing section of readthedocs, in docs/jwst/user_documentation/multiprocessing.rst.

These changes will need also some unit tests to exercise the new multiprocessing options. Do you want to attempt those yourself? If not, I can help.

[mperrin]

I'd appreciate your help with the tests, please, mostly since I'm not familiar with the test fixtures and other test infrastructure specific to the pipeline. I could probably figure it out based on your test code but I expect you will likely be faster at this part. Thanks!

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 85.78%. Comparing base (e9d87eb) to head (c90c35b).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main   #10391   +/-   ##
=======================================
  Coverage   85.77%   85.78%           
=======================================
  Files         372      372           
  Lines       40063    40081   +18     
=======================================
+ Hits        34364    34382   +18     
  Misses       5699     5699           

☔ 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.

[melanieclarke]

Testing locally, these changes look great.

Testing on my Mac with a NIRSpec IFU file with psf_optimal=True, processing time is improved from 7m to 1m with maximum_cores = "all", and memory use is not excessive.

I sent you a PR on your fork to add some tests and clean up a few software items here:
mperrin#1

[pllim]

• Ken had mention before that multiprocessing can swallow up some log messages. Is this a concern here?
• Given the description that multiprocessing is only even used under certain conditions, should there be extra checks to tell user that if they set >1 cores but really they don't have to?

[melanieclarke]

• Ken had mention before that multiprocessing can swallow up some log messages. Is this a concern here?

Not a concern here. There are no log messages internal to the code being multiprocessed.

• Given the description that multiprocessing is only even used under certain conditions, should there be extra checks to tell user that if they set >1 cores but really they don't have to?

It's not possible to tell from the input arguments if the multiprocessing will be needed or not. The decision is data-based. It does no harm to specify multiple cores anyway.

[emolter]

Two docs nitpicks and one question

[melanieclarke]

@mperrin - just checking in on this PR. There are a couple minor review comments to address. If you like, I can address them on your behalf so we can move forward.

[mperrin]

Hi @melanieclarke - thanks. I had meant to get to those, but have not had time yet due to other tasks and priorities. If you have spare bandwidth now then yes by all means you're welcome to go ahead. Thanks!

[melanieclarke]

Regression tests:
https://github.com/spacetelescope/RegressionTests/actions/runs/24531302622

All passing.