flexcompute / tidy3d

Fast electromagnetic solver (FDTD) at scale.
https://docs.flexcompute.com/projects/tidy3d/en/latest/
GNU Lesser General Public License v2.1
176 stars 40 forks source link

Speed up unit tests #1821

Open momchil-flex opened 2 months ago

momchil-flex commented 2 months ago

Our tests are starting to get a bit heavy to run, so I timed them. Here's all that take more than 1s.

225.51s call     tests/test_timing/import_tidy3d.py::test_import
43.32s call     tests/test_components/test_simulation.py::test_num_mediums
20.20s call     tests/test_components/test_autograd.py::test_autograd_async[<ALL>-field_vol]
17.94s call     tests/test_components/test_autograd.py::test_autograd_async[<ALL>-diff]
16.55s call     tests/test_components/test_autograd.py::test_autograd_async[<ALL>-mode]
15.43s call     tests/test_plugins/test_invdes.py::test_continue_run_from_file
15.17s call     tests/test_plugins/test_mode_solver.py::test_mode_solver_straight_vs_angled
14.91s call     tests/test_components/test_autograd.py::test_autograd_async[<ALL>-field_point]
11.74s call     tests/test_data/test_sim_data.py::test_plot[1.0]
10.89s call     tests/test_plugins/test_adjoint.py::test_adjoint_pipeline[False]
10.63s call     tests/test_plugins/test_adjoint.py::test_adjoint_pipeline[True]
10.11s call     tests/test_plugins/test_invdes.py::test_result_store_full_results_is_false
9.89s call     tests/test_plugins/test_adjoint.py::TestJaxComplexPolySlab::test_vertices_grads[10-5-base_vertices1]
9.79s call     tests/test_data/test_sim_data.py::test_plot[0]
9.73s call     tests/test_components/test_autograd.py::test_interp_objectives[flux-False]
9.02s call     tests/test_plugins/test_adjoint.py::test_adjoint_run_async[True]
8.97s call     tests/test_plugins/test_design.py::test_sweep[sweep_method0]
8.91s call     tests/test_components/test_autograd.py::test_interp_objectives[intensity-False]
7.97s call     tests/test_components/test_eme.py::test_eme_sim_data
7.55s call     tests/test_components/test_sidewall.py::test_intersection_with_inside_poly
7.33s call     tests/test_components/test_autograd.py::test_autograd_speed_num_structures
7.18s call     tests/test_components/test_autograd.py::test_interp_objectives[flux-True]
7.14s call     tests/test_components/test_scene.py::test_num_mediums
6.93s call     tests/test_plugins/test_dispersion_fitter.py::test_dispersion_set_wvg_range
6.83s call     tests/test_components/test_autograd.py::test_interp_objectives[intensity-True]
6.41s call     tests/test_components/test_simulation.py::test_warn_large_epsilon[1-11000-WARNING]
6.16s call     tests/test_plugins/test_adjoint.py::test_adjoint_run_async[False]
6.14s call     tests/test_plugins/test_adjoint.py::test_adjoint_pipeline_2d[True]
6.00s call     tests/test_plugins/test_adjoint.py::test_adjoint_pipeline_2d[False]
5.93s call     tests/test_plugins/test_invdes.py::test_continue_run_fns
5.50s call     tests/test_plugins/test_dispersion_fitter.py::test_dispersion_load_file
5.49s call     tests/test_data/test_monitor_data.py::test_mode_solver_data_sort
5.40s call     tests/test_components/test_autograd.py::test_autograd_objective[<ALL>-field_vol]
5.07s call     tests/test_components/test_custom.py::test_custom_lorentz[True]
4.82s call     tests/test_plugins/test_dispersion_fitter.py::test_lossy_dispersion
4.68s call     tests/test_components/test_autograd.py::test_autograd_objective[<ALL>-diff]
4.60s call     tests/test_components/test_autograd.py::test_autograd_objective[<ALL>-mode]
4.38s call     tests/test_components/test_autograd.py::test_autograd_async[custom_pole_res-<ALL>]
4.38s call     tests/test_plugins/test_adjoint.py::test_num_input_structures
4.36s call     tests/test_components/test_autograd.py::test_autograd_objective[<ALL>-field_point]
4.21s call     tests/test_plugins/test_mode_solver.py::test_group_index[True]
4.14s call     tests/test_components/test_autograd.py::test_autograd_async[geo_group-<ALL>]
4.00s call     tests/test_plugins/test_adjoint.py::TestJaxComplexPolySlab::test_vertices_grads[10-5-base_vertices0]
3.89s call     tests/test_components/test_custom.py::test_custom_sellmeier[True]
3.82s call     tests/test_plugins/test_adjoint.py::test_polyslab_2d[0]
3.77s call     tests/test_components/test_autograd.py::test_autograd_async[custom_med_vec-<ALL>]
3.73s call     tests/test_plugins/test_dispersion_fitter.py::test_lossless_dispersion
3.66s call     tests/test_components/test_autograd.py::test_autograd_async[polyslab-<ALL>]
3.39s call     tests/test_plugins/test_adjoint.py::test_polyslab_2d[10]
3.38s call     tests/test_web/test_webapi_eme.py::test_job
3.34s call     tests/test_web/test_webapi_eme.py::test_run
3.33s call     tests/test_web/test_webapi_heat.py::test_job
3.32s call     tests/test_web/test_webapi_heat.py::test_run
3.28s call     tests/test_plugins/test_adjoint.py::test_run_flux
3.19s call     tests/test_components/test_custom.py::test_custom_pole_residue[True]
3.14s call     tests/test_plugins/test_invdes.py::test_invdes_io
3.07s call     tests/test_components/test_custom.py::test_custom_debye[True]
2.91s call     tests/test_components/test_autograd.py::test_autograd_async[custom_med-<ALL>]
2.91s call     tests/test_components/test_simulation.py::test_warn_large_epsilon[50-1-WARNING]
2.91s call     tests/test_components/test_field_projection.py::test_proj_clientside
2.85s call     tests/test_components/test_custom.py::test_custom_anisotropic_medium[True]
2.82s call     tests/test_plugins/test_mode_solver.py::test_mode_solver_custom_medium[False]
2.75s call     tests/test_plugins/test_mode_solver.py::test_group_index[False]
2.55s call     tests/test_plugins/test_invdes.py::test_default_params
2.52s call     tests/test_plugins/test_mode_solver.py::test_mode_solver_simple[True]
2.50s call     tests/test_components/test_autograd.py::test_autograd_async[medium-<ALL>]
2.45s call     tests/test_plugins/test_adjoint.py::TestJaxComplexPolySlab::test_matches_complexpolyslab[10-5-base_vertices1]
2.44s call     tests/test_components/test_autograd.py::test_autograd_async[center_list-<ALL>]
2.42s call     tests/test_components/test_autograd.py::test_autograd_async[size_element-<ALL>]
2.39s call     tests/test_plugins/test_adjoint.py::test_jax_data_array
2.39s call     tests/test_components/test_simulation.py::test_sim_subsection[True-13]
2.35s call     tests/test_plugins/test_polyslab.py::test_many_sub_polyslabs
2.33s call     tests/test_components/test_autograd.py::test_autograd_async[pole_res-<ALL>]
2.27s call     tests/test_components/test_IO.py::test_validation_speed
2.25s call     tests/test_components/test_autograd.py::test_sim_full_ops[custom_med]
2.15s call     tests/test_plugins/test_invdes.py::test_warn_zero_grad
2.11s call     tests/test_plugins/test_dispersion_fitter.py::test_constant_loss_tangent
1.91s call     tests/test_components/test_custom.py::test_custom_drude[True]
1.86s call     tests/test_components/test_simulation.py::test_sim_subsection[True-1]
1.85s call     tests/test_components/test_scene.py::test_structure_alpha
1.81s call     tests/test_plugins/test_mode_solver.py::test_mode_solver_relative
1.79s call     tests/test_components/test_simulation.py::test_sim_subsection[False-1]
1.78s call     tests/test_components/test_simulation.py::test_sim_subsection[False-13]
1.74s call     tests/test_plugins/test_mode_solver.py::test_mode_solver_2D
1.72s call     tests/test_plugins/test_adjoint.py::TestJaxComplexPolySlab::test_vertices_grads[10-1-base_vertices1]
1.72s call     tests/test_components/test_parameter_perturbation.py::test_charge_perturbation
1.66s call     tests/test_components/test_custom.py::test_custom_isotropic_medium[True]
1.61s call     tests/test_plugins/test_adjoint.py::test_adjoint_setup_fwd
1.52s call     tests/test_plugins/test_adjoint.py::TestJaxComplexPolySlab::test_vertices_grads[10-0-base_vertices1]
1.41s call     tests/test_components/test_IO.py::test_simulation_load_export
1.40s call     tests/test_plugins/test_adjoint.py::test_pytreedef_errors
1.34s call     tests/test_components/test_simulation.py::test_run_time_spec
1.32s call     tests/test_plugins/test_mode_solver.py::test_mode_solver_simple[False]
1.30s call     tests/test_plugins/test_mode_solver.py::test_mode_solver_web_run_batch
1.28s call     tests/test_plugins/test_component_modeler.py::test_run_component_modeler
1.26s call     tests/test_plugins/test_invdes.py::test_result_data_multi
1.24s call     tests/test_plugins/test_adjoint.py::TestJaxComplexPolySlab::test_vertices_grads[10-1-base_vertices0]
1.24s call     tests/test_plugins/test_component_modeler.py::test_run_component_modeler_mappings
1.23s call     tests/test_plugins/test_component_modeler.py::test_to_from_file_batch
1.18s call     tests/test_data/test_datasets.py::test_triangular_dataset[test123]
1.17s call     tests/test_components/test_eme.py::test_eme_simulation
1.15s call     tests/test_components/test_source.py::test_plot_source_time
1.12s call     tests/test_plugins/test_adjoint.py::TestJaxComplexPolySlab::test_matches_complexpolyslab[10-5-base_vertices0]
1.12s call     tests/test_plugins/test_terminal_component_modeler.py::test_run_coaxial_component_modeler
1.11s call     tests/test_plugins/test_design.py::test_sweep[sweep_method3]
1.10s call     tests/test_plugins/test_design.py::test_sweep[sweep_method1]
1.10s call     tests/test_plugins/test_design.py::test_sweep[sweep_method2]
1.09s call     tests/test_plugins/test_component_modeler.py::test_mapping_exclusion
1.09s call     tests/test_components/test_autograd.py::test_autograd_objective[custom_pole_res-<ALL>]
1.05s call     tests/test_components/test_autograd.py::test_autograd_objective[polyslab-<ALL>]
1.00s call     tests/test_components/test_autograd.py::test_autograd_objective[custom_med_vec-<ALL>]

