Speedup import and add regression check for import time

[DanielYang59]

Summary

• Speedup import, to close pymatviz import cost way too high #209
• Add check for import time
• Update .pytest-split-durations (how to do this?)

Profile command: python -X importtime -c "import pymatviz" 2> pmv.log && tuna pmv.log

[DanielYang59]

I have a bad feeling about this PR, I think I might prefer to wait until the import fix in monty get merged:

• It's badly polluting the code base by lazily almost everything about pymatgen, because any import involving from monty.json import xxx could potentially slow down the import (if torch is installed).

---

Meanwhile for some reason every time I lazily import one package, there seems to be another package popping up from nowhere and make the total import time roughly the same. I think the reason is at the end the same costly packages are imported somewhere, and by refactoring iteratively, they're just broken into some smaller chunks (before refactor, they're all imported in the first module that needed them, then we keep relocating the current most expensive piece to another) so the most expensive chunk keep shrinking but the total import time is roughly unchanged.

Profile on main branch (without torch):

The tip of this branch:

[janosh]

Thanks a lot for digging into this, really appreciated! I suspect you're right, we're dependent on upstream PRs being merged to achieve significant improvements here.

[janosh]

i think it would be good to add performance regression tests to this PR and run them both on main and in this branch to compare. we should have them to catch future slow downs before they enter main. i posted some half-baked import time tests in #209 (comment) which might help. no worries if you're busy, i'll try and get to it myself eventually

[DanielYang59]

I will do this :) I was planning to add a runtime unit test as well (perhaps to core parts of pymatgen as well), now we only test the results, and runtime is not covered at all. Thanks again for bringing up this idea!

[DanielYang59]

Regarding the actual test methodology, I haven't got much finding about best practice on doing an import time test. Your method certainly works, what about spawning a subprocess (is there any pitfall about this, like subprocess overhead maybe?), something like:

import subprocess
import time


def measure_import_time(import_command: str) -> float:
    start_time = time.time()
    subprocess.run(['python', '-c', 'import pymatviz'], check=True)  # import_command
    return time.time() - start_time

---

I hope I'm correct here, I guess there's a slight pitfall with the sys.modules method though, because it only pops the outer module. Say if we have monty which imports torch indirectly/internally, if we just pop monty, the actual expensive torch would still remain in cache, so:

• if name == module_name or name.startswith(f"{module_name}.") doesn't seem to capture indirect imports?
• If we were to profile several modules at the same time (say pmv.a and pmv.b both import module), the shared internal module would still be cached.
• it's a bit hard to average across several runs?

---

Is there any advantage to use relative time (time.perf_counter()) over real-world time (time.time())? The latter seems more comprehensible to me (though at the end of the day we still need to collect the reference time from CI).

[janosh]

Your method certainly works, what about spawning a subprocess (is there any pitfall about this, like subprocess overhead maybe?), something like:

i'm not sure but i think subprocess should be fine since afaik it runs in a completely isolated process with separate import caching

Is there any advantage to use relative time (time.perf_counter()) over real-world time (time.time())?

yes, have a look at this video which explains the virtues of time.perf_counter

[DanielYang59]

Your method certainly works, what about spawning a subprocess (is there any pitfall about this, like subprocess overhead maybe?), something like:

i'm not sure but i think subprocess should be fine since afaik it runs in a completely isolated process with separate import caching

Thanks! Let's use subprocess for now, this should make things much easier (if there is not pitfall)

Is there any advantage to use relative time (time.perf_counter()) over real-world time (time.time())?

yes, have a look at this video which explains the virtues of time.perf_counter

Thanks for sharing, I would have a look later but I certainly trust your judgement.

By the way, how do we regenerate the .pytest-split-durations record (not now, after we finish everything, current import is still very slow and might need more tuning).

[DanielYang59]

Currently like 12% of the import time is used to import import plotly.figure_factory in ptable.plotly where it's only used once across the entire code base, I have lazily imported it ed48b13.

After that:

• The sympy import by pymatgen should also be done lazily bump monty to use the monty.json import speedup patch, add import time regression test, lazy load some rarely used but costly modules materialsproject/pymatgen#4128 (comment).
• scipy doesn't seem to have much room for improvement because it's used a lots of places.

