Skip to content

Commit

Permalink
New ModelManager class for managing model contexts
Browse files Browse the repository at this point in the history
We create a global instance of this model within this module.

We inherit from threading.local to ensure thread safety when working
with models on multiple threads. See pymc-devs#1552 for the reasoning. This is
already tested in `test_thread_safety`.
  • Loading branch information
thomasaarholt committed Aug 23, 2024
1 parent d47ca4c commit 69bfeac
Showing 1 changed file with 29 additions and 0 deletions.
29 changes: 29 additions & 0 deletions pymc/model/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,35 @@
]


class ModelManager(threading.local):
"""Keeps track of currently active model contexts.
A global instance of this is created in this module on import.
Use that instance, `MODEL_MANAGER` to inspect current contexts.
It inherits from threading.local so is thread-safe, if models
can be entered/exited within individual threads.
"""

def __init__(self):
self.active_contexts: list[Model] = []

@property
def current_context(self) -> Model | None:
"""Return the innermost context of any current contexts."""
return self.active_contexts[-1] if self.active_contexts else None

@property
def parent_context(self) -> Model | None:
"""Return the parent context to the active context, if any."""
return self.active_contexts[-2] if len(self.active_contexts) > 1 else None


# MODEL_MANAGER is instantiated at import, and serves as a truth for
# what any currently active model contexts are.
MODEL_MANAGER = ModelManager()


class ContextMeta(type):
"""Functionality for objects that put themselves in a context using
the `with` statement.
Expand Down

0 comments on commit 69bfeac

Please sign in to comment.