Speedup tests (#3696)

* Add kwarg `n_cpu`
* Set to use only single core in `test_cube_fit_after_unit_change` to avoid overhead from multiprocessing

* Move more tests to single cpu
* Also add coverage for multiprocessing vs single cpu in cube fitting backend test.

* Move another test to single cpu
* Model fitting strikes again!

* Move another test to single cpu
* Model fitting strikes again!

* Use fixture instead of for-loop
* Can be a tiny bit slower in serial but in a parallel environment, it's a no-brainer.

* Style fixes

* Add parallelization for pytest

* Example style fix

* Change load distribution technique
* Should be a bit better than the old since we only have two workers (for now)

* Add `cached_uri` logic

* Re-arranged uri_or_file logic

* Add durations flag for future test optimization/debugging

* Update jdaviz/core/tests/test_autoconfig.py

Co-authored-by: Kyle Conroy <[email protected]>

* Hardcoding n_cpu=1 instead of using a variable

* Style adjustment

* Re-use n_cpu for one of the tests so it's a variable again

* Overwrite uri instead of new variable

* Create ParallelMixin to use traitlet instead of kwarg
* Barebones framework for eventual work with parallelization

* Removed unnecessary import

* Use an integer traitlet instead of unicode
* To be consistent with current implementation and avoid the try except. This could change in the future

* n_cpu is now a traitlet passed directly into the model_fitting object

* Revert traitlet (keeping mixin)
* parallel_n_cpu is now (temporarily) an attribute of model_fitting
* traitlet implementation and ParallelMixin to come later
* Also now both fitting backends are parameterized for multiprocessing tests for proper coverage

* Fixed bug, properly access object now

* Leaving a TODO for future parallel work

* Moved comment about running in serial to first instance of calculate_fit
* Also force n_cpu to 1 in that test as well. I think I caught them all now.

* Style fixes

* Style fixes

* Remove ParallelMixin
* To be worked on eventually!

* Update template_mixin.py

Style fixes

* Update model_fitting.py

Taking out references to ParallelMixin

* Apply suggestions from code review

Removed leftover code (thanks Ricky!)

Co-authored-by: Ricky O'Steen <[email protected]>

* Update jdaviz/configs/default/plugins/model_fitting/tests/test_plugin.py

Co-authored-by: Ricky O'Steen <[email protected]>

* Apply suggestions from code review

Co-authored-by: Kyle Conroy <[email protected]>

---------

Co-authored-by: Kyle Conroy <[email protected]>
Co-authored-by: Ricky O'Steen <[email protected]>

Author: 3 people

jdaviz/configs/default/plugins/model_fitting/model_fitting.py

@@ -168,6 +168,7 @@ def __init__(self, *args, **kwargs):
         self.hub.subscribe(self, GlobalDisplayUnitChanged,
                            handler=self._on_global_display_unit_changed)
 
+        self.parallel_n_cpu = None
         self._set_relevant()
 
     @observe('dataset_items')
@@ -1185,7 +1186,8 @@ def _fit_model_to_spectrum(self, add_data):
                 models_to_fit,
                 self.model_equation,
                 run_fitter=True,
-                window=None
+                window=None,
+                n_cpu=self.parallel_n_cpu,
             )
         except AttributeError:
             msg = SnackbarMessage("Unable to fit: model equation may be invalid",
@@ -1299,7 +1301,8 @@ def _fit_model_to_cube(self, add_data):
                 models_to_fit,
                 self.model_equation,
                 run_fitter=True,
-                window=None
+                window=None,
+                n_cpu=self.parallel_n_cpu
             )
         except ValueError:
             snackbar_message = SnackbarMessage(

jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py

@@ -101,6 +101,9 @@ def test_parameter_retrieval(cubeviz_helper, spectral_cube_wcs):
     assert cubeviz_helper.app._get_display_unit('sb') == sb_unit
 
     plugin.create_model_component("Linear1D", "L")
+    # NOTE: Hardcoding n_cpu=1 to run in serial, it's slower to spool up
+    # multiprocessing for the size of the cube
+    plugin._obj.parallel_n_cpu = 1
     with warnings.catch_warnings():
         warnings.filterwarnings('ignore', message='Model is linear in parameters.*')
         plugin.calculate_fit()
@@ -119,8 +122,13 @@ def test_parameter_retrieval(cubeviz_helper, spectral_cube_wcs):
                              atol=1e-10 * sb_unit)
 
 
