Memory leak in tag processor (#24984)

* Cache `load_nlp_engine`

* Explicit dependency from `TagProcessor` on `SpacyNlpEngine`

Author: edg956

ingestion/src/metadata/pii/algorithms/presidio_utils.py

@@ -13,6 +13,7 @@
 """
 import inspect
 import logging
+from functools import cache
 from typing import Any, Callable, Iterable, List, Optional, Set, Type, Union, cast
 
 import spacy
@@ -43,6 +44,7 @@
 logger = pii_logger()
 
 
+@cache
 def load_nlp_engine(
     model_name: str = SPACY_EN_MODEL,
     supported_language: str = SUPPORTED_LANG,

ingestion/src/metadata/pii/tag_processor.py

@@ -12,6 +12,7 @@
     TagSource,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.pii.algorithms.presidio_utils import load_nlp_engine
 from metadata.pii.algorithms.tag_scoring import ScoreTagsForColumnService
 from metadata.pii.base_processor import AutoClassificationProcessor
 from metadata.pii.classification_manager import (
@@ -68,7 +69,9 @@ def __init__(
 
         # Service that runs analyzers
         if score_tags_for_column is None:
-            score_tags_for_column = ScoreTagsForColumnService()
+            score_tags_for_column = ScoreTagsForColumnService(
+                nlp_engine=load_nlp_engine()
+            )
         self.score_tags_for_column = score_tags_for_column
 
         logger.info(

ingestion/tests/unit/pii/algorithms/test_presidio_utils.py

@@ -8,7 +8,7 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import pytest
 from presidio_analyzer import EntityRecognizer, RecognizerResult
@@ -17,6 +17,7 @@
 from metadata.pii.algorithms.presidio_utils import (
     apply_confidence_threshold,
     build_analyzer_engine,
+    load_nlp_engine,
     set_presidio_logger_level,
 )
 from metadata.pii.algorithms.tags import PIITag
@@ -138,3 +139,88 @@ def test_threshold_of_zero_returns_all_results(self, mock_recognizer):
         )
 
         assert len(results) == 3
+
+
+@patch("metadata.pii.algorithms.presidio_utils._load_spacy_model")
+@patch("metadata.pii.algorithms.presidio_utils.SpacyNlpEngine")
+class TestLoadNlpEngine:
+    @staticmethod
+    def setup_method():
+        """Clear the cache before each test"""
+        load_nlp_engine.cache_clear()
+
+    @staticmethod
+    def teardown_method():
+        """Clear the cache after each test"""
+        load_nlp_engine.cache_clear()
+
+    def test_returns_same_instance_for_same_parameters(
+        self, mock_spacy_engine_class, mock_load_spacy
+    ):
+        """Test that calling load_nlp_engine with same parameters returns same instance"""
+        mock_engine = Mock()
+        mock_spacy_engine_class.return_value = mock_engine
+
+        engine1 = load_nlp_engine(model_name="en_core_web_sm", supported_language="en")
+        engine2 = load_nlp_engine(model_name="en_core_web_sm", supported_language="en")
+
+        assert engine1 is engine2
+        assert mock_spacy_engine_class.call_count == 1
+        assert mock_load_spacy.call_count == 1
+
+    def test_returns_different_instances_for_different_model_names(
+        self, mock_spacy_engine_class, mock_load_spacy
+    ):
+        """Test that different model names result in different instances"""
+        mock_engine1 = Mock()
+        mock_engine2 = Mock()
+        mock_spacy_engine_class.side_effect = [mock_engine1, mock_engine2]
+
+        engine1 = load_nlp_engine(model_name="en_core_web_sm", supported_language="en")
+        engine2 = load_nlp_engine(model_name="en_core_web_md", supported_language="en")
+
+        assert engine1 is not engine2
+        assert mock_spacy_engine_class.call_count == 2
+        assert mock_load_spacy.call_count == 2
+
+    def test_returns_different_instances_for_different_languages(
+        self, mock_spacy_engine_class, mock_load_spacy
+    ):
+        """Test that different languages result in different instances"""
+        mock_engine1 = Mock()
+        mock_engine2 = Mock()
+        mock_spacy_engine_class.side_effect = [mock_engine1, mock_engine2]
+
+        engine1 = load_nlp_engine(model_name="en_core_web_sm", supported_language="en")
+        engine2 = load_nlp_engine(model_name="en_core_web_sm", supported_language="fr")
+
+        assert engine1 is not engine2
+        assert mock_spacy_engine_class.call_count == 2
+
+    def test_cache_persists_across_multiple_calls(
+        self, mock_spacy_engine_class, mock_load_spacy
+    ):
+        """Test that cache works correctly across multiple calls"""
+        mock_engine = Mock()
+        mock_spacy_engine_class.return_value = mock_engine
+
+        engine1 = load_nlp_engine(model_name="en_core_web_sm", supported_language="en")
+        engine2 = load_nlp_engine(model_name="en_core_web_sm", supported_language="en")
+        engine3 = load_nlp_engine(model_name="en_core_web_sm", supported_language="en")
+
+        assert engine1 is engine2 is engine3
+        assert mock_spacy_engine_class.call_count == 1
+        assert mock_load_spacy.call_count == 1
+
+    def test_uses_default_parameters_when_not_provided(
+        self, mock_spacy_engine_class, mock_load_spacy
+    ):
+        """Test that default parameters work correctly with caching"""
+        mock_engine = Mock()
+        mock_spacy_engine_class.return_value = mock_engine
+
+        engine1 = load_nlp_engine()
+        engine2 = load_nlp_engine()
+
+        assert engine1 is engine2
+        assert mock_spacy_engine_class.call_count == 1