make shared_intermediates thread-safe #53
Conversation
jcmgray
changed the title
opt_einsum/sharing.py
Outdated
| def currently_sharing(): | ||
| """Check if we are currently sharing a cache -- thread specific. | ||
| """ | ||
| return threading.get_ident() in _SHARING_STACK |
There was a problem hiding this comment.
Hmm in Python 2, threading.get_ident does not exist
There was a problem hiding this comment.
Yes, just realized that! I'll have a quick look for an alternative.
opt_einsum/sharing.py
Outdated
|
|
||
| from .parser import alpha_canonicalize, parse_einsum_input | ||
|
|
||
| _SHARING_STACK = [] | ||
|
|
||
| _SHARING_STACK = {} |
There was a problem hiding this comment.
Consider making this a defaultdict(list) to avoid the setdefault complexity below?
opt_einsum/sharing.py
Outdated
|
|
||
| def _remove_sharing_cache(cache): | ||
| tid = threading.get_ident() | ||
| _SHARING_STACK[tid].remove(cache) |
There was a problem hiding this comment.
What is the advantage of .remove() over .pop()? If these are really stacks, then we should use .pop(). If these are not really stacks, then we should use sets or something.
There was a problem hiding this comment.
Yes should be exactly the same for the current behaviour, remove(cache) felt more explicit but is probably unnecessary, I'll update it.
|
Thanks for those comments @fritzo! I've made those changes and added a test compatible with both py2/py3. Ready to go from my end. |
|
It doesn't happen automatically, but threads sharing a cache is already possible if you explicitly supply the cache to use (they'll each add it their own 'stack', but will all still be referencing the same object). |
|
Good point, currently doing a docs pass and I will add that. |
Description
Makes
shared_intermediatesthread safe, by storing the caches in a dict referenced by thread id. All threads can (still) share the same cache among themselves, if it is supplied to them all. Resolves #51. cc @fritzo.Todos
Status