The tidy3d_import.py one imports 100 times, which seems excessive. It also doesn't really test anything, it only prints. Was it not supposed to be part of the test suite at all @daquintero, rather just for internal use?

I made a PR for the test_num_mediums one.

Apart from that there's a lot of adjoint/autograd stuff @tylerflex @yaugenst-flex , maybe some of it is inevitable, but maybe some can be improved? Not high priority.

Then there's some simulation data and mode solver ones, but again, not high priority.

yaugenst-flex commented 2 months ago

There definitely are some tests that just take longer but that still need to be tested. What we could do is mark some tests as slow (pytest.mark.slow), and then we can decide when those should be run?

At some point we should maybe also think about refactoring the tests...

daquinteroflex commented 2 months ago

Actually that's strange. I named that file import_tidy3d.py so that it would not get caught in the main pytest command. Possibly, when the docstring tests were upgraded this got included in the main test suite accidentally. The goal was to have that file whenever we wanted to benchmark manually but shouldn't be commonly run. Possibly this involves just updating the testing configuration settings to ignore that file and that'd work?

I'm checking now how to avoid it being included, but running into a weird __pycache__ local testing issue. Will comment again or PR when I work out how to remove it properly.

yaugenst-flex commented 2 months ago

Actually that's strange. I named that file import_tidy3d.py so that it would not get caught in the main pytest command.