-@pytest.mark.parametrize('unc', ('zeros', None))
-def test_fitting_backend(unc):
+@pytest.mark.parametrize(
+    ('n_cpu', 'unc'), [
+        (1, "zeros"),
+        (1, None),
+        (None, "zeros"),
+        (None, None),])
+def test_fitting_backend(n_cpu, unc):
     np.random.seed(42)
 
     x, y = build_spectrum()
@@ -142,24 +150,30 @@ def test_fitting_backend(unc):
 
     # Returns the initial model
     fm, fitted_spectrum = fb.fit_model_to_spectrum(spectrum, model_list, expression,
-                                                   run_fitter=False)
+                                                   run_fitter=False, n_cpu=n_cpu)
 
     parameters_expected = np.array([0.7, 4.65, 0.3, 2., 5.55, 0.3, -2.,
                                     8.15, 0.2, 1.])
     assert_allclose(fm.parameters, parameters_expected, atol=1e-5)
 
     # Returns the fitted model
     fm, fitted_spectrum = fb.fit_model_to_spectrum(spectrum, model_list, expression,
-                                                   run_fitter=True)
+                                                   run_fitter=True, n_cpu=n_cpu)
 
     parameters_expected = np.array([1.0104705, 4.58956282, 0.19590464, 2.39892026,
                                     5.49867754, 0.10834472, -1.66902953, 8.19714439,
                                     0.09535613, 3.99125545])
     assert_allclose(fm.parameters, parameters_expected, atol=1e-5)
 
 
-@pytest.mark.parametrize('unc', ('zeros', None))
-def test_cube_fitting_backend(cubeviz_helper, unc, tmp_path):
+# For coverage of serial vs multiprocessing and unc
+@pytest.mark.parametrize(
+    ('n_cpu', 'unc'), [
+        (1, "zeros"),
+        (1, None),
+        (None, "zeros"),
+        (None, None),])
+def test_cube_fitting_backend(cubeviz_helper, unc, n_cpu, tmp_path):
     np.random.seed(42)
 
     SIGMA = 0.1  # noise in data
@@ -214,9 +228,6 @@ def test_cube_fitting_backend(cubeviz_helper, unc, tmp_path):
     model_list = [g1f, g2f, g3f, zero_level]
     expression = "g1 + g2 + g3 + const1d"
 
-    n_cpu = None
-    # n_cpu = 1  # NOTE: UNCOMMENT TO DEBUG LOCALLY, AS NEEDED
-
     # Fit to all spaxels.
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore", message="The fit may be unsuccessful.*")
@@ -316,6 +327,7 @@ def test_results_table(specviz_helper, spectrum1d):
     uc = specviz_helper.plugins['Unit Conversion']
 
     mf.add_results.label = 'linear model'
+    mf._obj.parallel_n_cpu = 1
     with warnings.catch_warnings():
         warnings.filterwarnings('ignore', message='Model is linear in parameters.*')
         mf.calculate_fit(add_data=True)
@@ -401,6 +413,7 @@ def test_incompatible_units(specviz_helper, spectrum1d):
     mf.create_model_component('Linear1D')
 
     mf.add_results.label = 'model native units'
+    mf._obj.parallel_n_cpu = 1
     with warnings.catch_warnings():
         warnings.filterwarnings('ignore', message='Model is linear in parameters.*')
         mf.calculate_fit(add_data=True)
@@ -430,6 +443,8 @@ def test_cube_fit_with_nans(cubeviz_helper):
     mf = cubeviz_helper.plugins["Model Fitting"]
     mf.cube_fit = True
     mf.create_model_component("Const1D")
