Fix enabled flag for custom classification models (#23681)
Some checks are pending
CI / AMD64 Build (push) Waiting to run
CI / ARM Build (push) Waiting to run
CI / Jetson Jetpack 6 (push) Waiting to run
CI / AMD64 Extra Build (push) Blocked by required conditions
CI / ARM Extra Build (push) Blocked by required conditions
CI / Synaptics Build (push) Blocked by required conditions
CI / Assemble and push default build (push) Blocked by required conditions

* honor enabled flag for custom classification models

for both startup and dynamically, even though the UI doesn't currently have a way to toggle dynamically

* add test
This commit is contained in:
Josh Hawkins 2026-07-12 06:48:14 -05:00 committed by GitHub
parent da4037eb52
commit 5f6043aa92
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 141 additions and 18 deletions

View File

@ -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

View File

@ -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()