cache __TYPE_MAP and init submodules

[sneakers-the-rat]

Related to:
fix: #1255
fix: #1487
prior pr: #1050

Motivation

import slow

How to test the behavior?

import pynwb
# check ur watch meaningfully

Checklist

I'VE DONE NONE OF THESE AND I CAN EXPLAIN

• Did you update CHANGELOG.md with your changes?
• Have you checked our Contributing document?
• Have you ensured the PR clearly describes the problem and the solution?
• Is your contribution compliant with our coding style? This can be checked running flake8 from the source directory.
• Have you checked to ensure that there aren't other open Pull Requests for the same change?
• Have you included the relevant issue number using "Fix #XXX" notation where XXX is the issue number? By including "Fix #XXX" you allow GitHub to close issue #XXX when the PR is merged.

Narrative

sorry for the incomplete PR, i have to run out the door RIGHT NOW and didn't want to forget. opening as a draft and i'll come back later.

finishing up nwb_linkml and wanted to read through to see if there was any special behavior i was missing and was like why the heck does the import take like an hour.

on first import it's because hdmf imports the whole scipy sparse module and pynwb imports down that far in hdmf. that would be fixable using a TYPE_CHECKING check, but idk if that function signature decorator thing can handle delayed annotations like that. that's nbd it goes away.

most of the time is spent in the initial loading of the schema into __TYPE_MAP (and most of that is from probably unnecessary deepcopys, but that's another day's perf to win), with the remaining ~1/3 from the side effects of importing the other modules.

So i just mostly did what was talked about in the prior issues and pickled the __TYPE_MAP object before importing the rest of the modules bc i wasn't sure what the side effects were there, but if that's fine to also pickle then we could just copy/paste the module-level imports into the load_namespaces function.

I also was like "why the heck don't ya try to clone the submodules if you're gonna tell me i need to do it anyway" when i first imported so i also did that just bc it seems like a kind thing to do.

anyway more later and i'll test this too xoxo

[sneakers-the-rat]

(lmao i must hear the story here)

[sneakers-the-rat]

OK i tested whether importing all the submodules first before pickling sped up the import further, and it doesn't. That version is here: sneakers-the-rat@0a1514e

The profiling results on my computer (obviously variable) are:

• Before caching: 4.33s
• After caching: 1.05s
• Improvement: 3.28s (75%)

The reason that including the imported modules doesn't improve the speed much if any is that the main contributor to the slow imports is that a) there are import side effects, b) those import side effects include parsing yaml c) parsing yaml is slow. The module imports take very little time because the yaml is already parsed by the time they are imported, so imo those should remain uncached because unpickling objects with declaration-time side effects (the registration decorator) is super wonky and fragile, u can see the hack i had to do in that branch.

There are three remaining opportunities for future import perf improvements:

• hdmf basically duplicates the import side effects of pynwb, so a similar thing could be done there for an additional ~300ms off
• the YAMLSpecReader in hdmf uses the slower Python-based yaml parser because of the pure=True kwarg here: https://github.com/hdmf-dev/hdmf/blob/dfb1df79bce3e34c52af7996429c3f561d96390a/src/hdmf/spec/namespace.py#L203
  as far as I'm aware all the NWB schemas are yaml 1.1 compliant, so i am not sure there is a reason for doing that. Removing that and using the C-based loader saves an additional 150ms after caching.
• having all these import side effects seems to have lent itself to having a very big import tree - importing pynwb imports 765 modules (numpy and scipy, two very heavy packages, import 114 and 124 respectively). Fixing this would be a more involved refactoring where logic would have to be moved out of __init__.py files, import side effects would have to be turned into on demand loading, and then the package level imports could be pruned down. The remaining ~700ms of time is just the normal amount of time it takes to import 765 modules.

I'm marking this as ready for review without adding tests just to see if this approach would be acceptable, and if so then i'll add tests for the submodule cloning, caching, and cache invalidation.

[codecov]