[janosh]

By the way, how do we regenerate the .pytest-split-durations record (not now, after we finish everything, current import is still very slow and might need more tuning).

the command for that is

pytest --store-durations --durations-path tests/.pytest-split-durations

i should have documented that somewhere. probably best in test.yml where we set test-cmd: ...

[DanielYang59]

Thanks a ton for letting me know!

What I don't really understand is how do we generate this runtime record via GitHub CI runners (in contrast to run that command locally)?

Especially for pymatgen where:

• it's huge and I don't have that much resources to run all jobs locally
• tasks are split (I haven't tried myself, perhaps it would be combined automatically?)

I have a feeling that we may not want this test to run for each commit (because I don't expect we would introduce changes that vary import time very often), it's slowing down everything. Perhaps we only run it only when merging into main?

For this PR, I currently cannot think of more improvement to make as the left is pretty much core packages like scipy, pandas, matplotlib. Do we want to perhaps merge it (after I modify it to run at main branch only)

[janosh]

very thankful for this PR! 👍 didn't mean to keep it parked so long

What I don't really understand is how do we generate this runtime record via GitHub CI runners (in contrast to run that command locally)?

i think generating locally should be fine. what matters shouldn't be absolute values for test run times but how long a test takes relative to the others which will be similar in CI and locally

it's huge and I don't have that much resources to run all jobs locally

even pymatgen only takes 10 min or so to run the whole test suite locally for me. curious how long for you?

Perhaps we only run it only when merging into main?

is the test slow enough in CI that it would extend overall time to checks complete if it ran inside its own split? i just tried test_import_time locally and got 80.52s. will be more on GitHub but tests already take ~2 min.
either way, i think it's fine to only run this test on main

is this PR ready to go from your side? are we waiting on a monty release to get the new startup benefits?

[janosh]

are we waiting on a monty release to get the new startup benefits?

ah, saw your comment in pyproject

    # TODO: pmv doesn't actually depend on monty, however latest monty
    # includes a critical import patch, remove this after pmg bump dep
    "monty>=2024.10.21",

[DanielYang59]

i think generating locally should be fine.

even pymatgen only takes 10 min or so to run the whole test suite locally for me. curious how long for you?

Yep should takes about the same time just like pymatgen CI (macOS runner is running on M1 chip, each split takes around 1 min), as I'm usually coding on a M3 MacBook Air laptop. I'm just being curious as I think running on a dedicated runner would give better scalability and consistency :)

is the test slow enough in CI that it would extend overall time to checks complete if it ran inside its own split?

The original implementation was running 10 repeats, but I reduced repeat to 3 so should be good now.

[DanielYang59]

This part was intended to generate a directly copyable dict (and human-readable for manual inspection) for us to directly copy and paste to update the time recording, printing a dict directly gives a one-liner:

Copyable import time dictionary:
{'pymatviz': 2084.25, 'pymatviz.coordination': 2342.41, 'pymatviz.cumulative': 2299.73, 'pymatviz.histogram': 2443.11, 'pymatviz.phonons': 2235.57, 'pymatviz.powerups': 2172.71, 'pymatviz.ptable': 2286.77, 'pymatviz.rainclouds': 2702.03, 'pymatviz.rdf': 2331.98, 'pymatviz.relevance': 2256.29, 'pymatviz.sankey': 2313.12, 'pymatviz.scatter': 2312.48, 'pymatviz.structure_viz': 2330.39, 'pymatviz.sunburst': 2395.04, 'pymatviz.uncertainty': 2317.87, 'pymatviz.xrd': 2242.09}

[janosh]

i saw that. the thing is, if you paste that printed dict into a python file, ruff will auto-format it to multiple lines so having easier to read code seemed more important

[DanielYang59]

Fair point!

[DanielYang59]

Thanks for catching this. using perf_counter_ns was intended to avoid float precision lost (the former gives time in int, while the latter gives in float), but admittedly for our case the error should be negligible as: 1. it's far from the ns level; 2. we're rounding it anyway

Use perf_counter_ns() to avoid the precision loss caused by the float type.