Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

JP-3721: Simplify ModelContainer #8831

Merged
merged 117 commits into from
Oct 22, 2024
Merged

JP-3721: Simplify ModelContainer #8831

merged 117 commits into from
Oct 22, 2024

Conversation

emolter
Copy link
Collaborator

@emolter emolter commented Sep 26, 2024

Resolves JP-3721

Closes #8738

With the addition of ModelLibrary, ModelContainer can be substantially simplified. ModelLibrary certainly renders the return_open and save_open options obsolete, as well as the get_sections method, since if memory usage is a concern then ModelLibrary should be used instead. Some additional discussion about how the ModelContainer used to perform way to many tasks can be found on this innerspace page.

During discussions related to JP-3715, it has become increasingly clear that the strictness of ModelLibrary and the additional borrow/shelve code required to access models is not necessary/desired for many use-cases. For example, the calwebb_spec3 pipeline currently makes extensive use of ModelContainer but does not have the same memory issues as calwebb_image3 does. Learning to use ModelLibrary may also be a nuisance for users manipulating relatively small datasets.

Therefore, ModelContainer should not be removed, but instead become a lightweight, easy-to-use class for loading a list of models and association metadata. This PR proposed a container satisfying the following constraints:

  • Loads association metadata dictionary
  • Can be manipulated as if it were a list
  • Loads all datamodels into memory (i.e., use ModelLibrary instead if we care about memory usage)
  • Is the default data structure loaded by datamodels.open() for asn-type data but is no longer itself a datamodel
  • Can be used as a context manager, i.e., with datamodels.open(asn.json) as container has the expected behavior

Linked PR in stdatamodels: spacetelescope/stdatamodels#330
Linked PR in stpipe: spacetelescope/stpipe#190

Tasks

  • request a review from someone specific, to avoid making the maintainers review every PR
  • add a build milestone, i.e. Build 11.3 (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 below for change types)
    • 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
news fragment change types...
  • changes/<PR#>.general.rst: infrastructure or miscellaneous change
  • changes/<PR#>.docs.rst
  • changes/<PR#>.stpipe.rst
  • changes/<PR#>.datamodels.rst
  • changes/<PR#>.scripts.rst
  • changes/<PR#>.fits_generator.rst
  • changes/<PR#>.set_telescope_pointing.rst
  • changes/<PR#>.pipeline.rst

stage 1

  • changes/<PR#>.group_scale.rst
  • changes/<PR#>.dq_init.rst
  • changes/<PR#>.emicorr.rst
  • changes/<PR#>.saturation.rst
  • changes/<PR#>.ipc.rst
  • changes/<PR#>.firstframe.rst
  • changes/<PR#>.lastframe.rst
  • changes/<PR#>.reset.rst
  • changes/<PR#>.superbias.rst
  • changes/<PR#>.refpix.rst
  • changes/<PR#>.linearity.rst
  • changes/<PR#>.rscd.rst
  • changes/<PR#>.persistence.rst
  • changes/<PR#>.dark_current.rst
  • changes/<PR#>.charge_migration.rst
  • changes/<PR#>.jump.rst
  • changes/<PR#>.clean_flicker_noise.rst
  • changes/<PR#>.ramp_fitting.rst
  • changes/<PR#>.gain_scale.rst

stage 2

  • changes/<PR#>.assign_wcs.rst
  • changes/<PR#>.badpix_selfcal.rst
  • changes/<PR#>.msaflagopen.rst
  • changes/<PR#>.nsclean.rst
  • changes/<PR#>.imprint.rst
  • changes/<PR#>.background.rst
  • changes/<PR#>.extract_2d.rst
  • changes/<PR#>.master_background.rst
  • changes/<PR#>.wavecorr.rst
  • changes/<PR#>.srctype.rst
  • changes/<PR#>.straylight.rst
  • changes/<PR#>.wfss_contam.rst
  • changes/<PR#>.flatfield.rst
  • changes/<PR#>.fringe.rst
  • changes/<PR#>.pathloss.rst
  • changes/<PR#>.barshadow.rst
  • changes/<PR#>.photom.rst
  • changes/<PR#>.pixel_replace.rst
  • changes/<PR#>.resample_spec.rst
  • changes/<PR#>.residual_fringe.rst
  • changes/<PR#>.cube_build.rst
  • changes/<PR#>.extract_1d.rst
  • changes/<PR#>.resample.rst

