ezmsg-org / ezmsg

Pure-Python DAG-based high-performance SHM-backed pub-sub and multi-processing pattern
https://ezmsg.readthedocs.io/en/latest/
MIT License
15 stars 6 forks source link

Create and use fast_replace #155

Closed cboulay closed 1 week ago

griffinmilsap commented 1 week ago

Well, the performance doesn't lie; this is a pretty significant speedup on an operation we do a LOT.

Would it be possible to move replace to messages.util (or even messages.__init__) (instead of AxisArray), and make it possible to use a monkeypatch or possibly even an environment variable for all code to fall back to using original (slow) replace. Unless the idea is that replace is only to be used on AxisArrays, in which case it can remain in AxisArray and should be typehinted appropriately (with the TypeVar bound on AxisArray)

Additionally, the original dataclasses.py:replace has a check to deliver a helpful error message if its used on a non-dataclass object; I'll defer to you if that's a useful add here or not.

Once thats addressed, I'm happy to take this for a spin in a new tagged release on PyPi (with the new AxisArray w/ coords :D)

cboulay commented 1 week ago

move replace to messages.util

Done.

an environment variable for all code to fall back to using original (slow) replace

Done. from ezmsg.util.messages.util import replace will use fast_replace by default or it will use dataclasses.replace if "EZMSG_DISABLE_FAST_REPLACE" is in os.environ. User code still has the option to import fast_replace if they want to use that independent of the environment variable.

replace is only to be used on AxisArrays

The plan was to use it only on AxisArrays AND the different Axis classes. There is no single base class that captures them all so I couldn't bind the TypeVar anywhere. Now that it's moved, I suppose any dataclass can use it.

the original dataclasses.py:replace has a check

If we are foregoing checks for speed, I figure we might as well forego all checks.