+    mf._obj.parallel_n_cpu = 1
+
     with warnings.catch_warnings():
         warnings.filterwarnings('ignore', message='Model is linear in parameters.*')
         mf.calculate_fit()
@@ -456,6 +471,8 @@ def test_cube_fit_with_subset_and_nans(cubeviz_helper):
     mf.cube_fit = True
     mf.spectral_subset = 'Subset 1'
     mf.create_model_component("Const1D")
+    mf._obj.parallel_n_cpu = 1
+
     with warnings.catch_warnings():
         warnings.filterwarnings('ignore', message='Model is linear in parameters.*')
         mf.calculate_fit()
@@ -473,6 +490,7 @@ def test_fit_with_count_units(cubeviz_helper):
     mf = cubeviz_helper.plugins["Model Fitting"]
     mf.cube_fit = True
     mf.create_model_component("Const1D")
+    mf._obj.parallel_n_cpu = 1
 
     # ensures specutils.Spectrum.with_flux_unit has access to Jdaviz custom equivalencies for
     # PIX^2 unit
@@ -504,6 +522,7 @@ def test_cube_fit_after_unit_change(cubeviz_helper, solid_angle_unit):
     # Check that the parameter is using the current units when initialized
     assert mf._obj.component_models[0]['parameters'][0]['unit'] == f'MJy / {solid_angle_string}'
 
+    mf._obj.parallel_n_cpu = 1
     with warnings.catch_warnings():
         warnings.filterwarnings('ignore', message='Model is linear in parameters.*')
         mf.calculate_fit()
@@ -569,6 +588,7 @@ def test_deconf_mf_with_subset(deconfigged_helper):
     # ensure deconfigged app can access subsets in plugin
     mf.spectral_subset.selected = 'Subset 1'
     mf.add_results.label = 'linear model'
+    mf._obj.parallel_n_cpu = 1
     mf.create_model_component()
 
     with warnings.catch_warnings():

jdaviz/configs/default/plugins/model_fitting/tests/test_plugin.py

@@ -150,6 +150,8 @@ def test_register_cube_model(cubeviz_helper, spectrum1d_cube):
     modelfit_plugin.reestimate_model_parameters()
     assert modelfit_plugin._obj.results_label_default == 'model'
     assert modelfit_plugin._obj.results_label == test_label
+
+    modelfit_plugin._obj.parallel_n_cpu = 1
     with warnings.catch_warnings():
         warnings.filterwarnings('ignore', message='.*Model is linear in parameters.*')
         modelfit_plugin.calculate_fit()
@@ -177,6 +179,8 @@ def test_fit_cube_no_wcs(cubeviz_helper):
     mf.cube_fit = True
     # Need to manually reestimate the parameters to update the units
     mf.reestimate_model_parameters()
+
+    mf._obj.parallel_n_cpu = 1
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore", message="Model is linear in parameters.*")
         fitted_model, output_cube = mf.calculate_fit(add_data=True)

jdaviz/configs/specviz/plugins/line_analysis/tests/test_lineflux.py

@@ -5,8 +5,6 @@
 import pytest
 from specutils import Spectrum
 
-from jdaviz import Cubeviz
-
 import warnings
 warnings.filterwarnings('ignore')
 
@@ -73,24 +71,30 @@ def _calculate_line_flux(viz_helper):
 @pytest.mark.filterwarnings(r"ignore:.* contains multiple slashes")
 @pytest.mark.filterwarnings(r"ignore:.* apply_slider_redshift")
 @pytest.mark.parametrize('spectra_fluxunit', test_cases)
-def test_cubeviz_collapse_fluxunits(spectrum1d_cube_custom_fluxunit, spectra_fluxunit):
-    ''' Calculates line flux and checks the units for each collapse function '''
+@pytest.mark.parametrize('function', COLLAPSE_FUNCTIONS)
+def test_cubeviz_collapse_fluxunits(
+        cubeviz_helper,
+        spectrum1d_cube_custom_fluxunit,
+        function, spectra_fluxunit):
+    """
+    Calculates line flux and checks the units for each collapse function.
+    """
+
     data = spectrum1d_cube_custom_fluxunit(spectra_fluxunit)