stage 3

  • changes/<PR#>.assign_mtwcs.rst
  • changes/<PR#>.mrs_imatch.rst
  • changes/<PR#>.tweakreg.rst
  • changes/<PR#>.skymatch.rst
  • changes/<PR#>.exp_to_source.rst
  • changes/<PR#>.outlier_detection.rst
  • changes/<PR#>.tso_photometry.rst
  • changes/<PR#>.stack_refs.rst
  • changes/<PR#>.align_refs.rst
  • changes/<PR#>.klip.rst
  • changes/<PR#>.spectral_leak.rst
  • changes/<PR#>.source_catalog.rst
  • changes/<PR#>.combine_1d.rst
  • changes/<PR#>.ami.rst

other

  • changes/<PR#>.wfs_combine.rst
  • changes/<PR#>.white_light.rst
  • changes/<PR#>.cube_skymatch.rst
  • changes/<PR#>.engdb_tools.rst
  • changes/<PR#>.guider_cds.rst

braingram and others added 30 commits July 29, 2024 16:07
Copy link
Collaborator

@melanieclarke melanieclarke left a comment

Choose a reason for hiding this comment

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

A few comments below, relating to the default structure for the asn_table.

jwst/datamodels/container.py Outdated Show resolved Hide resolved
jwst/datamodels/container.py Outdated Show resolved Hide resolved
jwst/cube_build/data_types.py Show resolved Hide resolved
@emolter
Copy link
Collaborator Author

emolter commented Oct 17, 2024

@melanieclarke Thanks for your initial testing of this. I did some additional testing with the same setup and found the following:

I'm testing using regression test data from nirspec/ifu:

It appears this test also fails on main, at the line self.output_name = input_models.meta.asn_table.products[0].name with an "AttributeError: No attribute 'asn_table'". I was able to make the step at least run by getting the datatypes right in the EMPTY_ASN_TABLE, but this causes the output names to look strange: they start with an underscore, e.g., _g395h-f290lp_s3d.fits. Of course this is worth fixing, but I'm not sure what's desired here because I can't get the step to run on a list of models on main. Also, having never looked at CubeBuild in detail before, I have no idea how the file naming works, and it looks complicated with the whole DataTypes class and everything.

MasterBackgroundStep was indeed causing problems because a list of input models doesn't have dedicated "background" exposures, so I added a graceful exit (step skip with warning) in that case. Filenames look fine but do not contain any association ID.

OutlierDetectionStep appears to work fine on both the list of files and the ModelContainer instantiated from them after one small change to the logic of that step's guess_mode function. Filenames look fine but do not contain any association ID.

PixelReplaceStep appears to work fine on both the list of files and the ModelContainer instantiated from them, but the output filenames look like e.g. step_PixelReplaceStep_7_pixelreplacestep.fits.

RE changes to save(), stpipe expects to be able to pass a save_model_func into the save method, and since we can't change the way stpipe handles things, it's necessary to keep that here. I still slightly simplified the method and added a unit test that shows the intended behavior. But let me know if you think that behavior should still be changed.

I do plan to test the rest of these on main also, but I suspect at least some of these issues existed before.

I'm starting to wonder whether initializing a ModelContainer from a list should issue a warning or even be disallowed entirely - the whole point of ModelContainer is to hold asn metadata.

Starting new regression tests here to see if the most recent changes introduced more problems.

@melanieclarke
Copy link
Collaborator

I do plan to test the rest of these on main also, but I suspect at least some of these issues existed before.

Sorry if I was unclear - there are definitely already some issues on main with running multi-input steps on a list instead of an association, including the crash in cube_build. I was just thinking that maybe we could start sorting them out here, or at least make sure that the design of ModelContainer doesn't get in the way of fixing them, especially if we're going to claim in the docs that "a list is a perfectly valid input and output to a pipeline or step".

