fix(parser): memory leak in cache

 - a global cache was used and never cleared for tokens/parser results, which ended up creating a slow memory leak that eventually caused training to slow down (after ~2000 eps)

Author: justindujardin

libraries/mathy_python/mathy/core/parser.py

@@ -103,9 +103,6 @@ def contains(self, type: int) -> bool:
 _IS_EXP: TokenSet = TokenSet(TokenExponent)
 _IS_EQUAL: TokenSet = TokenSet(TokenEqual)
 
-_parse_cache: Dict[str, MathExpression] = {}
-_tokens_cache: Dict[str, List[Token]] = {}
-
 
 # NOTE: This cannot be shared between threads because it stores state in self.current_token and self.tokens
 class ExpressionParser:
@@ -146,27 +143,33 @@ class ExpressionParser:
     ```
     """
 
+    _parse_cache: Dict[str, MathExpression]
+    _tokens_cache: Dict[str, List[Token]]
+
     # Initialize the tokenizer.
     def __init__(self):
         self.tokenizer = Tokenizer()
+        self.clear_cache()
+
+    def clear_cache(self):
+        self._tokens_cache = {}
+        self._parse_cache = {}
 
     def tokenize(self, input_text: str):
-        global _tokens_cache
-        if input_text not in _tokens_cache:
-            _tokens_cache[input_text] = self.tokenizer.tokenize(input_text)
-        return _tokens_cache[input_text][:]
+        if input_text not in self._tokens_cache:
+            self._tokens_cache[input_text] = self.tokenizer.tokenize(input_text)
+        return self._tokens_cache[input_text][:]
 
     def parse(self, input_text: str) -> MathExpression:
         """Parse a string representation of an expression into a tree
         that can be later evaluated.
 
         Returns : The evaluatable expression tree.
         """
-        global _parse_cache
-        if input_text in _parse_cache:
-            return _parse_cache[input_text]
-        _parse_cache[input_text] = self._parse(self.tokenize(input_text))
-        return _parse_cache[input_text]
+        if input_text in self._parse_cache:
+            return self._parse_cache[input_text]
+        self._parse_cache[input_text] = self._parse(self.tokenize(input_text))
+        return self._parse_cache[input_text]
 
     def _parse(self, tokens: List[Token]) -> MathExpression:
         """Parse a given list of tokens into an expression tree"""

libraries/mathy_python/mathy/env.py

@@ -360,6 +360,7 @@ def get_initial_state(
         prob: MathyEnvProblem = self.problem_fn(config)
         self.valid_actions_mask_cache = dict()
         self.valid_rules_cache = dict()
+        self.parser.clear_cache()
         self.max_moves = self.max_moves_fn(prob, config)
 
         # Build and return the initial state