Implement subset splitting (#4965)

Author: itallix

application/backend/app/cli.py

@@ -9,9 +9,9 @@
 
 import click
 
-from app.core.models.task_type import TaskType
 from app.db import MigrationManager, get_db_session
 from app.db.schema import DatasetItemDB, LabelDB, ModelRevisionDB, PipelineDB, ProjectDB, SinkDB, SourceDB
+from app.models import TaskType
 from app.schemas import DisconnectedSinkConfig, DisconnectedSourceConfig, OutputFormat, SinkType, SourceType
 from app.schemas.model import TrainingStatus
 from app.schemas.pipeline import FixedRateDataCollectionPolicy

application/backend/app/core/jobs/exec/process_run.py

@@ -23,7 +23,6 @@
 from multiprocessing.connection import Connection
 from multiprocessing.context import SpawnProcess
 from multiprocessing.synchronize import Event
-from pathlib import Path
 
 from app.core.jobs.models import Done, ExecutionEvent, Failed, Job, JobType, Started
 from app.core.run import ExecutionContext, RunnableFactory, Runner
@@ -39,14 +38,12 @@ class ProcessRun:
 
     Args:
         ctx (mp.context.SpawnContext): Multiprocessing context for process & IPC.
-        data_dir (Path): Directory for job data.
         runnable_factory (RunnableFactory): Factory to create runnable job instances.
         job (Job): Job specification.
     """
 
-    def __init__(self, ctx: mp.context.SpawnContext, data_dir: Path, runnable_factory: RunnableFactory, job: Job):
+    def __init__(self, ctx: mp.context.SpawnContext, runnable_factory: RunnableFactory, job: Job):
         self._ctx = ctx
-        self._data_dir = data_dir
         self._runnable_factory = runnable_factory
         self._job = job
         self._parent, self._child = ctx.Pipe(duplex=False)
@@ -58,7 +55,6 @@ def start(self) -> "ProcessRun":
             target=_entrypoint,
             args=(
                 self._runnable_factory,
-                self._data_dir,
                 self._job.job_type,
                 self._job.params.model_dump_json(),
                 self._child,
@@ -118,7 +114,7 @@ async def stop(self, graceful_timeout: float = 6.0, term_timeout: float = 3.0, k
 
 
 def _entrypoint(
-    get_runnable: RunnableFactory, data_dir: Path, job_type: str, payload: str, conn: Connection, cancel_event: Event
+    get_runnable: RunnableFactory, job_type: str, payload: str, conn: Connection, cancel_event: Event
 ) -> None:
     """
     Entrypoint for the child process.
@@ -127,7 +123,6 @@ def _entrypoint(
 
     Args:
         get_runnable (RunnableFactory): Factory to create runnable job instance.
-        data_dir (Path): Directory for job data.
         job_type (str): Type of job to execute.
         payload (str): Serialized job parameters.
         conn (Connection): IPC connection to parent process.
@@ -150,7 +145,7 @@ def heartbeat():
 
     try:
         conn.send(Started())
-        runnable.run(ExecutionContext(payload=payload, data_dir=data_dir, report=report, heartbeat=heartbeat))
+        runnable.run(ExecutionContext(payload=payload, report=report, heartbeat=heartbeat))
         conn.send(Done())
     except CancelledExc:
         conn.send(Cancelled())
@@ -166,17 +161,15 @@ class ProcessRunnerFactory:
     Factory for creating process-based job runners.
 
     Args:
-        data_dir (Path): Directory for job data.
         runnable_factory (RunnableFactory): Factory to create runnable job instances.
 
     Methods:
         for_job(job: Job) -> Runner[Job, ExecutionEvent]: Create a ProcessRun instance for the given job.
     """
 
-    def __init__(self, data_dir: Path, runnable_factory: RunnableFactory) -> None:
+    def __init__(self, runnable_factory: RunnableFactory) -> None:
         # consider using native context for python 3.14 due to upgrade to 'fork_server' model
         self._ctx = mp.get_context("spawn")
-        self._data_dir = data_dir
         self._runnable_factory = runnable_factory
 
     def for_job(self, job: Job) -> Runner[Job, ExecutionEvent]:
@@ -189,4 +182,4 @@ def for_job(self, job: Job) -> Runner[Job, ExecutionEvent]:
         Returns:
             Runner[Job, ExecutionEvent]: Process-based job runner.
         """
-        return ProcessRun(self._ctx, self._data_dir, self._runnable_factory, job)
+        return ProcessRun(self._ctx, self._runnable_factory, job)

application/backend/app/core/jobs/exec/thread_run.py

@@ -101,7 +101,7 @@ class ThreadAwareExecutionContext(ExecutionContext):
             def __init__(self, runner: "ThreadRun"):
                 self.runner = runner
 
-            def report_progress(self, message: str = "training", progress: float = 0.0):
+            def report(self, message: str = "training", progress: float = 0.0):
                 if not self.runner._cancel_event.is_set():
                     self.runner._event_queue.put(Progress(message, progress))
 

application/backend/app/core/models/base.py renamed to application/backend/app/core/models.py

@@ -16,7 +16,6 @@
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from pathlib import Path
 from typing import Generic, Protocol, TypeVar
 
 ReportFn = Callable[[str, float], None]
@@ -26,14 +25,9 @@
 @dataclass(kw_only=True)
 class ExecutionContext:
     payload: str
-    data_dir: Path
     report: ReportFn
     heartbeat: HeartbeatFn
 
-    def report_progress(self, msg: str = "", progress: float = 0.0) -> None:
-        """Report progress of the execution."""
-        self.report(msg, progress)
-
 
 class Runnable(Protocol):  # ignore
     """Generic interface for activities executed by runners."""

application/backend/app/core/models/__init__.py

@@ -15,13 +15,14 @@
 
 from app.core.jobs import JobController, JobQueue, ProcessRunnerFactory
 from app.core.run import Runnable, RunnableFactory
-from app.db import MigrationManager
+from app.db import MigrationManager, get_db_session
 from app.scheduler import Scheduler
 from app.schemas.job import JobType
 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.subset_assignment import SubsetAssigner, SubsetService
 from app.settings import get_settings
 from app.webrtc.manager import WebRTCManager
 
@@ -30,23 +31,36 @@
 
 def setup_job_controller(data_dir: Path, max_parallel_jobs: int) -> tuple[JobQueue, JobController]:
     """
-    Set up job controller with queue and processing infrastructure.
+    Initializes and configures the job queue and job controller for managing parallel job execution.
 
-    Creates a job queue and controller with configured parallel job limits and training infrastructure
-    for job execution.
+    Sets up the infrastructure to run jobs concurrently and registers classes that comply with the Runnable protocol,
+    each associated with a job type and its required dependencies. These classes are executed in a context defined
+    by the runner factory.
 
     Args:
-        data_dir: Path to the data directory.
+        data_dir: Path to the directory containing data required for job execution.
         max_parallel_jobs (int): Maximum number of jobs that can run concurrently.
 
     Returns:
-        tuple[JobQueue, JobController]: A tuple containing the job queue instance and the configured job controller.
+        tuple[JobQueue, JobController]: The job queue and the configured job controller.
     """
     q = JobQueue()
     job_runnable_factory = RunnableFactory[JobType, Runnable]()
     base_weights_service = BaseWeightsService(data_dir=data_dir)
-    job_runnable_factory.register(JobType.TRAIN, partial(OTXTrainer, base_weights_service=base_weights_service))
-    process_runner_factory = ProcessRunnerFactory(data_dir, job_runnable_factory)
+    subset_service = SubsetService()
+    subset_assigner = SubsetAssigner()
+    job_runnable_factory.register(
+        JobType.TRAIN,
+        partial(
+            OTXTrainer,
+            base_weights_service=base_weights_service,
+            subset_service=subset_service,
+            subset_assigner=subset_assigner,
+            data_dir=data_dir,
+            db_session_factory=get_db_session,
+        ),
+    )
+    process_runner_factory = ProcessRunnerFactory(job_runnable_factory)
     job_controller = JobController(
         jobs_queue=q, runner_factory=process_runner_factory, max_parallel_jobs=max_parallel_jobs
     )

application/backend/app/core/run/runnable.py

@@ -4,6 +4,7 @@
 from .dataset_item import DatasetItem, DatasetItemAnnotation, DatasetItemFormat, DatasetItemSubset
 from .label import Label, LabelReference
 from .shape import FullImage, Point, Polygon, Rectangle, Shape
+from .task_type import TaskType
 
 __all__ = [
     "DatasetItem",
@@ -17,4 +18,5 @@
     "Polygon",
     "Rectangle",
     "Shape",
+    "TaskType",
 ]

application/backend/app/lifecycle.py

@@ -8,6 +8,7 @@
 from sqlalchemy.orm import Session
 
 from app.db.schema import DatasetItemDB, DatasetItemLabelDB
+from app.models import DatasetItemSubset
 
 
 class UpdateDatasetItemAnnotation(NamedTuple):
@@ -128,19 +129,20 @@ def get_subset(self, obj_id: str) -> str | None:
         )
         return self.db.scalar(stmt)
 
-    def set_subset(self, obj_id: str, subset: str) -> None:
+    def set_subset(self, obj_ids: set[str], subset: str) -> int:
         stmt = (
             update(DatasetItemDB)
             .where(
                 DatasetItemDB.project_id == self.project_id,
-                DatasetItemDB.id == obj_id,
+                DatasetItemDB.id.in_(obj_ids),
             )
             .values(
                 subset=subset,
                 updated_at=datetime.now(UTC),
             )
         )
-        self.db.execute(stmt)
+        result = self.db.execute(stmt)
+        return result.rowcount or 0
 
     def set_labels(self, dataset_item_id: str, label_ids: set[str]) -> None:
         self.delete_labels(dataset_item_id)
@@ -153,3 +155,23 @@ def set_labels(self, dataset_item_id: str, label_ids: set[str]) -> None:
     def delete_labels(self, dataset_item_id: str) -> None:
         stmt = delete(DatasetItemLabelDB).where(DatasetItemLabelDB.dataset_item_id == dataset_item_id)
         self.db.execute(stmt)
+
+    def list_unassigned_items(self) -> list[DatasetItemLabelDB]:
+        stmt = (
+            select(DatasetItemLabelDB)
+            .join(DatasetItemDB)
+            .where(
+                DatasetItemDB.project_id == self.project_id,
+                DatasetItemDB.subset == DatasetItemSubset.UNASSIGNED,
+            )
+        )
+        return list(self.db.scalars(stmt).all())
+
+    def get_subset_distribution(self) -> dict[str, int]:
+        stmt = (
+            select(DatasetItemDB.subset, func.count(DatasetItemDB.id).label("count"))
+            .where(DatasetItemDB.project_id == self.project_id)
+            .group_by(DatasetItemDB.subset)
+        )
+        result = self.db.execute(stmt)
+        return {row.subset: row.count for row in result}  # type: ignore[misc]