diff --git a/frigate/embeddings/maintainer.py b/frigate/embeddings/maintainer.py index c4f79d5736..042c2722c2 100644 --- a/frigate/embeddings/maintainer.py +++ b/frigate/embeddings/maintainer.py @@ -200,6 +200,9 @@ class EmbeddingMaintainer(threading.Thread): ) for model_config in self.config.classification.custom.values(): + if not model_config.enabled: + continue + self.realtime_processors.append( CustomStateClassificationProcessor( self.config, model_config, self.requestor, self.metrics @@ -332,6 +335,25 @@ class EmbeddingMaintainer(threading.Thread): for processor in self.post_processors: processor.update_config(topic, payload) + def _remove_custom_classification_processor(self, model_name: str) -> None: + """Shut down and drop any running processor for a custom model.""" + remaining = [] + for processor in self.realtime_processors: + if ( + isinstance( + processor, + ( + CustomStateClassificationProcessor, + CustomObjectClassificationProcessor, + ), + ) + and processor.model_config.name == model_name + ): + processor.shutdown() + else: + remaining.append(processor) + self.realtime_processors = remaining + def _handle_custom_classification_update( self, topic: str, model_config: Any ) -> None: @@ -339,23 +361,7 @@ class EmbeddingMaintainer(threading.Thread): model_name = topic.split("/")[-1] if model_config is None: - remaining = [] - for processor in self.realtime_processors: - if ( - isinstance( - processor, - ( - CustomStateClassificationProcessor, - CustomObjectClassificationProcessor, - ), - ) - and processor.model_config.name == model_name - ): - processor.shutdown() - else: - remaining.append(processor) - self.realtime_processors = remaining - + self._remove_custom_classification_processor(model_name) logger.info( f"Successfully removed classification processor for model: {model_name}" ) @@ -363,6 +369,13 @@ class EmbeddingMaintainer(threading.Thread): self.config.classification.custom[model_name] = model_config + # A disabled model must not run; tear down any existing processor and + # do not register a new one. + if not model_config.enabled: + self._remove_custom_classification_processor(model_name) + logger.info(f"Disabled classification processor for model: {model_name}") + return + # Check if processor already exists for processor in self.realtime_processors: if isinstance( @@ -702,7 +715,11 @@ class EmbeddingMaintainer(threading.Thread): and "license_plate" not in camera_config.objects.track ) - if not dedicated_lpr_enabled and len(self.config.classification.custom) == 0: + has_enabled_custom = any( + c.enabled for c in self.config.classification.custom.values() + ) + + if not dedicated_lpr_enabled and not has_enabled_custom: # no active features that use this data return diff --git a/frigate/test/test_classification_enabled.py b/frigate/test/test_classification_enabled.py new file mode 100644 index 0000000000..abb139cd08 --- /dev/null +++ b/frigate/test/test_classification_enabled.py @@ -0,0 +1,106 @@ +"""Tests that disabled custom classification models are not registered or run.""" + +import sys +import unittest +from unittest.mock import MagicMock + +# Mock TFLite before importing the maintainer / classification modules +_MOCK_MODULES = [ + "tflite_runtime", + "tflite_runtime.interpreter", + "ai_edge_litert", + "ai_edge_litert.interpreter", +] +for mod in _MOCK_MODULES: + if mod not in sys.modules: + sys.modules[mod] = MagicMock() + +from frigate.data_processing.real_time.custom_classification import ( # noqa: E402 + CustomObjectClassificationProcessor, +) +from frigate.embeddings.maintainer import EmbeddingMaintainer # noqa: E402 + + +class TestCustomClassificationEnabledGating(unittest.TestCase): + """A model with enabled: false must not keep a processor registered.""" + + def _make_maintainer(self) -> EmbeddingMaintainer: + # Bypass the heavy __init__; only the attributes touched by the + # config update path are needed for these tests. + maintainer = EmbeddingMaintainer.__new__(EmbeddingMaintainer) + maintainer.realtime_processors = [] + maintainer.config = MagicMock() + maintainer.config.classification.custom = {} + maintainer.requestor = MagicMock() + maintainer.metrics = MagicMock() + maintainer.event_metadata_publisher = MagicMock() + return maintainer + + def _make_model_config(self, name: str, enabled: bool) -> MagicMock: + model_config = MagicMock() + model_config.name = name + model_config.enabled = enabled + model_config.state_config = None + return model_config + + def _make_processor(self, name: str) -> MagicMock: + processor = MagicMock(spec=CustomObjectClassificationProcessor) + processor.model_config = MagicMock() + processor.model_config.name = name + return processor + + def test_disabled_update_tears_down_existing_processor(self): + """Toggling a running model to disabled shuts down and drops its processor.""" + maintainer = self._make_maintainer() + processor = self._make_processor("atli") + maintainer.realtime_processors = [processor] + + maintainer._handle_custom_classification_update( + "config/classification/custom/atli", + self._make_model_config("atli", enabled=False), + ) + + processor.shutdown.assert_called_once() + self.assertEqual(maintainer.realtime_processors, []) + + def test_disabled_update_does_not_register_processor(self): + """A disabled model that has no processor is never registered.""" + maintainer = self._make_maintainer() + + maintainer._handle_custom_classification_update( + "config/classification/custom/atli", + self._make_model_config("atli", enabled=False), + ) + + self.assertEqual(maintainer.realtime_processors, []) + + def test_disabled_update_leaves_other_processors_untouched(self): + """Disabling one model must not affect other running processors.""" + maintainer = self._make_maintainer() + other = self._make_processor("simbi") + maintainer.realtime_processors = [other] + + maintainer._handle_custom_classification_update( + "config/classification/custom/atli", + self._make_model_config("atli", enabled=False), + ) + + other.shutdown.assert_not_called() + self.assertEqual(maintainer.realtime_processors, [other]) + + def test_removed_model_tears_down_processor(self): + """A None payload (model deleted) still shuts down its processor.""" + maintainer = self._make_maintainer() + processor = self._make_processor("atli") + maintainer.realtime_processors = [processor] + + maintainer._handle_custom_classification_update( + "config/classification/custom/atli", None + ) + + processor.shutdown.assert_called_once() + self.assertEqual(maintainer.realtime_processors, []) + + +if __name__ == "__main__": + unittest.main()