logix-project / logix

AI Logging for Interpretability and Explainability🔬
Apache License 2.0
76 stars 6 forks source link

Major refactoring with shared states across handlers #42

Closed sangkeun00 closed 9 months ago

sangkeun00 commented 9 months ago

Previously, communication between different handlers was performed directly in each handler, which led to reduced modularity. Given that the content of most communication is various states (e.g. Hessian states, current log), we decided to create a new AnaLogState that stores all these relevant log states. All states in AnaLogState follow the nested dictionary data structure, and thus various operations such as synchronization are implemented in a recursive manner.

With this new design, AnaLog instantiates one AnaLogState and pass this for initializing all other handlers. When needed, each handler directly updates this state, and the updated information would be immediately available to all other handlers.

I am still not sure whether it's beneficial for LoggingHandler to directly and immediately pass the current log to StorageHandler (there may be some optimization opportunities with this approach), or wait until __exit__ and pass the aggregated log to only once. In this refactoring, I followed the later stage by additionally implementing add_on_exit in StorageHandler, but kept the original add in case we revert our strategy.