🔧 chore(inspect): add early stopping to training job (#3217)

* Add early stopping

Signed-off-by: Ashwin Vaidya <[email protected]>

* Fix test

Signed-off-by: Ashwin Vaidya <[email protected]>

---------

Signed-off-by: Ashwin Vaidya <[email protected]>

Author: ashwinvaidya17

application/backend/src/services/training_service.py

@@ -10,7 +10,10 @@
 from anomalib.deploy import ExportType
 from anomalib.engine import Engine
 from anomalib.loggers import AnomalibTensorBoardLogger
+from anomalib.metrics import AUROC, F1Score
+from anomalib.metrics.evaluator import Evaluator
 from anomalib.models import get_model
+from lightning.pytorch.callbacks import EarlyStopping
 from loguru import logger
 
 from pydantic_models import Job, JobStatus, JobType, Model
@@ -203,7 +206,18 @@ def _train_model(
         )
 
         # Initialize anomalib model and engine
-        anomalib_model = get_model(model=model.name)
+        anomalib_model = get_model(
+            model=model.name,
+            evaluator=Evaluator(
+                val_metrics=[AUROC(fields=["anomaly_map", "gt_mask"], prefix="pixel_", strict=False)],
+                test_metrics=[
+                    AUROC(fields=["pred_score", "gt_label"], prefix="image_"),
+                    F1Score(fields=["pred_label", "gt_label"], prefix="image_"),
+                    AUROC(fields=["anomaly_map", "gt_mask"], prefix="pixel_", strict=False),
+                    F1Score(fields=["pred_mask", "gt_mask"], prefix="pixel_", strict=False),
+                ],
+            ),
+        )
 
         trackio = TrackioLogger(project=str(model.project_id), name=model.name)
         tensorboard = AnomalibTensorBoardLogger(save_dir=global_log_config.tensorboard_log_path, name=name)
@@ -212,7 +226,10 @@ def _train_model(
             logger=[trackio, tensorboard],
             devices=[0],  # Only single GPU training is supported for now
             max_epochs=max_epochs,
-            callbacks=[GetiInspectProgressCallback(synchronization_parameters)],
+            callbacks=[
+                GetiInspectProgressCallback(synchronization_parameters),
+                EarlyStopping(monitor="pixel_AUROC", mode="max", patience=5),
+            ],
             accelerator=training_device,
         )
 

application/backend/tests/unit/services/test_training_service.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import asyncio
 from typing import Any
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
 from uuid import uuid4
 
 import pytest
@@ -274,7 +274,7 @@ def test_train_model_success(
 
         # Verify all components were called correctly
         fxt_mock_anomalib_components["folder_class"].assert_called_once()
-        fxt_mock_anomalib_components["get_model"].assert_called_once_with(model=fxt_model.name)
+        fxt_mock_anomalib_components["get_model"].assert_called_once_with(model=fxt_model.name, evaluator=ANY)
 
         # Verify Engine was called with expected parameters
         fxt_mock_anomalib_components["engine_class"].assert_called_once()