Codecov Report

Attention: Patch coverage is 67.50000% with 13 lines in your changes missing coverage. Please review.

Project coverage is 91.60%. Comparing base (570fb3b) to head (070cb66).

Files	Patch %	Lines
src/pynwb/__init__.py	67.50%	11 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##              dev    #1931      +/-   ##
==========================================
- Coverage   91.91%   91.60%   -0.32%     
==========================================
  Files          27       27              
  Lines        2648     2681      +33     
  Branches      691      699       +8     
==========================================
+ Hits         2434     2456      +22     
- Misses        141      151      +10     
- Partials       73       74       +1     

Flag	Coverage Δ
integration	72.32% <50.00%> (-0.49%)	⬇️
unit	83.51% <50.00%> (-0.48%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[sneakers-the-rat]

Huh, so the test that's failing is this one

pynwb/tests/back_compat/test_import_structure.py

Line 6 in 570fb3b

class TestImportStructure(TestCase):

which makes sense because I removed __resources from the module-level scope. I figured that would be OK since it is a dunder attribute, a "really private, really don't try and access these from outside the scope" name, and so i'm sort of surprised that they would be in that API consistency test?

Lmk what would be better, updating that test (i think that none of those dunder attributes should be in there?) or making it so __resources is back in the module scope. I think that might just be a code style thing, in general i try and avoid using module level globals, esp in this case where it's a single use, quickly computed thing, but it would really be no problem to put it back in module globals.

[sneakers-the-rat]

OK moved __resources back to the module scope, and I added a simple test for the submodule cloning - incidentally that requires one to remove the cache files, so it implicitly functions as a test for the creation but not the use of the cache files. I am not sure the best way to test that the cache file is actually used - i could do something like emit a warning there, but that seems noisy, or i could add some kind of env var flag, but that also seems inelegant, so lmk what y'all would like there.

during testing i discovered that one also needs to cache __NS_CATALOG, which makes sense because it looks like one wraps the other. I restored the global version from the one that is saved in the cached __TYPE_MAP so we only need one cache file. Added a very simple test that their id matches.

One more thing is that while linting i noticed that the ruff excludes path points to the wrong directory for nwb-schema so i adjusted that too.

[sneakers-the-rat]

gee for a simple thing this sure does involve a lot of fighting lol

OK so i had written some tests for this but since we're talking about import side effects testing them is hard.

• it turns out that __TYPE_MAP also stores the generated classes, so a test for import effects of the top-level pynwb fails because on reimport the references to the objects are regenerated, so then isinstance() and equality checks fail and those seem to cascade through hdmf in a way i can't really tell where the problem comes in. so anyway that breaks the rest of the tests
• calling the _load_core_namespace manually to trigger a creation of the cache file again doesn't work because by the time the import tests run there has been an example extension test run, but it doesn't clean up after itself in the type map so there is an unpickleable class (because it doesn't have a namespace, it's (TestExtension.test_lab_meta.<locals>.MyTestMetaData), whomp whomp.
• the package is not installed in editable mode from tox, so it's no longer in a git repo when the tests run, so testing cloning submodules fails there (but succeeds when run from the repo)

fortunately....

• the tests all run and pass when using the cached TYPE_MAP, so those import/symbol identity inconsistencies are not true for the pickled file, presumably because unpickling runs the code they come from and the names are bound then.
• the submodule cloning code has the same error path - if the submodule directory doesn't exist, throw error. so it's at worse net neutral in terms of bugs (actually a little bit better - if you delete the schema directory the type map still loads correctly),

so I just scrapped the tests because this is ultimately a very simple change and i did not budget this much time to this very simple change.

if that's not gonna cut it feel free to close this, just trying to do a little QoL contribs while i work with this package so not mission critical, fine for me to just have it as a fork.

[stephprince]

Hi @sneakers-the-rat, thanks for all the work and profiling! Really appreciate the walkthrough.

@rly said he can take a look at this so I'll let him review in detail.