Training job: prepare OTX training config (#5028)

itallix · leoll2 · web-flow · commit fadcddeeb1b5 · 2025-11-20T15:52:54.000Z
Co-authored-by: Leonardo Lai &lt;leonardo.lai@intel.com&gt;
diff --git a/application/backend/app/lifecycle.py b/application/backend/app/lifecycle.py
@@ -19,11 +19,12 @@
 from app.db import MigrationManager, get_db_session
 from app.scheduler import Scheduler
 from app.schemas.job import JobType
-from app.services import DatasetService, LabelService, ModelService
+from app.services import DatasetService, LabelService, ModelService, TrainingConfigurationService
 from app.services.base_weights_service import BaseWeightsService
 from app.services.data_collect import DataCollector
 from app.services.event.event_bus import EventBus
 from app.services.training import OTXTrainer
+from app.services.training.otx_trainer import TrainingDependencies
 from app.services.training.subset_assignment import SubsetAssigner, SubsetService
 from app.settings import get_settings
 from app.webrtc.manager import WebRTCManager
@@ -46,23 +47,20 @@ def setup_job_controller(data_dir: Path, max_parallel_jobs: int) -> tuple[JobQue
     """
     q = JobQueue()
     job_runnable_factory = RunnableFactory[JobType, Runnable]()
-    base_weights_service = BaseWeightsService(data_dir=data_dir)
-    subset_service = SubsetService()
-    subset_assigner = SubsetAssigner()
-    label_service = LabelService()
-    model_service = ModelService()
-    dataset_service = DatasetService(data_dir=data_dir, label_service=label_service)
     job_runnable_factory.register(
         JobType.TRAIN,
         partial(
             OTXTrainer,
-            base_weights_service=base_weights_service,
-            subset_service=subset_service,
-            subset_assigner=subset_assigner,
-            dataset_service=dataset_service,
-            model_service=model_service,
-            data_dir=data_dir,
-            db_session_factory=get_db_session,
+            training_deps=TrainingDependencies(
+                base_weights_service=BaseWeightsService(data_dir=data_dir),
+                subset_service=SubsetService(),
+                subset_assigner=SubsetAssigner(),
+                dataset_service=DatasetService(data_dir=data_dir, label_service=LabelService()),
+                model_service=ModelService(),
+                training_configuration_service=TrainingConfigurationService(),
+                data_dir=data_dir,
+                db_session_factory=get_db_session,
+            ),
         ),
     )
     process_runner_factory = ProcessRunnerFactory(job_runnable_factory)
diff --git a/application/backend/app/services/__init__.py b/application/backend/app/services/__init__.py
@@ -22,6 +22,7 @@
 from .sink_service import SinkService
 from .source_service import SourceService, SourceUpdateService
 from .system_service import SystemService
+from .training_configuration_service import TrainingConfigurationService
 from .video_stream_service import VideoStreamService
 
 __all__ = [
@@ -46,5 +47,6 @@
     "SourceService",
     "SourceUpdateService",
     "SystemService",
+    "TrainingConfigurationService",
     "VideoStreamService",
 ]
diff --git a/application/backend/app/services/training/models.py b/application/backend/app/services/training/models.py
@@ -6,6 +6,7 @@
 from uuid import UUID, uuid4
 
 from loguru import logger
+from pydantic import Field
 
 from app.core.jobs import Job, JobParams, JobType
 from app.schemas.project import TaskBase
@@ -17,7 +18,7 @@ class TrainingParams(JobParams):
     model_architecture_id: str
     parent_model_revision_id: UUID | None = None
     task: TaskBase
-    model_id: UUID = uuid4()  # Reserve the ID for the model to be created for this training job
+    model_id: UUID = Field(default_factory=uuid4)
 
 
 class ProjectJob(Job):
diff --git a/application/backend/app/services/training/otx_trainer.py b/application/backend/app/services/training/otx_trainer.py
@@ -8,16 +8,26 @@
 from pathlib import Path
 from uuid import UUID
 
+import yaml
 from datumaro.experimental import Dataset
 from datumaro.experimental.fields import Subset
 from loguru import logger
+from otx.types.task import OTXTaskType
 from sqlalchemy.orm import Session
 
+from app.core.jobs import JobType
 from app.core.run import ExecutionContext
-from app.models import DatasetItemAnnotationStatus
+from app.models import DatasetItemAnnotationStatus, TaskType
+from app.models.training_configuration.configuration import TrainingConfiguration
 from app.schemas.model import TrainingStatus
 from app.schemas.project import TaskBase
-from app.services import BaseWeightsService, DatasetService, ModelRevisionMetadata, ModelService
+from app.services import (
+    BaseWeightsService,
+    DatasetService,
+    ModelRevisionMetadata,
+    ModelService,
+    TrainingConfigurationService,
+)
 
 from .base import Trainer, step
 from .models import TrainingParams
@@ -26,6 +36,23 @@
 MODEL_WEIGHTS_PATH = "model_weights_path"
 
 
+# TODO: Consider adopting some lightweight DI framework
+# As the number of constructor dependencies grows and start violating ruff rules, we should evaluate DI frameworks like:
+# - dependency-injector (https://python-dependency-injector.ets-labs.org/)
+# - injector (https://github.com/python-injector/injector)
+# - python-inject (https://github.com/ivankorobkov/python-inject)
+@dataclass(frozen=True)
+class TrainingDependencies:
+    data_dir: Path
+    base_weights_service: BaseWeightsService
+    subset_service: SubsetService
+    dataset_service: DatasetService
+    model_service: ModelService
+    training_configuration_service: TrainingConfigurationService
+    subset_assigner: SubsetAssigner
+    db_session_factory: Callable[[], AbstractContextManager[Session]]
+
+
 @dataclass(frozen=True)
 class DatasetInfo:
     training: Dataset
@@ -39,22 +66,17 @@ class OTXTrainer(Trainer):
 
     def __init__(
         self,
-        data_dir: Path,
-        base_weights_service: BaseWeightsService,
-        subset_service: SubsetService,
-        dataset_service: DatasetService,
-        model_service: ModelService,
-        subset_assigner: SubsetAssigner,
-        db_session_factory: Callable[[], AbstractContextManager[Session]],
+        training_deps: TrainingDependencies,
     ):
         super().__init__()
-        self._data_dir = data_dir
-        self._base_weights_service = base_weights_service
-        self._subset_service = subset_service
-        self._dataset_service = dataset_service
-        self._model_service = model_service
-        self._subset_assigner = subset_assigner
-        self._db_session_factory = db_session_factory
+        self._data_dir = training_deps.data_dir
+        self._base_weights_service = training_deps.base_weights_service
+        self._subset_service = training_deps.subset_service
+        self._dataset_service = training_deps.dataset_service
+        self._model_service = training_deps.model_service
+        self._training_configuration_service = training_deps.training_configuration_service
+        self._subset_assigner = training_deps.subset_assigner
+        self._db_session_factory = training_deps.db_session_factory
 
     @step("Prepare Model Weights")
     def prepare_weights(self, training_params: TrainingParams) -> Path:
@@ -127,19 +149,28 @@ def create_training_dataset(self, project_id: UUID, task: TaskBase) -> DatasetIn
                 revision_id=self._dataset_service.save_revision(project_id, dm_dataset),
             )
 
-    @step("Prepare Model Metadata")
+    @step("Prepare Model and Training Configuration")
     def prepare_model(self, training_params: TrainingParams, dataset_revision_id: UUID) -> None:
         if training_params.project_id is None:
             raise ValueError("Project ID must be provided for model preparation")
         with self._db_session_factory() as db:
+            self._training_configuration_service.set_db_session(db)
             self._model_service.set_db_session(db)
+            configuration = self._training_configuration_service.get_training_configuration(
+                project_id=training_params.project_id,
+                model_architecture_id=training_params.model_architecture_id,
+            )
+            config_path = self.__build_model_config_path(
+                self._data_dir, training_params.project_id, training_params.model_id
+            )
+            self.__persist_configuration(configuration, config_path, training_params.task)
             self._model_service.create_revision(
                 ModelRevisionMetadata(
                     model_id=training_params.model_id,
                     project_id=training_params.project_id,
                     architecture_id=training_params.model_architecture_id,
                     parent_revision_id=training_params.parent_model_revision_id,
-                    training_configuration=None,  # TODO: to be set when config is added
+                    training_configuration=configuration,
                     dataset_revision_id=dataset_revision_id,
                     training_status=TrainingStatus.NOT_STARTED,
                 )
@@ -172,5 +203,31 @@ def run(self, ctx: ExecutionContext) -> None:
         self.train_model(training_params)
 
     @staticmethod
-    def __build_model_weights_path(data_dir: Path, project_id: UUID, model_id: UUID) -> Path:
-        return data_dir / "projects" / str(project_id) / "models" / str(model_id) / "model.pth"
+    def __base_model_path(data_dir: Path, project_id: UUID, model_id: UUID) -> Path:
+        return data_dir / "projects" / str(project_id) / "models" / str(model_id)
+
+    @classmethod
+    def __build_model_weights_path(cls, data_dir: Path, project_id: UUID, model_id: UUID) -> Path:
+        return cls.__base_model_path(data_dir, project_id, model_id) / "model.pth"
+
+    @classmethod
+    def __build_model_config_path(cls, data_dir: Path, project_id: UUID, model_id: UUID) -> Path:
+        return cls.__base_model_path(data_dir, project_id, model_id) / "config.yaml"
+
+    @staticmethod
+    def __persist_configuration(configuration: TrainingConfiguration, config_path: Path, task: TaskBase) -> None:
+        extended_config = configuration.model_dump(exclude_none=True)
+        extended_config["job_type"] = JobType.TRAIN.value
+        match task.task_type:
+            case TaskType.CLASSIFICATION:
+                if task.exclusive_labels:
+                    extended_config["sub_task_type"] = OTXTaskType.MULTI_CLASS_CLS.value
+                else:
+                    extended_config["sub_task_type"] = OTXTaskType.MULTI_LABEL_CLS.value
+            case TaskType.DETECTION:
+                extended_config["sub_task_type"] = OTXTaskType.DETECTION.value
+            case TaskType.INSTANCE_SEGMENTATION:
+                extended_config["sub_task_type"] = OTXTaskType.INSTANCE_SEGMENTATION.value
+        config_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(config_path, "w") as f:
+            yaml.dump(extended_config, f, default_flow_style=False)
diff --git a/application/backend/app/services/training_configuration_service.py b/application/backend/app/services/training_configuration_service.py
@@ -7,16 +7,15 @@
 from app.models.training_configuration.configuration import PartialTrainingConfiguration, TrainingConfiguration
 from app.repositories import ModelRevisionRepository, ProjectRepository
 from app.repositories.training_configuration_repo import TrainingConfigurationRepository
-from app.services import ResourceNotFoundError, ResourceType
+from app.services import BaseSessionManagedService, ResourceNotFoundError, ResourceType
 from app.services.tools import ConfigurationOverlayTools
 from app.supported_models import SupportedModels
 from app.supported_models.default_models import DefaultModels
 
 
-class TrainingConfigurationService:
-    def __init__(self, db_session: Session) -> None:
-        self._db_session = db_session
-        self._training_config_repo = TrainingConfigurationRepository(db_session)
+class TrainingConfigurationService(BaseSessionManagedService):
+    def __init__(self, db_session: Session | None = None) -> None:
+        super().__init__(db_session)
 
     def get_training_configuration(
         self,
@@ -63,7 +62,7 @@ def _get_by_model_revision_id(self, project_id: UUID, model_revision_id: UUID) -
         Returns:
             TrainingConfiguration: The training configuration object.
         """
-        model = ModelRevisionRepository(str(project_id), self._db_session).get_by_id(str(model_revision_id))
+        model = ModelRevisionRepository(str(project_id), self.db_session).get_by_id(str(model_revision_id))
         if not model:
             raise ResourceNotFoundError(ResourceType.MODEL, str(model_revision_id))
         return TrainingConfiguration.model_validate(model.training_configuration)
@@ -79,7 +78,7 @@ def _get_by_model_architecture_id(self, project_id: UUID, model_architecture_id:
         Returns:
             TrainingConfiguration: The training configuration object.
         """
-        stored_config = TrainingConfigurationRepository(self._db_session).get_by_project_and_model_architecture(
+        stored_config = TrainingConfigurationRepository(self.db_session).get_by_project_and_model_architecture(
             project_id=str(project_id),
             model_architecture_id=model_architecture_id,
         )
@@ -102,7 +101,7 @@ def _get_default_configuration(self, project_id: UUID) -> TrainingConfiguration:
         Returns:
             TrainingConfiguration: The default training configuration object.
         """
-        project = ProjectRepository(self._db_session).get_by_id(str(project_id))
+        project = ProjectRepository(self.db_session).get_by_id(str(project_id))
         if not project:
             raise ResourceNotFoundError(ResourceType.PROJECT, str(project_id))
 
@@ -133,7 +132,7 @@ def update_training_configuration(
         Returns:
             TrainingConfiguration: The updated training configuration object.
         """
-        project = ProjectRepository(self._db_session).get_by_id(str(project_id))
+        project = ProjectRepository(self.db_session).get_by_id(str(project_id))
         if not project:
             raise ResourceNotFoundError(ResourceType.PROJECT, str(project_id))
 
@@ -150,7 +149,8 @@ def update_training_configuration(
 
         validated_updated_config = PartialTrainingConfiguration(**updated_config)  # type: ignore[arg-type]
 
-        self._training_config_repo.create_or_update(
+        training_config_repo = TrainingConfigurationRepository(self.db_session)
+        training_config_repo.create_or_update(
             project_id=str(project_id),
             model_architecture_id=model_architecture_id,
             configuration_data=validated_updated_config.model_dump(),
diff --git a/application/backend/tests/unit/services/training/test_otx_trainer.py b/application/backend/tests/unit/services/training/test_otx_trainer.py