I'm starting to wonder whether initializing a ModelContainer from a list should issue a warning or even be disallowed entirely - the whole point of ModelContainer is to hold asn metadata.

I would prefer to keep the ability to make a container from a list, if at all possible. Right now, feeding a list of filenames to cube_build makes a ModelContainer anyway, which makes sense so that the rest of the code can treat the input consistently. Making an asn file is a lot of overhead, sometimes it's nice to just be able to run a step on a list of input.

For less weird output filenames, we could maybe set the asn name to the basename for the first member model. But I don't think we need to have a complete answer now - these things don't currently work on main anyway, so if it's possible to support them, we can gather more input on what the output naming conventions should be. If it just doesn't make sense to support running these steps without an association, we should just throw an exception to make that clear to the user.

@emolter
Copy link
Collaborator Author

emolter commented Oct 18, 2024

"a list is a perfectly valid input and output to a pipeline or step"

This comment was meant primarily for developers, as in, "when deciding how a step is supposed to work, remember that you have the option to pass a list in and out of steps, you don't always need to turn it into a container". I can clarify this. Although it seems like that might be a bad idea currently because of the way stpipe tries to save step output...

For less weird output filenames, we could maybe set the asn name to the basename for the first member model

Seems like a good idea, will try.

I don't think we need to have a complete answer now - these things don't currently work on main anyway

Good to know where the goalpost is. If you're ok with it, I'm going to try the default asn name thing, and then just make it so all four of the steps I tested above run on both a list-initialized container and a list of filenames. Then I'll open a separate issue or issues for additional issues I notice.

Regression tests are all passing except the two unrelated tweakreg catalog ones.

@melanieclarke
Copy link
Collaborator

This comment was meant primarily for developers, as in, "when deciding how a step is supposed to work, remember that you have the option to pass a list in and out of steps, you don't always need to turn it into a container". I can clarify this. Although it seems like that might be a bad idea currently because of the way stpipe tries to save step output...

Yeah, maybe we should just remove that line for now.

Good to know where the goalpost is. If you're ok with it, I'm going to try the default asn name thing, and then just make it so all four of the steps I tested above run on both a list-initialized container and a list of filenames. Then I'll open a separate issue or issues for additional issues I notice.

Sounds good, thanks!

@emolter
Copy link
Collaborator Author

emolter commented Oct 18, 2024

The suggested addition of a default name for list-like input seems to fix the file name issues, or at least make the output names less silly-looking, for all but the pixel replace step. I opened an issue for that one: #8904.

The following script now runs successfully for both a list of filenames and a container initialized with a list of filenames:

# data from test_bigdata/jwst-pipeline/dev/nirspec/ifu/
fn = glob.glob(stem + "jw01249*_cal.fits")
container = datamodels.ModelContainer(fn)
# container = fn

CubeBuildStep.call(container, save_results=True)
MasterBackgroundStep.call(container, save_results=True)
OutlierDetectionStep.call(container, save_results=True)
PixelReplaceStep.call(container, save_results=True, skip=False)

Copy link
Collaborator

@melanieclarke melanieclarke left a comment

Choose a reason for hiding this comment

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

This all looks good to me now - thanks for all the updates!

We should probably run regression tests one more time, but I'll approve from my perspective now.

jwst/cube_build/data_types.py Show resolved Hide resolved
@emolter
Copy link
Collaborator Author

emolter commented Oct 18, 2024

Tracking down unit test failures now, all of which seem to have been introduced by my most recent commit.

@emolter
Copy link
Collaborator Author

emolter commented Oct 18, 2024

new regtests started here

edit: passing, but I'll wait until at least Monday after stand-up to merge this in case more reviews may come in

@emolter emolter requested a review from tapastro October 18, 2024 20:27
Copy link
Contributor

@tapastro tapastro left a comment

Choose a reason for hiding this comment

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

👍

@emolter emolter merged commit e0ef7c5 into spacetelescope:main Oct 22, 2024
31 checks passed
@emolter emolter deleted the JP-3721 branch October 22, 2024 20:46
hayescr pushed a commit to hayescr/jwst that referenced this pull request Oct 29, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Simplify ModelContainer
5 participants