PERF: delay rarely used imports (netcdf4, importlib, multiprocessing, tarfile, tomllib, tomli_w)

[neutrinoceros]

PR Summary

• import netCDF4 is responsible for ~6% startup overhead, let's delay it when we can afford it.

• import importlib.metadata is about ~3% (we still need to import it to determine matplotlib's version at runtime, which we'll be able to avoid when support for mpl<3.5 is dropped)

• multiprocessing, tarfile, tomllib, tomli_w are (collectively) about ~1%

nothing life-changing here but every % counts (yt's CLI still takes a couple seconds to do anything).

[neutrinoceros]

blocked by #4515

[matthewturk]

Do we have a breakdown of import costs? Maybe it could be a good way to look at splitting stuff up?

[neutrinoceros]

I profile startup with tuna (python -m pip install tuna)

This command opens an interactive flame graph showing everything that happens on import yt

python -X importtime -c "import yt" 2> import.log && python -m tuna import.log

unyt (together with sympy) still largely dominates the startup time, but we can separate yt from its most expensive dependencies with a more complex invoke:

python -X importtime -c "import numpy; import matplotlib; import sympy; import unyt; import yt" 2> import.log && python -m tuna import.log

I highly recommend running this profile with Python 3.11, and latest versions of sympy, unyt, and cmyt, all of which contain optimisations that make yt's startup slightly faster.

Also note that when switching branch/making changes, the measured startup time include byte code compilation, so results should not be considered realistic on the very first run.

Splitting up modules is unlikely to bring direct gains. Unyt is by far the biggest offender, but I think I've made every possible "easy" optimisations between v2.8 and v2.9.5 already.

[neutrinoceros]

I think I've made every possible "easy" optimisations between v2.8 and v2.9.5 already.

yt-project/unyt#429 👀

[neutrinoceros]

currently blocked by #4540

[chrishavlin]

So I've looked at this a few times and I'm feeling on the fence: when it's just one or two uses within a module, nesting those imports inside function calls is OK by me, particularly if it is an import of an external package. Where I'm uncertain is for our internal imports that are used throughout a module: e.g., I think the readability of parallel_analysis_tools suffers quite a bit by nesting all those from yt.data_objects.image_array import ImageArray and I'd like to know just how much of a speedup we're getting to justify the conflict with our style guide (which is admittedly vague -- maybe this is enough of a "good reason" to not import up top?).

If there are parts of yt that both take a long time to load and are not used in the majority of sessions, then I think it'd be worth looking at a more formal lazy import framework using something like importlib.util.LazyLoader. If the import is something that is used in most sessions then I'd personally stick with keeping the import cost up front in the initial import yt (e.g., is ImageArray ever not used when someone uses yt? this is a genuine question).

So I think I'd add an approval here but would appreciate a bit more discussion before this goes in cause it feels like a shift in coding style -- it seems hard to maintain these gains going forward without recommending nested imports in new code, because it certainly doesn't seem tenable for folks to simply remember which imports should be nested and which don't need it.

[neutrinoceros]

Thank you for having a look. I'm okay to drop parts of this PR, as the gains are absolutely not worth any level of controversy :)

#4539 is a much more efficient, and completely different approach, towards the same goal, so if that one is okay I'll be happy to re-evaluate the present patch.

[chrishavlin]

Well I don't think this strays into controversial :) It's more a practical question of whether the gains you get here will be easily maintained down the line. I think the external imports are easy enough, but not so sure about internal yt imports. The very specific yt imports (like your changes to plot_modifications) probably are pretty easy, but it seems very likely that someone would write some new unrelated code at some point that imports ImageArray and inadvertently un-delays your delay.

[neutrinoceros]

Yeah that's the biggest problem with this technique, it's very easy to ruin hours of effort with just one from matplotlib import pyplot as plt, it is much preferable to gain time by reducing side effects. I find that in practice, these are easier to spot when other effects (costly imports) are neutralised.

[neutrinoceros]

@chrishavlin so I did my homework: here's a box plot of import times for the main branch against this branch

I ran import yt 200 times on each version. The first commit by itself seems to correspond to a statistically significant gain, but that isn't true for commits 2 and 3, which I think are the ones that got you on the fence, so I'm happy to drop them if that's satisfactory to you :)

[chrishavlin]

Thanks for breaking it down by commit! Ya, I think with just the first commit I'd hit approve right away :)

Does look like the monkeypatch adjustment in yt/tests/test_external_frontends.py should probably be kept from the 2nd commit though. And I'm OK with keeping the changes to external imports (matplotlib and csv) from your 2nd/3rd commits, but if they're not actually getting you a speedup, might not be worth picking out those changes to include here.

[neutrinoceros]

revamped ! I kept the first commit, and followed your suggestions on what to keep fro the 2nd and 3rd

[chrishavlin]

Great! LGTM now!