-    for function in COLLAPSE_FUNCTIONS:
-        # Initialize Cubeviz with specific data and collapse function
-        cubeviz_helper = Cubeviz()
-        data_label = "Test Cube"
-        cubeviz_helper.load_data(data, data_label=data_label)
-        cubeviz_helper.app.get_viewer('spectrum-viewer').state.function = function
-
-        lineflux_result = _calculate_line_flux(cubeviz_helper)
-        autocollapsed_spectrum_unit = (cubeviz_helper.
-                                       specviz.get_spectra()["Spectrum (sum)"].flux.unit)
-        # Futureproofing: Eventually Cubeviz autocollapse will change the flux units of the
-        # spectra depending on whether the spectrum was collapsed OVER SPAXELS or not. Only
-        # do the assertion check if we KNOW what the expected lineflux results should be
-        if autocollapsed_spectrum_unit in expected_lineflux_results.keys():
-            assert u.Unit(lineflux_result['unit']) == expected_lineflux_results[spectra_fluxunit]
+
+    # Initialize Cubeviz with specific data and collapse function
+    data_label = "Test Cube"
+    cubeviz_helper.load_data(data, data_label=data_label)
+    cubeviz_helper.app.get_viewer('spectrum-viewer').state.function = function
+
+    lineflux_result = _calculate_line_flux(cubeviz_helper)
+    autocollapsed_spectrum_unit = (cubeviz_helper.
+                                   specviz.get_spectra()["Spectrum (sum)"].flux.unit)
+    # Futureproofing: Eventually Cubeviz autocollapse will change the flux units of the
+    # spectra depending on whether the spectrum was collapsed OVER SPAXELS or not. Only
+    # do the assertion check if we KNOW what the expected lineflux results should be
+    if autocollapsed_spectrum_unit in expected_lineflux_results.keys():
+        assert u.Unit(lineflux_result['unit']) == expected_lineflux_results[spectra_fluxunit]
 
 
 @pytest.mark.filterwarnings(r"ignore:.* contains multiple slashes")

jdaviz/core/tests/test_autoconfig.py

@@ -60,11 +60,15 @@ def test_launcher(tmp_path):
 
     # Test with real files
     for uri, config in AUTOCONFIG_EXAMPLES:
+        uri = cached_uri(uri)
         if uri.startswith("mast:"):
             download_path = str(tmp_path / Path(uri).name)
             Observations.download_file(uri, local_path=download_path)
         elif uri.startswith("http"):
             download_path = download_file(uri, cache=True, timeout=100)
+        else:
+            # cached local file
+            download_path = uri
         launcher.filepath = download_path
         # In testing, setting the filepath will stall until identifying is complete
         # No need to be concerned for race condition here

tox.ini

@@ -62,6 +62,7 @@ deps =
     devdeps: git+https://github.com/glue-viz/glue-astronomy.git
     devdeps: git+https://github.com/widgetti/solara.git
     devdeps: git+https://github.com/astropy/specreduce.git
+    pytest-xdist
 
 # The following indicates which extras_require from pyproject.toml will be installed
 extras =
@@ -73,8 +74,8 @@ extras =
 commands =
     jupyter --paths
     pip freeze
-    !cov: pytest --pyargs jdaviz {toxinidir}/docs --ignore=jdaviz/qt.py {posargs}
-    cov: pytest --pyargs jdaviz {toxinidir}/docs --cov jdaviz --cov-config={toxinidir}/pyproject.toml --ignore=jdaviz/qt.py {posargs}
+    !cov: pytest -n auto --dist loadscope --pyargs jdaviz {toxinidir}/docs --ignore=jdaviz/qt.py {posargs} --durations=50
+    cov: pytest -n auto --dist loadscope --pyargs jdaviz {toxinidir}/docs --cov jdaviz --cov-config={toxinidir}/pyproject.toml --ignore=jdaviz/qt.py {posargs} --durations=50
     cov: coverage xml -o {toxinidir}/coverage.xml
 
 pip_pre =