In our pyproject.toml, we set

[tool.pytest.ini_options]
python_files = "*.py"

Meaning, we are globbing for any python file, and the test is then discovered due to it being named test_. However, if this import thing is not actually a test but a script, maybe we should not put it under tests/?

daquinteroflex commented 2 months ago

Oh yeah, that was a doctest change. Nice, we should move it to scripts/ then. I can include this in the docs PR I'm trying to finish polishing atm.

yaugenst-flex commented 2 months ago

Alternatively, we could mark it, by registering a marker in pytest.ini:

[pytest]
markers =
    optional: mark a test as optional.

Mark the test:

@pytest.mark.optional
def test_import():
    ...

And then we can exclude this marker from running by default and run it via pytest -m optional We can do the same for slow tests.

momchil-flex commented 2 months ago

Oh yeah, that was a doctest change. Nice, we should move it to scripts/ then. I can include this in the docs PR I'm trying to finish polishing atm.

Maybe let's just have a tiny PR for that so it already saves ~3mins for the tests? :)

daquinteroflex commented 2 months ago

PR opened with the move!

daquinteroflex commented 2 months ago

To add to this issue.

Yannick suggested also to improve what we're caching in the Github actions. I think we can cache the poetry virtual environment but this did cause issues beforehand during the first poetry implementation, but I think can be revisited.

momchil-flex commented 2 months ago

I can leave this open for future improvements, but I think nothing else is planned for now?