File size: 1,878 Bytes
59f1501 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
"""
This module provides thread-safe code context management for TorchDynamo using weak references.
The CodeContextDict class maintains a mapping between Python code objects and their associated
context data, using weak references to automatically clean up entries when code objects are
garbage collected. This prevents memory leaks while allowing context data to be associated
with code objects throughout their lifecycle.
Key features:
- Thread-safe context storage and retrieval
- Automatic cleanup using weak references
- Safe context management for Python code objects
- Memory-leak prevention
Example usage:
code_obj = compile('x = 1', '<string>', 'exec')
# Store context
context = code_context.get_context(code_obj)
context['metadata'] = {'optimized': True}
# Retrieve context
if code_context.has_context(code_obj):
ctx = code_context.get_context(code_obj)
# Use context data...
# Remove context
ctx = code_context.pop_context(code_obj)
"""
import types
from typing import Any
from .utils import ExactWeakKeyDictionary
class CodeContextDict:
def __init__(self) -> None:
self.code_context: ExactWeakKeyDictionary = ExactWeakKeyDictionary()
def has_context(self, code: types.CodeType) -> bool:
return code in self.code_context
def get_context(self, code: types.CodeType) -> dict[str, Any]:
ctx = self.code_context.get(code)
if ctx is None:
ctx = {}
self.code_context[code] = ctx
return ctx
def pop_context(self, code: types.CodeType) -> dict[str, Any]:
ctx = self.get_context(code)
self.code_context._remove_id(id(code))
return ctx
def clear(self) -> None:
self.code_context.clear()
code_context: CodeContextDict = CodeContextDict()
|