Apply filtering parameters (#1046)

Co-authored-by: Copilot <[email protected]>

Author: maxxgx

interactive_ai/services/director/app/service/job_submission/job_creation_helpers.py

@@ -14,6 +14,7 @@
 
 from communication.exceptions import JobCreationFailedException
 from features.feature_flag import FeatureFlag
+from service.configuration_service import ConfigurationService
 
 from geti_types import ID, ProjectIdentifier
 from iai_core.entities.dataset_storage import DatasetStorage
@@ -51,6 +52,56 @@ class JobName(str, Enum):
     MODEL_TEST = "Model testing"
 
 
+@dataclass
+class FilteringParameters:
+    """
+    Helper class to hold filtering parameters for optimize/test jobs.
+    """
+
+    min_annotation_size: int | None
+    max_annotation_size: int | None
+    min_number_of_annotations: int | None
+    max_number_of_annotations: int | None
+
+
+def _get_filtering_parameters(project: Project, model: Model) -> FilteringParameters:
+    """
+    Returns filtering parameters for training jobs based on the project and model.
+
+    :param project: Project containing the task node
+    :param model: Model to be trained
+    :return: FilteringParameters instance with the filtering parameters
+    """
+    training_config = ConfigurationService.get_full_training_configuration(
+        project_identifier=project.identifier,
+        task_id=model.model_storage.task_node_id,
+        model_manifest_id=model.model_storage.model_manifest_id,
+    )
+    filtering_params = training_config.global_parameters.dataset_preparation.filtering
+    return FilteringParameters(
+        min_annotation_size=(
+            filtering_params.min_annotation_pixels.min_annotation_pixels
+            if filtering_params.min_annotation_pixels and filtering_params.min_annotation_pixels.enable
+            else None
+        ),
+        max_annotation_size=(
+            filtering_params.max_annotation_pixels.max_annotation_pixels
+            if filtering_params.max_annotation_pixels and filtering_params.max_annotation_pixels.enable
+            else None
+        ),
+        min_number_of_annotations=(
+            filtering_params.min_annotation_objects.min_annotation_objects
+            if filtering_params.min_annotation_objects and filtering_params.min_annotation_objects.enable
+            else None
+        ),
+        max_number_of_annotations=(
+            filtering_params.max_annotation_objects.max_annotation_objects
+            if filtering_params.max_annotation_objects and filtering_params.max_annotation_objects.enable
+            else None
+        ),
+    )
+
+
 def get_model_storage_for_task(
     project_identifier: ProjectIdentifier,
     task_node_id: ID,
@@ -239,11 +290,14 @@ def create_payload(
 
         :returns: a dict representing the job payload
         """
+        filter_parameters = _get_filtering_parameters(project=self.project, model=self.model)
         return {
             "project_id": self.project.id_,
             "model_test_result_id": self.model_test_result.id_,
-            "min_annotation_size": self.min_annotation_size,
-            "max_number_of_annotations": self.max_number_of_annotations,
+            "min_annotation_size": filter_parameters.min_annotation_size,
+            "max_annotation_size": filter_parameters.max_annotation_size,
+            "min_number_of_annotations": filter_parameters.min_number_of_annotations,
+            "max_number_of_annotations": filter_parameters.max_number_of_annotations,
         }
 
     def create_metadata(self) -> dict:
@@ -285,9 +339,6 @@ class OptimizationJobData:
     training_dataset_storage: DatasetStorage
     model: Model
     optimization_type: ModelOptimizationType
-    optimization_parameters: dict
-    min_annotation_size: int | None = None
-    max_number_of_annotations: int | None = None
 
     @property
     def job_type(self) -> str:
@@ -308,7 +359,6 @@ def create_key(self) -> str:
             "dataset_storage_id": self.training_dataset_storage.id_,
             "model_storage_id": self.model.model_storage.id_,
             "type": self.job_type,
-            "optimization_parameters": self.optimization_parameters,
         }
         return _serialize_job_key(job_key)
 
@@ -318,6 +368,7 @@ def create_payload(self) -> dict:
 
         :return: dict with the job's metadata
         """
+        filter_parameters = _get_filtering_parameters(project=self.project, model=self.model)
         return {
             "project_id": self.project.id_,
             "dataset_storage_id": self.training_dataset_storage.id_,
@@ -327,8 +378,10 @@ def create_payload(self) -> dict:
             # However, leaving this flag itself can make it easier to debug the future problem (e.g., CVS-142877).
             # You can still launch the job with `enable_optimize_from_dataset_shard=False` from Flyte console.
             "enable_optimize_from_dataset_shard": True,
-            "min_annotation_size": self.min_annotation_size,
-            "max_number_of_annotations": self.max_number_of_annotations,
+            "min_annotation_size": filter_parameters.min_annotation_size,
+            "max_annotation_size": filter_parameters.max_annotation_size,
+            "min_number_of_annotations": filter_parameters.min_number_of_annotations,
+            "max_number_of_annotations": filter_parameters.max_number_of_annotations,
         }
 
     def create_metadata(self) -> dict:

interactive_ai/services/director/app/service/job_submission/optimize.py

@@ -1,20 +1,13 @@
 # Copyright (C) 2022-2025 Intel Corporation
 # LIMITED EDGE SOFTWARE DISTRIBUTION LICENSE
 
-from typing import TYPE_CHECKING, cast
 
 from service.job_submission.base import JobParams, ModelJobSubmitter
 from service.job_submission.job_creation_helpers import OPTIMIZE_JOB_PRIORITY, JobDuplicatePolicy, OptimizationJobData
 
 from geti_types import ID
-from iai_core.configuration.elements.component_parameters import ComponentType
-from iai_core.configuration.elements.dataset_manager_parameters import DatasetManagementConfig
 from iai_core.entities.model import Model, ModelOptimizationType
 from iai_core.entities.project import Project
-from iai_core.repos import ConfigurableParametersRepo
-
-if TYPE_CHECKING:
-    from iai_core.configuration.elements.optimization_parameters import POTOptimizationParameters
 
 
 class ModelOptimizationJobSubmitter(ModelJobSubmitter):
@@ -34,32 +27,12 @@ def prepare_data(  # type: ignore
         :param author: ID of the user submitting the job
         :return: ID of the optimization job that has been submitted.
         """
-        optimization_params = self._get_model_optimization_parameters(model=model)
-        config_repo = ConfigurableParametersRepo(project.identifier)
-        dataset_manager_config = config_repo.get_or_create_component_parameters(
-            data_instance_of=DatasetManagementConfig,
-            component=ComponentType.PIPELINE_DATASET_MANAGER,
-        )
-        min_annotation_size = (
-            None
-            if dataset_manager_config.minimum_annotation_size == -1
-            else dataset_manager_config.minimum_annotation_size
-        )
-        max_number_of_annotations = (
-            None
-            if dataset_manager_config.maximum_number_of_annotations == -1
-            else dataset_manager_config.maximum_number_of_annotations
-        )
-
         optimization_job_data = OptimizationJobData(
             workspace_id=project.workspace_id,
             project=project,
             training_dataset_storage=project.get_training_dataset_storage(),
             model=model,
             optimization_type=optimization_type,
-            optimization_parameters=optimization_params,
-            min_annotation_size=min_annotation_size,
-            max_number_of_annotations=max_number_of_annotations,
         )
         return JobParams(
             priority=OPTIMIZE_JOB_PRIORITY,
@@ -74,18 +47,3 @@ def prepare_data(  # type: ignore
             gpu_num_required=optimization_job_data.gpu_num_required,
             cancellable=True,
         )
-
-    @staticmethod
-    def _get_model_optimization_parameters(model: Model) -> dict:
-        """
-        Return the model optimization parameters as a dict.
-
-        :param model: the model to extract the optimization parameters from
-        :returns: the optimization parameters for the model
-        """
-        pot_parameters = cast(
-            "POTOptimizationParameters",
-            model.configuration.configurable_parameters.pot_parameters,  # type: ignore[attr-defined]
-        )
-
-        return {"stat_subset_size": pot_parameters.stat_subset_size}

interactive_ai/services/director/app/service/job_submission/test.py

@@ -4,11 +4,8 @@
 from service.job_submission.job_creation_helpers import MODEL_TEST_JOB_PRIORITY, JobDuplicatePolicy, ModelTestJobData
 
 from geti_types import ID
-from iai_core.configuration.elements.component_parameters import ComponentType
-from iai_core.configuration.elements.dataset_manager_parameters import DatasetManagementConfig
 from iai_core.entities.model_test_result import ModelTestResult
 from iai_core.entities.project import Project
-from iai_core.repos import ConfigurableParametersRepo
 
 
 class ModelTestingJobSubmitter(ModelJobSubmitter):
@@ -24,21 +21,6 @@ def prepare_data(  # type: ignore
         :return: ID of the model test job that has been submitted to the jobs client
         """
         model = model_test_result.get_model()
-        config_repo = ConfigurableParametersRepo(project.identifier)
-        dataset_manager_config = config_repo.get_or_create_component_parameters(
-            data_instance_of=DatasetManagementConfig,
-            component=ComponentType.PIPELINE_DATASET_MANAGER,
-        )
-        min_annotation_size = (
-            None
-            if dataset_manager_config.minimum_annotation_size == -1
-            else dataset_manager_config.minimum_annotation_size
-        )
-        max_number_of_annotations = (
-            None
-            if dataset_manager_config.maximum_number_of_annotations == -1
-            else dataset_manager_config.maximum_number_of_annotations
-        )
         model_testing_job_data = ModelTestJobData(
             model_test_result=model_test_result,
             model=model,
@@ -49,8 +31,6 @@ def prepare_data(  # type: ignore
             task_node=next(
                 task for task in project.get_trainable_task_nodes() if task.id_ == model.model_storage.task_node_id
             ),
-            min_annotation_size=min_annotation_size,
-            max_number_of_annotations=max_number_of_annotations,
         )
 
         return JobParams(

interactive_ai/services/director/tests/integration/communication/test_integration_optimization_controller.py

@@ -6,6 +6,7 @@
 
 from communication.controllers.optimization_controller import OptimizationController
 from communication.exceptions import NotReadyForOptimizationException
+from service.configuration_service import ConfigurationService
 
 from geti_types import ID
 from grpc_interfaces.job_submission.client import GRPCJobsClient
@@ -16,7 +17,14 @@
 
 
 class TestOptimizationController:
-    def test_start_pot_job(self, fxt_db_project_service, fxt_pot_hyperparameters, fxt_mock_jobs_client) -> None:
+    def test_start_pot_job(
+        self,
+        fxt_db_project_service,
+        fxt_pot_hyperparameters,
+        fxt_mock_jobs_client,
+        fxt_training_configuration_task_level,
+        fxt_global_parameters,
+    ) -> None:
         """
         <b>Description:</b>
         Test POT optimization job creation
@@ -50,8 +58,17 @@ def test_start_pot_job(self, fxt_db_project_service, fxt_pot_hyperparameters, fx
         model.model_storage.model_template.entrypoints = EntryPoints(base="Base interface", nncf="NNCF interface")
         model_repo = ModelRepo(model.model_storage_identifier)
         model_repo.save(model)
+        filtering_params = fxt_global_parameters.dataset_preparation.filtering
+
         # Act
-        with patch.object(GRPCJobsClient, "submit", return_value=None):
+        with (
+            patch.object(GRPCJobsClient, "submit", return_value=None),
+            patch.object(
+                ConfigurationService,
+                "get_full_training_configuration",
+                return_value=fxt_training_configuration_task_level,
+            ) as mock_get_full_training_configuration,
+        ):
             result = OptimizationController().start_optimization(
                 project_id=project.id_,
                 model_storage_id=model.model_storage.id_,
@@ -60,13 +77,17 @@ def test_start_pot_job(self, fxt_db_project_service, fxt_pot_hyperparameters, fx
             )
 
         # Assert
+        mock_get_full_training_configuration.assert_called_once_with(
+            project_identifier=project.identifier,
+            task_id=model.model_storage.task_node_id,
+            model_manifest_id=model.model_storage.model_manifest_id,
+        )
         fxt_mock_jobs_client._jobs_client.submit.assert_called_once_with(
             priority=1,
             job_name="Optimization",
             job_type="optimize_pot",
             key=f'{{"dataset_storage_id": "{dataset_storage.id_}",'
             f' "model_storage_id": "{model.model_storage.id_}",'
-            f' "optimization_parameters": {{"stat_subset_size": 300}},'
             f' "project_id": "{project.id_}",'
             f' "type": "optimize_pot",'
             f' "workspace_id": "{project.workspace_id}"}}',
@@ -76,8 +97,10 @@ def test_start_pot_job(self, fxt_db_project_service, fxt_pot_hyperparameters, fx
                 "model_storage_id": model.model_storage.id_,
                 "project_id": project.id_,
                 "enable_optimize_from_dataset_shard": True,
-                "max_number_of_annotations": None,
-                "min_annotation_size": None,
+                "max_number_of_annotations": filtering_params.max_annotation_objects.max_annotation_objects,
+                "min_number_of_annotations": filtering_params.min_annotation_objects.min_annotation_objects,
+                "min_annotation_size": filtering_params.min_annotation_pixels.min_annotation_pixels,
+                "max_annotation_size": filtering_params.max_annotation_pixels.max_annotation_pixels,
             },
             metadata={
                 "base_model_id": model.id_,