60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
"""
|
|
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()
|