Save initial model revision in a training job (#5024)

itallix · web-flow · commit fcbaf56493ef · 2025-11-18T08:17:15.000Z
diff --git a/application/backend/app/lifecycle.py b/application/backend/app/lifecycle.py
@@ -19,7 +19,7 @@
 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
+from app.services import DatasetService, LabelService, ModelService
 from app.services.base_weights_service import BaseWeightsService
 from app.services.data_collect import DataCollector
 from app.services.event.event_bus import EventBus
@@ -50,6 +50,7 @@ def setup_job_controller(data_dir: Path, max_parallel_jobs: int) -> tuple[JobQue
     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,
@@ -59,6 +60,7 @@ def setup_job_controller(data_dir: Path, max_parallel_jobs: int) -> tuple[JobQue
             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,
         ),
diff --git a/application/backend/app/services/__init__.py b/application/backend/app/services/__init__.py
@@ -15,7 +15,7 @@
 from .dispatch_service import DispatchService
 from .label_service import LabelService
 from .metrics_service import MetricsService
-from .model_service import ModelService
+from .model_service import ModelRevisionMetadata, ModelService
 from .pipeline_metrics_service import PipelineMetricsService
 from .pipeline_service import PipelineService
 from .project_service import ProjectService
@@ -32,6 +32,7 @@
     "DispatchService",
     "LabelService",
     "MetricsService",
+    "ModelRevisionMetadata",
     "ModelService",
     "PipelineMetricsService",
     "PipelineService",
diff --git a/application/backend/app/services/dataset_service.py b/application/backend/app/services/dataset_service.py
@@ -375,7 +375,7 @@ def _get_image_path(item: DatasetItem) -> str:
             get_image_path=_get_image_path,
         )
 
-    def save_revision(self, project_id: UUID, dataset: dm.Dataset) -> None:
+    def save_revision(self, project_id: UUID, dataset: dm.Dataset) -> UUID:
         """
         Saves the dataset as a new revision.
 
@@ -387,7 +387,7 @@ def save_revision(self, project_id: UUID, dataset: dm.Dataset) -> None:
             dataset: The Datumaro dataset to export.
 
         Returns:
-            None
+            UUID: The UUID of the newly created dataset revision.
         """
         revision_repo = DatasetRevisionRepository(db=self.db_session)
         revision_db = revision_repo.save(
@@ -403,3 +403,4 @@ def save_revision(self, project_id: UUID, dataset: dm.Dataset) -> None:
             export_images=True,
             as_zip=True,
         )
+        return UUID(revision_db.id)
diff --git a/application/backend/app/services/model_service.py b/application/backend/app/services/model_service.py
@@ -1,24 +1,35 @@
 # Copyright (C) 2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
+from dataclasses import dataclass
 from uuid import UUID
 
 from sqlalchemy.exc import IntegrityError
-from sqlalchemy.orm import Session
 
-from app.repositories import ModelRevisionRepository, ProjectRepository
+from app.db.schema import ModelRevisionDB
+from app.models.training_configuration.configuration import TrainingConfiguration
+from app.repositories import LabelRepository, ModelRevisionRepository, ProjectRepository
 from app.schemas.model import Model as ModelSchema
+from app.schemas.model import TrainingStatus
 
-from .base import ResourceInUseError, ResourceNotFoundError, ResourceType
+from .base import BaseSessionManagedService, ResourceInUseError, ResourceNotFoundError, ResourceType
 from .mappers.model_revision_mapper import ModelRevisionMapper
 from .parent_process_guard import parent_process_only
 
 
-class ModelService:
-    """Service to register and activate models"""
+@dataclass(frozen=True)
+class ModelRevisionMetadata:
+    model_id: UUID
+    project_id: UUID
+    architecture_id: str
+    parent_revision_id: UUID | None
+    dataset_revision_id: UUID | None
+    training_status: TrainingStatus
+    training_configuration: TrainingConfiguration | None = None
+
 
-    def __init__(self, db_session: Session) -> None:
-        self._db_session = db_session
+class ModelService(BaseSessionManagedService):
+    """Service to register and activate models"""
 
     def get_model_by_id(self, project_id: UUID, model_id: UUID) -> ModelSchema:
         """
@@ -39,7 +50,7 @@ def get_model_by_id(self, project_id: UUID, model_id: UUID) -> ModelSchema:
             ResourceNotFoundError: If the project with the given project_id does not exist,
                 or if no model with the given model_id is found within the project.
         """
-        project_repo = ProjectRepository(self._db_session)
+        project_repo = ProjectRepository(self.db_session)
         # Prefer using a JOIN here since the list of model revisions per project is not large,
         # and it allows us to check for project existence and fetch the model in a single query.
         project = project_repo.get_by_id(str(project_id))
@@ -72,10 +83,10 @@ def delete_model_by_id(self, project_id: UUID, model_id: UUID) -> None:
             ResourceInUseError: If the model cannot be deleted due to integrity constraints
                 (e.g., the model is referenced by other entities).
         """
-        project_repo = ProjectRepository(self._db_session)
+        project_repo = ProjectRepository(self.db_session)
         if not project_repo.exists(str(project_id)):
             raise ResourceNotFoundError(ResourceType.PROJECT, str(project_id))
-        model_rev_repo = ModelRevisionRepository(self._db_session)
+        model_rev_repo = ModelRevisionRepository(self.db_session)
         try:
             # TODO: delete model artifacts from filesystem when implemented
             deleted = model_rev_repo.delete(str(model_id))
@@ -102,8 +113,40 @@ def list_models(self, project_id: UUID) -> list[ModelSchema]:
         Raises:
             ResourceNotFoundError: If the project with the given project_id does not exist.
         """
-        project_repo = ProjectRepository(self._db_session)
+        project_repo = ProjectRepository(self.db_session)
         project = project_repo.get_by_id(str(project_id))
         if not project:
             raise ResourceNotFoundError(ResourceType.PROJECT, str(project_id))
         return [ModelRevisionMapper.to_schema(model_rev_db) for model_rev_db in project.model_revisions]
+
+    def create_revision(self, metadata: ModelRevisionMetadata) -> None:
+        """
+        Create and persist a new model revision for the given project metadata.
+
+        Reads the project's label definitions, serializes them into a dict format,
+        combines them with the provided metadata into a new model revision record,
+        and saves it to the database.
+
+        Args:
+            metadata (ModelRevisionMetadata): Metadata used to create the new model revision
+                including project id, architecture, optional parent revision id,
+                dataset revision id, training status and optional training
+                configuration.
+        """
+        label_repo = LabelRepository(project_id=str(metadata.project_id), db=self.db_session)
+        labels_schema_rev = {"labels": [{"name": label.name, "id": label.id} for label in label_repo.list_all()]}
+        model_revision_repo = ModelRevisionRepository(self.db_session)
+        model_revision_repo.save(
+            ModelRevisionDB(
+                id=str(metadata.model_id),
+                project_id=str(metadata.project_id),
+                architecture=metadata.architecture_id,
+                parent_revision=str(metadata.parent_revision_id) if metadata.parent_revision_id else None,
+                training_status=metadata.training_status,
+                training_configuration=metadata.training_configuration.model_dump()
+                if metadata.training_configuration
+                else {},
+                training_dataset_id=str(metadata.dataset_revision_id),
+                label_schema_revision=labels_schema_rev,
+            )
+        )
diff --git a/application/backend/app/services/training/otx_trainer.py b/application/backend/app/services/training/otx_trainer.py
@@ -4,6 +4,7 @@
 import time
 from collections.abc import Callable
 from contextlib import AbstractContextManager
+from dataclasses import dataclass
 from pathlib import Path
 from uuid import UUID
 
@@ -14,14 +15,25 @@
 
 from app.core.run import ExecutionContext
 from app.models import DatasetItemAnnotationStatus
-from app.services import BaseWeightsService, DatasetService
+from app.schemas.model import TrainingStatus
+from app.schemas.project import TaskBase
+from app.services import BaseWeightsService, DatasetService, ModelRevisionMetadata, ModelService
 
 from .base import Trainer, step
+from .models import TrainingParams
 from .subset_assignment import SplitRatios, SubsetAssigner, SubsetService
 
 MODEL_WEIGHTS_PATH = "model_weights_path"
 
 
+@dataclass(frozen=True)
+class DatasetInfo:
+    training: Dataset
+    validation: Dataset
+    testing: Dataset
+    revision_id: UUID
+
+
 class OTXTrainer(Trainer):
     """OTX-specific trainer implementation."""
 
@@ -31,6 +43,7 @@ def __init__(
         base_weights_service: BaseWeightsService,
         subset_service: SubsetService,
         dataset_service: DatasetService,
+        model_service: ModelService,
         subset_assigner: SubsetAssigner,
         db_session_factory: Callable[[], AbstractContextManager[Session]],
     ):
@@ -39,26 +52,22 @@ def __init__(
         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._training_dataset: Dataset | None = None
-        self._validation_dataset: Dataset | None = None
-        self._testing_dataset: Dataset | None = None
 
     @step("Prepare Model Weights")
-    def prepare_weights(self) -> Path:
+    def prepare_weights(self, training_params: TrainingParams) -> Path:
         """
         Prepare weights for training based on training parameters.
 
         If a parent model revision ID is provided, it fetches the weights from the parent model.
         Otherwise, it retrieves the base weights for the specified model architecture.
         """
-        if self._training_params is None:
-            raise ValueError("Training parameters not set")
-        parent_model_revision_id = self._training_params.parent_model_revision_id
-        task = self._training_params.task
-        model_architecture_id = self._training_params.model_architecture_id
-        project_id = self._training_params.project_id
+        parent_model_revision_id = training_params.parent_model_revision_id
+        task = training_params.task
+        model_architecture_id = training_params.model_architecture_id
+        project_id = training_params.project_id
         if parent_model_revision_id is None:
             return self._base_weights_service.get_local_weights_path(
                 task=task.task_type, model_manifest_id=model_architecture_id
@@ -74,17 +83,11 @@ def prepare_weights(self) -> Path:
         return weights_path
 
     @step("Assign Dataset Subsets")
-    def assign_subsets(self) -> None:
+    def assign_subsets(self, project_id: UUID) -> None:
         """Assigning subsets to all unassigned dataset items in the project dataset."""
-        if self._training_params is None:
-            raise ValueError("Training parameters not set")
-        project_id = self._training_params.project_id
-        self.report_progress("Retrieving unassigned items")
-        if project_id is None:
-            raise ValueError("Project ID must be provided for subset assignment")
-
         with self._db_session_factory() as db:
             self._subset_service.set_db_session(db)
+            self.report_progress("Retrieving unassigned items")
             unassigned_items = self._subset_service.get_unassigned_items_with_labels(project_id)
 
             if not unassigned_items:
@@ -112,30 +115,41 @@ def assign_subsets(self) -> None:
         self.report_progress(f"Successfully assigned {len(assignments)} items to subsets")
 
     @step("Create Training Dataset")
-    def create_training_dataset(self) -> None:
+    def create_training_dataset(self, project_id: UUID, task: TaskBase) -> DatasetInfo:
         """Create datasets for training, validation, and testing."""
-        if self._training_params is None:
-            raise ValueError("Training parameters not set")
-        project_id = self._training_params.project_id
-        if project_id is None:
-            raise ValueError("Project ID must be provided")
-        task = self._training_params.task
-
         with self._db_session_factory() as db:
             self._dataset_service.set_db_session(db)
             dm_dataset = self._dataset_service.get_dm_dataset(project_id, task, DatasetItemAnnotationStatus.REVIEWED)
-            self._training_dataset = dm_dataset.filter_by_subset(Subset.TRAINING)
-            self._validation_dataset = dm_dataset.filter_by_subset(Subset.VALIDATION)
-            self._testing_dataset = dm_dataset.filter_by_subset(Subset.TESTING)
-            self._dataset_service.save_revision(project_id, dm_dataset)
+            return DatasetInfo(
+                training=dm_dataset.filter_by_subset(Subset.TRAINING),
+                validation=dm_dataset.filter_by_subset(Subset.VALIDATION),
+                testing=dm_dataset.filter_by_subset(Subset.TESTING),
+                revision_id=self._dataset_service.save_revision(project_id, dm_dataset),
+            )
+
+    @step("Prepare Model Metadata")
+    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._model_service.set_db_session(db)
+            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
+                    dataset_revision_id=dataset_revision_id,
+                    training_status=TrainingStatus.NOT_STARTED,
+                )
+            )
 
     @step("Train Model with OTX")
-    def train_model(self) -> None:
+    def train_model(self, training_params: TrainingParams) -> None:
         """Execute OTX model training."""
-        if self._training_params is None:
-            raise ValueError("Training parameters not set")
         # Simulate training with progress reporting
-        job_id = self._training_params.job_id
+        job_id = training_params.job_id
         step_count = 20
         for i in range(step_count):
             time.sleep(1)
@@ -145,12 +159,17 @@ def train_model(self) -> None:
 
     def run(self, ctx: ExecutionContext) -> None:
         self._ctx = ctx
-        self._training_params = self._get_training_params(ctx)
-
-        self.prepare_weights()
-        self.assign_subsets()
-        self.create_training_dataset()
-        self.train_model()
+        training_params = self._get_training_params(ctx)
+        project_id = training_params.project_id
+        if project_id is None:
+            raise ValueError("Project ID must be provided in training parameters")
+        task = training_params.task
+
+        self.prepare_weights(training_params)
+        self.assign_subsets(project_id)
+        dataset_info = self.create_training_dataset(project_id, task)
+        self.prepare_model(training_params, dataset_info.revision_id)
+        self.train_model(training_params)
 
     @staticmethod
     def __build_model_weights_path(data_dir: Path, project_id: UUID, model_id: UUID) -> Path:
diff --git a/application/backend/tests/integration/services/test_dataset_service.py b/application/backend/tests/integration/services/test_dataset_service.py
@@ -1385,13 +1385,11 @@ def test_save_revision(
         project, db_dataset_items = fxt_project_with_subset_items
         dataset = fxt_dataset_service.get_dm_dataset(project.id, project.task, DatasetItemAnnotationStatus.REVIEWED)
 
-        fxt_dataset_service.save_revision(
+        revision_id = fxt_dataset_service.save_revision(
             project_id=project.id,
             dataset=dataset,
         )
 
         # Verify that a revision entry was created
-        db_revisions = db_session.query(DatasetRevisionDB).all()
-        assert len(db_revisions) == 1
-        revision_id = db_revisions[0].id
-        assert (fxt_projects_dir / str(project.id) / "dataset_revisions" / revision_id / "dataset.zip").exists()
+        assert db_session.get(DatasetRevisionDB, str(revision_id)) is not None
+        assert (fxt_projects_dir / str(project.id) / "dataset_revisions" / str(revision_id) / "dataset.zip").exists()
diff --git a/application/backend/tests/unit/services/training/test_otx_trainer.py b/application/backend/tests/unit/services/training/test_otx_trainer.py
